Auto merge of rust-lang#111813 - scottmcm:pretty-mir, r=cjgillot

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

This doesn't matter at all for actual codegen, but after spending some time reading pre-codegen MIR, I was wishing I didn't have to jump around so much in reading post-inlining code.

So this add two passes that are off by default for every mir level, but can be enabled (`-Zmir-enable-passes=+ReorderBasicBlocks,+ReorderLocals`) for humans.

Author: bors

compiler/rustc_middle/src/mir/terminator.rs

@@ -105,7 +105,7 @@ pub struct Terminator<'tcx> {
     pub kind: TerminatorKind<'tcx>,
 }
 
-pub type Successors<'a> = impl Iterator<Item = BasicBlock> + 'a;
+pub type Successors<'a> = impl DoubleEndedIterator<Item = BasicBlock> + 'a;
 pub type SuccessorsMut<'a> =
     iter::Chain<std::option::IntoIter<&'a mut BasicBlock>, slice::IterMut<'a, BasicBlock>>;
 

compiler/rustc_middle/src/mir/traversal.rs

@@ -149,7 +149,7 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
         //      B     C
         //      |     |
         //      |     |
-        //      D     |
+        //      |     D
         //       \   /
         //        \ /
         //         E
@@ -159,26 +159,26 @@ impl<'a, 'tcx> Postorder<'a, 'tcx> {
         //
         // When the first call to `traverse_successor` happens, the following happens:
         //
-        //     [(B, [D]),  // `B` taken from the successors of `A`, pushed to the
-        //                 // top of the stack along with the successors of `B`
-        //      (A, [C])]
+        //     [(C, [D]),  // `C` taken from the successors of `A`, pushed to the
+        //                 // top of the stack along with the successors of `C`
+        //      (A, [B])]
         //
-        //     [(D, [E]),  // `D` taken from successors of `B`, pushed to stack
-        //      (B, []),
-        //      (A, [C])]
+        //     [(D, [E]),  // `D` taken from successors of `C`, pushed to stack
+        //      (C, []),
+        //      (A, [B])]
         //
         //     [(E, []),   // `E` taken from successors of `D`, pushed to stack
         //      (D, []),
-        //      (B, []),
-        //      (A, [C])]
+        //      (C, []),
+        //      (A, [B])]
         //
         // Now that the top of the stack has no successors we can traverse, each item will
-        // be popped off during iteration until we get back to `A`. This yields [E, D, B].
+        // be popped off during iteration until we get back to `A`. This yields [E, D, C].
         //
-        // When we yield `B` and call `traverse_successor`, we push `C` to the stack, but
+        // When we yield `C` and call `traverse_successor`, we push `B` to the stack, but
         // since we've already visited `E`, that child isn't added to the stack. The last
-        // two iterations yield `C` and finally `A` for a final traversal of [E, D, B, C, A]
-        while let Some(&mut (_, ref mut iter)) = self.visit_stack.last_mut() && let Some(bb) = iter.next() {
+        // two iterations yield `B` and finally `A` for a final traversal of [E, D, C, B, A]
+        while let Some(&mut (_, ref mut iter)) = self.visit_stack.last_mut() && let Some(bb) = iter.next_back() {
             if self.visited.insert(bb) {
                 if let Some(term) = &self.basic_blocks[bb].terminator {
                     self.visit_stack.push((bb, term.successors()));

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"),
         ],

compiler/rustc_mir_transform/src/prettify.rs

@@ -0,0 +1,150 @@
+//! 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;
+
+/// Rearranges the basic blocks into a *reverse post-order*.
+///
+/// Thus after this pass, all the successors of a block are later than it in the
+/// `IndexVec`, unless that successor is a back-edge (such as from a loop).
+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);
+    }
+}
+
+/// Rearranges the locals into *use* order.
+///
+/// Thus after this pass, a local with a smaller [`Location`] where it was first
+/// assigned or referenced will have a smaller number.
+///
+/// (Does not reorder arguments nor the [`RETURN_PLACE`].)
+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 local in (0..=body.arg_count).map(Local::from_usize) {
+            finder.track(local);
+        }
+
+        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 };
+
+        for local in (0..=body.arg_count).map(Local::from_usize) {
+            debug_assert_eq!(updater.map[local], local);
+        }
+
+        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) {
+        // Exclude non-uses to keep `StorageLive` from controlling where we put
+        // a `Local`, since it might not actually be assigned until much later.
+        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];
+    }
+}

src/tools/compiletest/src/runtest.rs

@@ -2045,7 +2045,10 @@ impl<'test> TestCx<'test> {
                 if let Some(pass) = &self.props.mir_unit_test {
                     rustc.args(&["-Zmir-opt-level=0", &format!("-Zmir-enable-passes=+{}", pass)]);
                 } else {
-                    rustc.arg("-Zmir-opt-level=4");
+                    rustc.args(&[
+                        "-Zmir-opt-level=4",
+                        "-Zmir-enable-passes=+ReorderBasicBlocks,+ReorderLocals",
+                    ]);
                 }
 
                 let mir_dump_dir = self.get_mir_dump_dir();

tests/mir-opt/deref-patterns/string.foo.PreCodegen.after.mir

@@ -3,41 +3,33 @@
 fn foo(_1: Option<String>) -> i32 {
     debug s => _1;                       // in scope 0 at $DIR/string.rs:+0:12: +0:13
     let mut _0: i32;                     // return place in scope 0 at $DIR/string.rs:+0:34: +0:37
-    let mut _2: &std::string::String;    // in scope 0 at $DIR/string.rs:+2:14: +2:17
-    let mut _3: &str;                    // in scope 0 at $DIR/string.rs:+2:14: +2:17
-    let mut _4: bool;                    // in scope 0 at $DIR/string.rs:+2:14: +2:17
-    let mut _5: isize;                   // in scope 0 at $DIR/string.rs:+2:9: +2:18
-    let _6: std::option::Option<std::string::String>; // in scope 0 at $DIR/string.rs:+3:9: +3:10
-    let mut _7: bool;                    // in scope 0 at $DIR/string.rs:+5:1: +5:2
+    let mut _2: bool;                    // in scope 0 at $DIR/string.rs:+5:1: +5:2
+    let mut _3: isize;                   // in scope 0 at $DIR/string.rs:+2:9: +2:18
+    let mut _4: &std::string::String;    // in scope 0 at $DIR/string.rs:+2:14: +2:17
+    let mut _5: &str;                    // in scope 0 at $DIR/string.rs:+2:14: +2:17
+    let mut _6: bool;                    // in scope 0 at $DIR/string.rs:+2:14: +2:17
+    let _7: std::option::Option<std::string::String>; // in scope 0 at $DIR/string.rs:+3:9: +3:10
     scope 1 {
-        debug s => _6;                   // in scope 1 at $DIR/string.rs:+3:9: +3:10
+        debug s => _7;                   // in scope 1 at $DIR/string.rs:+3:9: +3:10
     }
 
     bb0: {
-        _7 = const false;                // scope 0 at $DIR/string.rs:+1:11: +1:12
-        _7 = const true;                 // scope 0 at $DIR/string.rs:+1:11: +1:12
-        _5 = discriminant(_1);           // scope 0 at $DIR/string.rs:+1:11: +1:12
-        switchInt(move _5) -> [1: bb2, otherwise: bb1]; // scope 0 at $DIR/string.rs:+1:5: +1:12
+        _2 = const false;                // scope 0 at $DIR/string.rs:+1:11: +1:12
+        _2 = const true;                 // scope 0 at $DIR/string.rs:+1:11: +1:12
+        _3 = discriminant(_1);           // scope 0 at $DIR/string.rs:+1:11: +1:12
+        switchInt(move _3) -> [1: bb1, otherwise: bb5]; // scope 0 at $DIR/string.rs:+1:5: +1:12
     }
 
     bb1: {
-        StorageLive(_6);                 // scope 0 at $DIR/string.rs:+3:9: +3:10
-        _7 = const false;                // scope 0 at $DIR/string.rs:+3:9: +3:10
-        _6 = move _1;                    // scope 0 at $DIR/string.rs:+3:9: +3:10
-        _0 = const 4321_i32;             // scope 1 at $DIR/string.rs:+3:14: +3:18
-        drop(_6) -> [return: bb6, unwind unreachable]; // scope 0 at $DIR/string.rs:+3:17: +3:18
-    }
-
-    bb2: {
-        _2 = &((_1 as Some).0: std::string::String); // scope 0 at $DIR/string.rs:+2:14: +2:17
-        _3 = <String as Deref>::deref(move _2) -> [return: bb3, unwind unreachable]; // scope 0 at $DIR/string.rs:+2:14: +2:17
+        _4 = &((_1 as Some).0: std::string::String); // scope 0 at $DIR/string.rs:+2:14: +2:17
+        _5 = <String as Deref>::deref(move _4) -> [return: bb2, unwind unreachable]; // scope 0 at $DIR/string.rs:+2:14: +2:17
                                          // mir::Constant
                                          // + span: $DIR/string.rs:9:14: 9:17
                                          // + literal: Const { ty: for<'a> fn(&'a String) -> &'a <String as Deref>::Target {<String as Deref>::deref}, val: Value(<ZST>) }
     }
 
-    bb3: {
-        _4 = <str as PartialEq>::eq(_3, const "a") -> [return: bb4, unwind unreachable]; // scope 0 at $DIR/string.rs:+2:14: +2:17
+    bb2: {
+        _6 = <str as PartialEq>::eq(_5, const "a") -> [return: bb3, unwind unreachable]; // scope 0 at $DIR/string.rs:+2:14: +2:17
                                          // mir::Constant
                                          // + span: $DIR/string.rs:9:14: 9:17
                                          // + literal: Const { ty: for<'a, 'b> fn(&'a str, &'b str) -> bool {<str as PartialEq>::eq}, val: Value(<ZST>) }
@@ -46,29 +38,37 @@ fn foo(_1: Option<String>) -> i32 {
                                          // + literal: Const { ty: &str, val: Value(Slice(..)) }
     }
 
+    bb3: {
+        switchInt(move _6) -> [0: bb5, otherwise: bb4]; // scope 0 at $DIR/string.rs:+2:14: +2:17
+    }
+
     bb4: {
-        switchInt(move _4) -> [0: bb1, otherwise: bb5]; // scope 0 at $DIR/string.rs:+2:14: +2:17
+        _0 = const 1234_i32;             // scope 0 at $DIR/string.rs:+2:22: +2:26
+        goto -> bb7;                     // scope 0 at $DIR/string.rs:+2:22: +2:26
     }
 
     bb5: {
-        _0 = const 1234_i32;             // scope 0 at $DIR/string.rs:+2:22: +2:26
-        goto -> bb9;                     // scope 0 at $DIR/string.rs:+2:22: +2:26
+        StorageLive(_7);                 // scope 0 at $DIR/string.rs:+3:9: +3:10
+        _2 = const false;                // scope 0 at $DIR/string.rs:+3:9: +3:10
+        _7 = move _1;                    // scope 0 at $DIR/string.rs:+3:9: +3:10
+        _0 = const 4321_i32;             // scope 1 at $DIR/string.rs:+3:14: +3:18
+        drop(_7) -> [return: bb6, unwind unreachable]; // scope 0 at $DIR/string.rs:+3:17: +3:18
     }
 
     bb6: {
-        StorageDead(_6);                 // scope 0 at $DIR/string.rs:+3:17: +3:18
-        goto -> bb9;                     // scope 0 at $DIR/string.rs:+3:17: +3:18
+        StorageDead(_7);                 // scope 0 at $DIR/string.rs:+3:17: +3:18
+        goto -> bb7;                     // scope 0 at $DIR/string.rs:+3:17: +3:18
     }
 
     bb7: {
-        return;                          // scope 0 at $DIR/string.rs:+5:2: +5:2
+        switchInt(_2) -> [0: bb9, otherwise: bb8]; // scope 0 at $DIR/string.rs:+5:1: +5:2
     }
 
     bb8: {
-        drop(_1) -> [return: bb7, unwind unreachable]; // scope 0 at $DIR/string.rs:+5:1: +5:2
+        drop(_1) -> [return: bb9, unwind unreachable]; // scope 0 at $DIR/string.rs:+5:1: +5:2
     }
 
     bb9: {
-        switchInt(_7) -> [0: bb7, otherwise: bb8]; // scope 0 at $DIR/string.rs:+5:1: +5:2
+        return;                          // scope 0 at $DIR/string.rs:+5:2: +5:2
     }
 }

tests/mir-opt/inline/cycle.g.Inline.diff

@@ -8,8 +8,8 @@
 +     let mut _5: ();                      // in scope 0 at $DIR/cycle.rs:6:5: 6:8
 +     scope 1 (inlined f::<fn() {main}>) { // at $DIR/cycle.rs:12:5: 12:12
 +         debug g => _2;                   // in scope 1 at $DIR/cycle.rs:5:6: 5:7
-+         let _3: ();                      // in scope 1 at $DIR/cycle.rs:6:5: 6:8
-+         let mut _4: &fn() {main};        // in scope 1 at $DIR/cycle.rs:6:5: 6:6
++         let mut _3: &fn() {main};        // in scope 1 at $DIR/cycle.rs:6:5: 6:6
++         let _4: ();                      // in scope 1 at $DIR/cycle.rs:6:5: 6:8
 +         scope 2 (inlined <fn() {main} as Fn<()>>::call - shim(fn() {main})) { // at $DIR/cycle.rs:6:5: 6:8
 +         }
 +     }
@@ -25,16 +25,16 @@
 -                                          // mir::Constant
                                            // + span: $DIR/cycle.rs:12:7: 12:11
                                            // + literal: Const { ty: fn() {main}, val: Value(<ZST>) }
-+         StorageLive(_3);                 // scope 0 at $DIR/cycle.rs:+1:5: +1:12
-+         StorageLive(_4);                 // scope 1 at $DIR/cycle.rs:6:5: 6:6
-+         _4 = &_2;                        // scope 1 at $DIR/cycle.rs:6:5: 6:6
++         StorageLive(_4);                 // scope 0 at $DIR/cycle.rs:+1:5: +1:12
++         StorageLive(_3);                 // scope 1 at $DIR/cycle.rs:6:5: 6:6
++         _3 = &_2;                        // scope 1 at $DIR/cycle.rs:6:5: 6:6
 +         StorageLive(_5);                 // scope 1 at $DIR/cycle.rs:6:5: 6:8
 +         _5 = const ();                   // scope 1 at $DIR/cycle.rs:6:5: 6:8
-+         _3 = move (*_4)() -> [return: bb4, unwind: bb2]; // scope 2 at $SRC_DIR/core/src/ops/function.rs:LL:COL
++         _4 = move (*_3)() -> [return: bb4, unwind: bb2]; // scope 2 at $SRC_DIR/core/src/ops/function.rs:LL:COL
       }
 
       bb1: {
-+         StorageDead(_3);                 // scope 0 at $DIR/cycle.rs:+1:5: +1:12
++         StorageDead(_4);                 // scope 0 at $DIR/cycle.rs:+1:5: +1:12
 +         StorageDead(_2);                 // scope 0 at $DIR/cycle.rs:+1:5: +1:12
           StorageDead(_1);                 // scope 0 at $DIR/cycle.rs:+1:12: +1:13
           _0 = const ();                   // scope 0 at $DIR/cycle.rs:+0:8: +2:2
@@ -51,7 +51,7 @@
 + 
 +     bb4: {
 +         StorageDead(_5);                 // scope 1 at $DIR/cycle.rs:6:5: 6:8
-+         StorageDead(_4);                 // scope 1 at $DIR/cycle.rs:6:7: 6:8
++         StorageDead(_3);                 // scope 1 at $DIR/cycle.rs:6:7: 6:8
 +         drop(_2) -> bb1;                 // scope 1 at $DIR/cycle.rs:7:1: 7:2
       }
   }