rustc_mir_transform/
prettify.rs
1use 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
14pub(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
44pub(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 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 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 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 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}