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