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;
12
13use crate::PassPolicy;
14
15pub(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
41pub(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 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 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 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 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}