don't try to remove assignments in SimplifyComparisonIntegral

Author: Mark-Simulacrum

compiler/rustc_mir_transform/src/simplify_comparison_integral.rs

@@ -76,30 +76,24 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyComparisonIntegral {
                 _ => unreachable!(),
             }
 
-            // delete comparison statement if it the value being switched on was moved, which means
-            // it can not be used later on
-            if opt.can_remove_bin_op_stmt {
-                bb.statements[opt.bin_op_stmt_idx].make_nop(true);
-            } else {
-                // if the integer being compared to a const integral is being moved into the
-                // comparison, e.g `_2 = Eq(move _3, const 'x');`
-                // we want to avoid making a double move later on in the switchInt on _3.
-                // So to avoid `switchInt(move _3) -> ['x': bb2, otherwise: bb1];`,
-                // we convert the move in the comparison statement to a copy.
-
-                // unwrap is safe as we know this statement is an assign
-                let (_, rhs) = bb.statements[opt.bin_op_stmt_idx].kind.as_assign_mut().unwrap();
-
-                use Operand::*;
-                match rhs {
-                    Rvalue::BinaryOp(_, box (left @ Move(_), Constant(_))) => {
-                        *left = Copy(opt.to_switch_on);
-                    }
-                    Rvalue::BinaryOp(_, box (Constant(_), right @ Move(_))) => {
-                        *right = Copy(opt.to_switch_on);
-                    }
-                    _ => (),
+            // if the integer being compared to a const integral is being moved into the
+            // comparison, e.g `_2 = Eq(move _3, const 'x');`
+            // we want to avoid making a double move later on in the switchInt on _3.
+            // So to avoid `switchInt(move _3) -> ['x': bb2, otherwise: bb1];`,
+            // we convert the move in the comparison statement to a copy.
+
+            // unwrap is safe as we know this statement is an assign
+            let (_, rhs) = bb.statements[opt.bin_op_stmt_idx].kind.as_assign_mut().unwrap();
+
+            use Operand::*;
+            match rhs {
+                Rvalue::BinaryOp(_, box (left @ Move(_), Constant(_))) => {
+                    *left = Copy(opt.to_switch_on);
+                }
+                Rvalue::BinaryOp(_, box (Constant(_), right @ Move(_))) => {
+                    *right = Copy(opt.to_switch_on);
                 }
+                _ => (),
             }
 
             let terminator = bb.terminator();
@@ -187,7 +181,6 @@ impl<'tcx> OptimizationFinder<'_, 'tcx> {
                                     Some(OptimizationInfo {
                                         bin_op_stmt_idx: stmt_idx,
                                         bb_idx,
-                                        can_remove_bin_op_stmt: discr.is_move(),
                                         to_switch_on,
                                         branch_value_scalar,
                                         branch_value_ty,
@@ -238,11 +231,8 @@ fn find_branch_value_info<'tcx>(
 struct OptimizationInfo<'tcx> {
     /// Basic block to apply the optimization
     bb_idx: BasicBlock,
-    /// Statement index of Eq/Ne assignment that can be removed. None if the assignment can not be
-    /// removed - i.e the statement is used later on
+    /// Statement index of Eq/Ne assignment
     bin_op_stmt_idx: usize,
-    /// Can remove Eq/Ne assignment
-    can_remove_bin_op_stmt: bool,
     /// Place that needs to be switched on. This place is of type integral
     to_switch_on: Place<'tcx>,
     /// Constant to use in switch target value

tests/mir-opt/if_condition_int.dont_remove_moved_comparison.SimplifyComparisonIntegral.diff

@@ -8,8 +8,7 @@
       let mut _4: i32;
 
       bb0: {
--         _2 = Eq(copy _1, const 17_i8);
-+         nop;
+          _2 = Eq(copy _1, const 17_i8);
           _3 = copy _2 as i32 (IntToInt);
 -         switchInt(move _2) -> [1: bb1, otherwise: bb2];
 +         switchInt(copy _1) -> [17: bb1, otherwise: bb2];

tests/mir-opt/if_condition_int.opt_char.SimplifyComparisonIntegral.diff

@@ -11,9 +11,8 @@
           StorageLive(_2);
           StorageLive(_3);
           _3 = copy _1;
--         _2 = Eq(copy _1, const 'x');
+          _2 = Eq(copy _1, const 'x');
 -         switchInt(move _2) -> [0: bb2, otherwise: bb1];
-+         nop;
 +         switchInt(copy _1) -> [120: bb1, otherwise: bb2];
       }
 

tests/mir-opt/if_condition_int.opt_i8.SimplifyComparisonIntegral.diff

@@ -11,9 +11,8 @@
           StorageLive(_2);
           StorageLive(_3);
           _3 = copy _1;
--         _2 = Eq(copy _1, const 42_i8);
+          _2 = Eq(copy _1, const 42_i8);
 -         switchInt(move _2) -> [0: bb2, otherwise: bb1];
-+         nop;
 +         switchInt(copy _1) -> [42: bb1, otherwise: bb2];
       }
 

tests/mir-opt/if_condition_int.opt_multiple_ifs.SimplifyComparisonIntegral.diff

@@ -13,9 +13,8 @@
           StorageLive(_2);
           StorageLive(_3);
           _3 = copy _1;
--         _2 = Eq(copy _1, const 42_u32);
+          _2 = Eq(copy _1, const 42_u32);
 -         switchInt(move _2) -> [0: bb2, otherwise: bb1];
-+         nop;
 +         switchInt(copy _1) -> [42: bb1, otherwise: bb2];
       }
 
@@ -30,9 +29,8 @@
           StorageLive(_4);
           StorageLive(_5);
           _5 = copy _1;
--         _4 = Ne(copy _1, const 21_u32);
+          _4 = Ne(copy _1, const 21_u32);
 -         switchInt(move _4) -> [0: bb4, otherwise: bb3];
-+         nop;
 +         switchInt(copy _1) -> [21: bb4, otherwise: bb3];
       }
 

tests/mir-opt/if_condition_int.opt_negative.SimplifyComparisonIntegral.diff

@@ -11,9 +11,8 @@
           StorageLive(_2);
           StorageLive(_3);
           _3 = copy _1;
--         _2 = Eq(copy _1, const -42_i32);
+          _2 = Eq(copy _1, const -42_i32);
 -         switchInt(move _2) -> [0: bb2, otherwise: bb1];
-+         nop;
 +         switchInt(copy _1) -> [4294967254: bb1, otherwise: bb2];
       }
 

tests/mir-opt/if_condition_int.opt_u32.SimplifyComparisonIntegral.diff

@@ -11,9 +11,8 @@
           StorageLive(_2);
           StorageLive(_3);
           _3 = copy _1;
--         _2 = Eq(copy _1, const 42_u32);
+          _2 = Eq(copy _1, const 42_u32);
 -         switchInt(move _2) -> [0: bb2, otherwise: bb1];
-+         nop;
 +         switchInt(copy _1) -> [42: bb1, otherwise: bb2];
       }
 

tests/mir-opt/if_condition_int.rs

@@ -85,7 +85,7 @@ fn opt_multiple_ifs(x: u32) -> u32 {
 }
 
 // EMIT_MIR if_condition_int.dont_remove_comparison.SimplifyComparisonIntegral.diff
-// test that we optimize, but do not remove the b statement, as that is used later on
+// the switchInt can be optimized but the b statement can't be removed as it's used later on
 fn dont_remove_comparison(a: i8) -> i32 {
     // CHECK-LABEL: fn dont_remove_comparison(
     // CHECK: [[b:_.*]] = Eq(copy _1, const 17_i8);
@@ -104,8 +104,17 @@ fn dont_remove_comparison(a: i8) -> i32 {
 }
 
 // EMIT_MIR if_condition_int.dont_remove_moved_comparison.SimplifyComparisonIntegral.diff
+// like dont_remove_comparison above, but with switchInt(move _N) - regression test for #158206
 #[custom_mir(dialect = "runtime")]
 fn dont_remove_moved_comparison(a: i8) -> i32 {
+    // CHECK-LABEL: fn dont_remove_moved_comparison(
+    // CHECK: [[b:_.*]] = Eq(copy _1, const 17_i8);
+    // CHECK: [[cast:_.*]] = copy [[b]] as i32 (IntToInt);
+    // CHECK: switchInt(copy _1) -> [17: [[BB1:bb.*]], otherwise: [[BB2:bb.*]]];
+    // CHECK: [[BB1]]:
+    // CHECK: _0 = copy [[cast]];
+    // CHECK: [[BB2]]:
+    // CHECK: _0 = Add(copy [[cast]], const 1_i32);
     mir! {
         let b: bool;
         let c: i32;