rustc_mir_transform/
prettify.rs1use 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
16pub(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
42pub(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 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 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 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 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}