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