Skip to main content

rustc_mir_transform/
prettify.rs

1//! These two passes provide no value to the compiler, so are off at every level.
2//!
3//! However, they can be enabled on the command line
4//! (`-Zmir-enable-passes=+ReorderBasicBlocks,+ReorderLocals`)
5//! to make the MIR easier to read for humans.
6
7use rustc_index::bit_set::DenseBitSet;
8use rustc_index::{IndexSlice, IndexVec};
9use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
10use rustc_middle::mir::*;
11use rustc_middle::ty::TyCtxt;
12use rustc_session::Session;
13
14use crate::PassPolicy;
15
16/// Rearranges the basic blocks into a *reverse post-order*.
17///
18/// Thus after this pass, all the successors of a block are later than it in the
19/// `IndexVec`, unless that successor is a back-edge (such as from a loop).
20pub(super) struct ReorderBasicBlocks;
21
22impl<'tcx> crate::MirPass<'tcx> for ReorderBasicBlocks {
23    fn policy(&self, _session: &Session) -> PassPolicy {
24        PassPolicy::optional_non_optimization(false)
25    }
26
27    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
28        let rpo: IndexVec<BasicBlock, BasicBlock> =
29            body.basic_blocks.reverse_postorder().iter().copied().collect();
30        if rpo.iter().is_sorted() {
31            return;
32        }
33
34        let mut updater = BasicBlockUpdater { map: rpo.invert_bijective_mapping(), tcx };
35        debug_assert_eq!(updater.map[START_BLOCK], START_BLOCK);
36        updater.visit_body(body);
37
38        permute(body.basic_blocks.as_mut(), &updater.map);
39    }
40}
41
42/// Rearranges the locals into *use* order.
43///
44/// Thus after this pass, a local with a smaller [`Location`] where it was first
45/// assigned or referenced will have a smaller number.
46///
47/// (Does not reorder arguments nor the [`RETURN_PLACE`].)
48pub(super) struct ReorderLocals;
49
50impl<'tcx> crate::MirPass<'tcx> for ReorderLocals {
51    fn policy(&self, _session: &Session) -> PassPolicy {
52        PassPolicy::optional_non_optimization(false)
53    }
54
55    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
56        let mut finder = LocalFinder {
57            map: IndexVec::new(),
58            seen: DenseBitSet::new_empty(body.local_decls.len()),
59        };
60
61        // We can't reorder the return place or the arguments
62        for local in (0..=body.arg_count).map(Local::from_usize) {
63            finder.track(local);
64        }
65
66        for (bb, bbd) in body.basic_blocks.iter_enumerated() {
67            finder.visit_basic_block_data(bb, bbd);
68        }
69
70        // Track everything in case there are some locals that we never saw,
71        // such as in non-block things like debug info or in non-uses.
72        for local in body.local_decls.indices() {
73            finder.track(local);
74        }
75
76        if finder.map.iter().is_sorted() {
77            return;
78        }
79
80        let mut updater = LocalUpdater { map: finder.map.invert_bijective_mapping(), tcx };
81
82        for local in (0..=body.arg_count).map(Local::from_usize) {
83            debug_assert_eq!(updater.map[local], local);
84        }
85
86        updater.visit_body_preserves_cfg(body);
87
88        permute(&mut body.local_decls, &updater.map);
89    }
90}
91
92fn permute<I: rustc_index::Idx + Ord, T>(data: &mut IndexVec<I, T>, map: &IndexSlice<I, I>) {
93    // FIXME: It would be nice to have a less-awkward way to apply permutations,
94    // but I don't know one that exists. `sort_by_cached_key` has logic for it
95    // internally, but not in a way that we're allowed to use here.
96    let mut enumerated: Vec<_> = std::mem::take(data).into_iter_enumerated().collect();
97    enumerated.sort_by_key(|p| map[p.0]);
98    *data = enumerated.into_iter().map(|p| p.1).collect();
99}
100
101struct BasicBlockUpdater<'tcx> {
102    map: IndexVec<BasicBlock, BasicBlock>,
103    tcx: TyCtxt<'tcx>,
104}
105
106impl<'tcx> MutVisitor<'tcx> for BasicBlockUpdater<'tcx> {
107    fn tcx(&self) -> TyCtxt<'tcx> {
108        self.tcx
109    }
110
111    fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, _location: Location) {
112        terminator.successors_mut(|succ| *succ = self.map[*succ]);
113    }
114}
115
116struct LocalFinder {
117    map: IndexVec<Local, Local>,
118    seen: DenseBitSet<Local>,
119}
120
121impl LocalFinder {
122    fn track(&mut self, l: Local) {
123        if self.seen.insert(l) {
124            self.map.push(l);
125        }
126    }
127}
128
129impl<'tcx> Visitor<'tcx> for LocalFinder {
130    fn visit_local(&mut self, l: Local, context: PlaceContext, _location: Location) {
131        // Exclude non-uses to keep `StorageLive` from controlling where we put
132        // a `Local`, since it might not actually be assigned until much later.
133        if context.is_use() {
134            self.track(l);
135        }
136    }
137}
138
139struct LocalUpdater<'tcx> {
140    map: IndexVec<Local, Local>,
141    tcx: TyCtxt<'tcx>,
142}
143
144impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> {
145    fn tcx(&self) -> TyCtxt<'tcx> {
146        self.tcx
147    }
148
149    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
150        *l = self.map[*l];
151    }
152}