Move constants into the SwitchInt to be easier to see

Author: scottmcm

compiler/rustc_mir_transform/src/lib.rs

@@ -108,6 +108,7 @@ mod simplify_branches;
 mod simplify_comparison_integral;
 mod single_use_consts;
 mod sroa;
+mod switch_const;
 mod unreachable_enum_branching;
 mod unreachable_prop;
 mod validate;
@@ -598,6 +599,8 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &simplify::SimplifyLocals::AfterGVN,
             &dataflow_const_prop::DataflowConstProp,
             &single_use_consts::SingleUseConsts,
+            // 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,56 @@
+//! A pass that makes `SwitchInt`-on-`const` more obvious to later code.
+
+use rustc_middle::bug;
+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/single_use_consts.rs

@@ -29,9 +29,7 @@ fn match_const<T: MyTrait>() -> &'static str {
 fn if_const_debug<T: MyTrait>() -> i32 {
     // CHECK-LABEL: fn if_const_debug(
     // CHECK: my_bool => const <T as MyTrait>::ASSOC_BOOL;
-    // FIXME: `if` forces a temporary (unlike `match`), so the const isn't direct
-    // CHECK: _3 = const <T as MyTrait>::ASSOC_BOOL;
-    // CHECK: switchInt(move _3)
+    // CHECK: switchInt(const <T as MyTrait>::ASSOC_BOOL)
     let my_bool = T::ASSOC_BOOL;
     do_whatever();
     if my_bool { 7 } else { 42 }