Move constants into the SwitchInt to be easier to see

Author: scottmcm

compiler/rustc_mir_transform/src/lib.rs

@@ -110,6 +110,7 @@ pub mod simplify;
 mod simplify_branches;
 mod simplify_comparison_integral;
 mod sroa;
+mod switch_const;
 mod uninhabited_enum_branching;
 mod unreachable_prop;
 
@@ -600,6 +601,8 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &simplify::SimplifyLocals::AfterGVN,
             &dataflow_const_prop::DataflowConstProp,
             &const_debuginfo::ConstDebugInfo,
+            // GVN & ConstProp often don't fixup unevaluatable constants
+            &switch_const::SwitchConst,
             &o1(simplify_branches::SimplifyConstCondition::AfterConstProp),
             &jump_threading::JumpThreading,
             &early_otherwise_branch::EarlyOtherwiseBranch,

compiler/rustc_mir_transform/src/switch_const.rs

@@ -0,0 +1,55 @@
+//! A pass that makes `SwitchInt`-on-`const` more obvious to later code.
+
+use rustc_middle::mir::*;
+use rustc_middle::ty::TyCtxt;
+
+/// A `MirPass` for simplifying `if T::CONST`.
+///
+/// Today, MIR building for things like `if T::IS_ZST` introduce a constant
+/// for the copy of the bool, so it ends up in MIR as
+/// `_1 = CONST; switchInt (move _1)` or `_2 = CONST; switchInt (_2)`.
+///
+/// This pass is very specifically targeted at *exactly* those patterns.
+/// It can absolutely be replaced with a more general pass should we get one that
+/// we can run in low optimization levels, but at the time of writing even in
+/// optimized builds this wasn't simplified.
+#[derive(Default)]
+pub struct SwitchConst;
+
+impl<'tcx> MirPass<'tcx> for SwitchConst {
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        for block in body.basic_blocks.as_mut_preserves_cfg() {
+            let switch_local = if let TerminatorKind::SwitchInt { discr, .. } =
+                &block.terminator().kind
+                && let Some(place) = discr.place()
+                && let Some(local) = place.as_local()
+            {
+                local
+            } else {
+                continue;
+            };
+
+            let new_operand = if let Some(statement) = block.statements.last()
+                && let StatementKind::Assign(place_and_rvalue) = &statement.kind
+                && let Some(local) = place_and_rvalue.0.as_local()
+                && local == switch_local
+                && let Rvalue::Use(operand) = &place_and_rvalue.1
+                && let Operand::Constant(_) = operand
+            {
+                operand.clone()
+            } else {
+                continue;
+            };
+
+            if !tcx.consider_optimizing(|| format!("SwitchConst: switchInt(move {switch_local:?}"))
+            {
+                break;
+            }
+
+            let TerminatorKind::SwitchInt { discr, .. } = &mut block.terminator_mut().kind else {
+                bug!("Somehow wasn't a switchInt any more?")
+            };
+            *discr = new_operand;
+        }
+    }
+}

tests/mir-opt/pre-codegen/if_associated_const.check_bool.PreCodegen.after.mir

@@ -2,12 +2,9 @@
 
 fn check_bool() -> u32 {
     let mut _0: u32;
-    let mut _1: bool;
 
     bb0: {
-        StorageLive(_1);
-        _1 = const <T as TraitWithBool>::FLAG;
-        switchInt(move _1) -> [0: bb1, otherwise: bb2];
+        switchInt(const <T as TraitWithBool>::FLAG) -> [0: bb1, otherwise: bb2];
     }
 
     bb1: {
@@ -21,7 +18,6 @@ fn check_bool() -> u32 {
     }
 
     bb3: {
-        StorageDead(_1);
         return;
     }
 }

tests/mir-opt/pre-codegen/if_associated_const.check_int.PreCodegen.after.mir

@@ -2,12 +2,9 @@
 
 fn check_int() -> u32 {
     let mut _0: u32;
-    let mut _1: i32;
 
     bb0: {
-        StorageLive(_1);
-        _1 = const <T as TraitWithInt>::VALUE;
-        switchInt(_1) -> [1: bb1, 2: bb2, 3: bb3, otherwise: bb4];
+        switchInt(const <T as TraitWithInt>::VALUE) -> [1: bb1, 2: bb2, 3: bb3, otherwise: bb4];
     }
 
     bb1: {
@@ -31,7 +28,6 @@ fn check_int() -> u32 {
     }
 
     bb5: {
-        StorageDead(_1);
         return;
     }
 }

tests/mir-opt/pre-codegen/slice_iter.enumerated_loop.PreCodegen.after.panic-abort.mir

@@ -4,22 +4,22 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
     debug slice => _1;
     debug f => _2;
     let mut _0: ();
-    let mut _13: std::slice::Iter<'_, T>;
+    let mut _12: std::slice::Iter<'_, T>;
+    let mut _13: std::iter::Enumerate<std::slice::Iter<'_, T>>;
     let mut _14: std::iter::Enumerate<std::slice::Iter<'_, T>>;
-    let mut _15: std::iter::Enumerate<std::slice::Iter<'_, T>>;
-    let mut _16: &mut std::iter::Enumerate<std::slice::Iter<'_, T>>;
-    let mut _17: std::option::Option<(usize, &T)>;
-    let mut _18: isize;
-    let mut _21: &impl Fn(usize, &T);
-    let mut _22: (usize, &T);
-    let _23: ();
+    let mut _15: &mut std::iter::Enumerate<std::slice::Iter<'_, T>>;
+    let mut _16: std::option::Option<(usize, &T)>;
+    let mut _17: isize;
+    let mut _20: &impl Fn(usize, &T);
+    let mut _21: (usize, &T);
+    let _22: ();
     scope 1 {
-        debug iter => _15;
-        let _19: usize;
-        let _20: &T;
+        debug iter => _14;
+        let _18: usize;
+        let _19: &T;
         scope 2 {
-            debug i => _19;
-            debug x => _20;
+            debug i => _18;
+            debug x => _19;
         }
     }
     scope 3 (inlined core::slice::<impl [T]>::iter) {
@@ -28,19 +28,18 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
             debug slice => _1;
             let _3: usize;
             let mut _5: std::ptr::NonNull<[T]>;
-            let mut _8: bool;
+            let mut _8: *mut T;
             let mut _9: *mut T;
-            let mut _10: *mut T;
-            let mut _12: *const T;
+            let mut _11: *const T;
             scope 5 {
                 debug len => _3;
                 let _7: std::ptr::NonNull<T>;
                 scope 6 {
                     debug ptr => _7;
                     scope 7 {
-                        let _11: *const T;
+                        let _10: *const T;
                         scope 8 {
-                            debug end_or_len => _11;
+                            debug end_or_len => _10;
                         }
                         scope 14 (inlined without_provenance::<T>) {
                             debug addr => _3;
@@ -51,7 +50,7 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
                             debug self => _7;
                         }
                         scope 17 (inlined std::ptr::mut_ptr::<impl *mut T>::add) {
-                            debug self => _9;
+                            debug self => _8;
                             debug count => _3;
                             scope 18 {
                             }
@@ -77,17 +76,17 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
         }
     }
     scope 19 (inlined <std::slice::Iter<'_, T> as Iterator>::enumerate) {
-        debug self => _13;
+        debug self => _12;
         scope 20 (inlined Enumerate::<std::slice::Iter<'_, T>>::new) {
-            debug iter => _13;
+            debug iter => _12;
         }
     }
     scope 21 (inlined <Enumerate<std::slice::Iter<'_, T>> as IntoIterator>::into_iter) {
-        debug self => _14;
+        debug self => _13;
     }
 
     bb0: {
-        StorageLive(_13);
+        StorageLive(_12);
         StorageLive(_3);
         StorageLive(_7);
         StorageLive(_4);
@@ -99,62 +98,59 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
         _6 = _4 as *const T (PtrToPtr);
         _7 = NonNull::<T> { pointer: _6 };
         StorageDead(_5);
-        StorageLive(_11);
-        StorageLive(_8);
-        _8 = const <T as std::mem::SizedTypeProperties>::IS_ZST;
-        switchInt(move _8) -> [0: bb1, otherwise: bb2];
+        StorageLive(_10);
+        switchInt(const <T as std::mem::SizedTypeProperties>::IS_ZST) -> [0: bb1, otherwise: bb2];
     }
 
     bb1: {
-        StorageLive(_10);
         StorageLive(_9);
-        _9 = _4 as *mut T (PtrToPtr);
-        _10 = Offset(_9, _3);
+        StorageLive(_8);
+        _8 = _4 as *mut T (PtrToPtr);
+        _9 = Offset(_8, _3);
+        StorageDead(_8);
+        _10 = move _9 as *const T (PointerCoercion(MutToConstPointer));
         StorageDead(_9);
-        _11 = move _10 as *const T (PointerCoercion(MutToConstPointer));
-        StorageDead(_10);
         goto -> bb3;
     }
 
     bb2: {
-        _11 = _3 as *const T (Transmute);
+        _10 = _3 as *const T (Transmute);
         goto -> bb3;
     }
 
     bb3: {
-        StorageDead(_8);
-        StorageLive(_12);
-        _12 = _11;
-        _13 = std::slice::Iter::<'_, T> { ptr: _7, end_or_len: move _12, _marker: const ZeroSized: PhantomData<&T> };
-        StorageDead(_12);
+        StorageLive(_11);
+        _11 = _10;
+        _12 = std::slice::Iter::<'_, T> { ptr: _7, end_or_len: move _11, _marker: const ZeroSized: PhantomData<&T> };
         StorageDead(_11);
+        StorageDead(_10);
         StorageDead(_6);
         StorageDead(_4);
         StorageDead(_7);
         StorageDead(_3);
-        _14 = Enumerate::<std::slice::Iter<'_, T>> { iter: _13, count: const 0_usize };
-        StorageDead(_13);
-        StorageLive(_15);
-        _15 = _14;
+        _13 = Enumerate::<std::slice::Iter<'_, T>> { iter: _12, count: const 0_usize };
+        StorageDead(_12);
+        StorageLive(_14);
+        _14 = _13;
         goto -> bb4;
     }
 
     bb4: {
-        StorageLive(_17);
         StorageLive(_16);
-        _16 = &mut _15;
-        _17 = <Enumerate<std::slice::Iter<'_, T>> as Iterator>::next(move _16) -> [return: bb5, unwind unreachable];
+        StorageLive(_15);
+        _15 = &mut _14;
+        _16 = <Enumerate<std::slice::Iter<'_, T>> as Iterator>::next(move _15) -> [return: bb5, unwind unreachable];
     }
 
     bb5: {
-        StorageDead(_16);
-        _18 = discriminant(_17);
-        switchInt(move _18) -> [0: bb6, 1: bb8, otherwise: bb10];
+        StorageDead(_15);
+        _17 = discriminant(_16);
+        switchInt(move _17) -> [0: bb6, 1: bb8, otherwise: bb10];
     }
 
     bb6: {
-        StorageDead(_17);
-        StorageDead(_15);
+        StorageDead(_16);
+        StorageDead(_14);
         drop(_2) -> [return: bb7, unwind unreachable];
     }
 
@@ -163,19 +159,19 @@ fn enumerated_loop(_1: &[T], _2: impl Fn(usize, &T)) -> () {
     }
 
     bb8: {
-        _19 = (((_17 as Some).0: (usize, &T)).0: usize);
-        _20 = (((_17 as Some).0: (usize, &T)).1: &T);
+        _18 = (((_16 as Some).0: (usize, &T)).0: usize);
+        _19 = (((_16 as Some).0: (usize, &T)).1: &T);
+        StorageLive(_20);
+        _20 = &_2;
         StorageLive(_21);
-        _21 = &_2;
-        StorageLive(_22);
-        _22 = (_19, _20);
-        _23 = <impl Fn(usize, &T) as Fn<(usize, &T)>>::call(move _21, move _22) -> [return: bb9, unwind unreachable];
+        _21 = (_18, _19);
+        _22 = <impl Fn(usize, &T) as Fn<(usize, &T)>>::call(move _20, move _21) -> [return: bb9, unwind unreachable];
     }
 
     bb9: {
-        StorageDead(_22);
         StorageDead(_21);
-        StorageDead(_17);
+        StorageDead(_20);
+        StorageDead(_16);
         goto -> bb4;
     }