MIR: opt-in normalization of `BasicBlock` and `Local` numbering

2 files changed, 139 insertions, 0 deletions

diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 65864dc016f..54c138b6fbd 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/compiler/rustc_mir_transform/src/lib.rs
@@ -3,6 +3,7 @@
 #![deny(rustc::diagnostic_outside_of_impl)]
 #![feature(box_patterns)]
 #![feature(drain_filter)]
+#![feature(is_sorted)]
 #![feature(let_chains)]
 #![feature(map_try_insert)]
 #![feature(min_specialization)]
@@ -84,6 +85,7 @@ mod match_branches;
 mod multiple_return_terminators;
 mod normalize_array_len;
 mod nrvo;
+mod prettify;
 mod ref_prop;
 mod remove_noop_landing_pads;
 mod remove_storage_markers;
@@ -581,6 +583,9 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &large_enums::EnumSizeOpt { discrepancy: 128 },
             // Some cleanup necessary at least for LLVM and potentially other codegen backends.
             &add_call_guards::CriticalCallEdges,
+            // Cleanup for human readability, off by default.
+            &prettify::ReorderBasicBlocks,
+            &prettify::ReorderLocals,
             // Dump the end result for testing and debugging purposes.
             &dump_mir::Marker("PreCodegen"),
         ],
diff --git a/compiler/rustc_mir_transform/src/prettify.rs b/compiler/rustc_mir_transform/src/prettify.rs
new file mode 100644
index 00000000000..d3637405edc
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/prettify.rs
@@ -0,0 +1,134 @@
+//! These two passes provide no value to the compiler, so are off at every level.
+//!
+//! However, they can be enabled on the command line
+//! (`-Zmir-enable-passes=+ReorderBasicBlocks,+ReorderLocals`)
+//! to make the MIR easier to read for humans.
+
+use crate::MirPass;
+use rustc_index::{bit_set::BitSet, IndexSlice, IndexVec};
+use rustc_middle::mir::visit::{MutVisitor, PlaceContext, Visitor};
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+use rustc_session::Session;
+
+pub struct ReorderBasicBlocks;
+
+impl<'tcx> MirPass<'tcx> for ReorderBasicBlocks {
+    fn is_enabled(&self, _session: &Session) -> bool {
+        false
+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let rpo: IndexVec<BasicBlock, BasicBlock> =
+            body.basic_blocks.postorder().iter().copied().rev().collect();
+        if rpo.iter().is_sorted() {
+            return;
+        }
+
+        let mut updater = BasicBlockUpdater { map: rpo.invert_bijective_mapping(), tcx };
+        debug_assert_eq!(updater.map[START_BLOCK], START_BLOCK);
+        updater.visit_body(body);
+
+        permute(body.basic_blocks.as_mut(), &updater.map);
+    }
+}
+
+pub struct ReorderLocals;
+
+impl<'tcx> MirPass<'tcx> for ReorderLocals {
+    fn is_enabled(&self, _session: &Session) -> bool {
+        false
+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let mut finder =
+            LocalFinder { map: IndexVec::new(), seen: BitSet::new_empty(body.local_decls.len()) };
+
+        // We can't reorder the return place or the arguments
+        for i in 0..=body.arg_count {
+            finder.track(Local::from_usize(i));
+        }
+
+        for (bb, bbd) in body.basic_blocks.iter_enumerated() {
+            finder.visit_basic_block_data(bb, bbd);
+        }
+
+        // track everything in case there are some locals that we never saw,
+        // such as in non-block things like debug info or in non-uses.
+        for local in body.local_decls.indices() {
+            finder.track(local);
+        }
+
+        if finder.map.iter().is_sorted() {
+            return;
+        }
+
+        let mut updater = LocalUpdater { map: finder.map.invert_bijective_mapping(), tcx };
+        debug_assert_eq!(updater.map[RETURN_PLACE], RETURN_PLACE);
+        updater.visit_body_preserves_cfg(body);
+
+        permute(&mut body.local_decls, &updater.map);
+    }
+}
+
+fn permute<I: rustc_index::Idx + Ord, T>(data: &mut IndexVec<I, T>, map: &IndexSlice<I, I>) {
+    // FIXME: It would be nice to have a less-awkward way to apply permutations,
+    // but I don't know one that exists.  `sort_by_cached_key` has logic for it
+    // internally, but not in a way that we're allowed to use here.
+    let mut enumerated: Vec<_> = std::mem::take(data).into_iter_enumerated().collect();
+    enumerated.sort_by_key(|p| map[p.0]);
+    *data = enumerated.into_iter().map(|p| p.1).collect();
+}
+
+struct BasicBlockUpdater<'tcx> {
+    map: IndexVec<BasicBlock, BasicBlock>,
+    tcx: TyCtxt<'tcx>,
+}
+
+impl<'tcx> MutVisitor<'tcx> for BasicBlockUpdater<'tcx> {
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn visit_terminator(&mut self, terminator: &mut Terminator<'tcx>, _location: Location) {
+        for succ in terminator.successors_mut() {
+            *succ = self.map[*succ];
+        }
+    }
+}
+
+struct LocalFinder {
+    map: IndexVec<Local, Local>,
+    seen: BitSet<Local>,
+}
+
+impl LocalFinder {
+    fn track(&mut self, l: Local) {
+        if self.seen.insert(l) {
+            self.map.push(l);
+        }
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for LocalFinder {
+    fn visit_local(&mut self, l: Local, context: PlaceContext, _location: Location) {
+        if context.is_use() {
+            self.track(l);
+        }
+    }
+}
+
+struct LocalUpdater<'tcx> {
+    pub map: IndexVec<Local, Local>,
+    pub tcx: TyCtxt<'tcx>,
+}
+
+impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> {
+    fn tcx(&self) -> TyCtxt<'tcx> {
+        self.tcx
+    }
+
+    fn visit_local(&mut self, l: &mut Local, _: PlaceContext, _: Location) {
+        *l = self.map[*l];
+    }
+}