Don't optimize across storage markers in SimplifyComparisonIntegral

compiler/rustc_mir_transform/src/simplify_comparison_integral.rs

@@ -3,8 +3,7 @@ use std::iter;
 use rustc_middle::bug;
 use rustc_middle::mir::interpret::Scalar;
 use rustc_middle::mir::{
-    BasicBlock, BinOp, Body, Operand, Place, Rvalue, Statement, StatementKind, SwitchTargets,
-    TerminatorKind,
+    BasicBlock, BinOp, Body, Operand, Place, Rvalue, StatementKind, SwitchTargets, TerminatorKind,
 };
 use rustc_middle::ty::{Ty, TyCtxt};
 use tracing::trace;
@@ -16,7 +15,6 @@ use crate::ssa::SsaLocals;
 ///
 /// ```ignore (MIR)
 /// _3 = Eq(move _4, const 43i32);
-/// StorageDead(_4);
 /// switchInt(_3) -> [false: bb2, otherwise: bb3];
 /// ```
 ///
@@ -39,8 +37,6 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyComparisonIntegral {
         let ssa = SsaLocals::new(tcx, body, typing_env);
         let helper = OptimizationFinder { body };
         let opts = helper.find_optimizations(&ssa);
-        let mut storage_deads_to_insert = vec![];
-        let mut storage_deads_to_remove: Vec<(usize, BasicBlock)> = vec![];
         for opt in opts {
             trace!("SUCCESS: Applying {:?}", opt);
             // replace terminator with a switchInt that switches on the integer directly
@@ -96,30 +92,6 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyComparisonIntegral {
                 _ => (),
             }
 
-            let terminator = bb.terminator();
-
-            // remove StorageDead (if it exists) being used in the assign of the comparison
-            for (stmt_idx, stmt) in bb.statements.iter().enumerate() {
-                if !matches!(
-                    stmt.kind,
-                    StatementKind::StorageDead(local) if local == opt.to_switch_on.local
-                ) {
-                    continue;
-                }
-                storage_deads_to_remove.push((stmt_idx, opt.bb_idx));
-                // if we have StorageDeads to remove then make sure to insert them at the top of
-                // each target
-                for bb_idx in new_targets.all_targets() {
-                    storage_deads_to_insert.push((
-                        *bb_idx,
-                        Statement::new(
-                            terminator.source_info,
-                            StatementKind::StorageDead(opt.to_switch_on.local),
-                        ),
-                    ));
-                }
-            }
-
             let [bb_cond, bb_otherwise] = match new_targets.all_targets() {
                 [a, b] => [*a, *b],
                 e => bug!("expected 2 switch targets, got: {:?}", e),
@@ -131,14 +103,6 @@ impl<'tcx> crate::MirPass<'tcx> for SimplifyComparisonIntegral {
             terminator.kind =
                 TerminatorKind::SwitchInt { discr: Operand::Copy(opt.to_switch_on), targets };
         }
-
-        for (idx, bb_idx) in storage_deads_to_remove {
-            body.basic_blocks_mut()[bb_idx].statements[idx].make_nop(true);
-        }
-
-        for (idx, stmt) in storage_deads_to_insert {
-            body.basic_blocks_mut()[idx].statements.insert(0, stmt);
-        }
     }
 
     fn is_required(&self) -> bool {
@@ -175,6 +139,20 @@ impl<'tcx> OptimizationFinder<'_, 'tcx> {
                                     let (branch_value_scalar, branch_value_ty, to_switch_on) =
                                         find_branch_value_info(left, right, ssa)?;
 
+                                    // The transformation adds a use of `to_switch_on` at the
+                                    // terminator. Both storage markers make the local uninitialized,
+                                    // so either invalidates the value used by the comparison.
+                                    if bb.statements[stmt_idx + 1..].iter().any(|stmt| {
+                                        matches!(
+                                            stmt.kind,
+                                            StatementKind::StorageLive(local)
+                                                | StatementKind::StorageDead(local)
+                                                if local == to_switch_on.local
+                                        )
+                                    }) {
+                                        return None;
+                                    }
+
                                     Some(OptimizationInfo {
                                         bin_op_stmt_idx: stmt_idx,
                                         bb_idx,

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

@@ -0,0 +1,27 @@
+- // MIR for `dont_opt_storage_live_after_comparison` before SimplifyComparisonIntegral
++ // MIR for `dont_opt_storage_live_after_comparison` after SimplifyComparisonIntegral
+
+  fn dont_opt_storage_live_after_comparison(_1: bool) -> () {
+      let mut _0: ();
+      let mut _2: u32;
+      let mut _3: bool;
+
+      bb0: {
+          StorageLive(_2);
+          goto -> bb1;
+      }
+
+      bb1: {
+          _2 = copy _1 as u32 (IntToInt);
+          _3 = Eq(copy _2, const 42_u32);
+          StorageDead(_2);
+          StorageLive(_2);
+          switchInt(move _3) -> [1: bb1, otherwise: bb2];
+      }
+
+      bb2: {
+          StorageDead(_2);
+          return;
+      }
+  }
+

tests/mir-opt/if_condition_int.rs

@@ -139,6 +139,40 @@ fn dont_remove_moved_comparison(a: i8) -> i32 {
     }
 }
 
+// EMIT_MIR if_condition_int.dont_opt_storage_live_after_comparison.SimplifyComparisonIntegral.diff
+// Regression test for https://github.com/rust-lang/rust/issues/158231.
+#[custom_mir(dialect = "runtime")]
+fn dont_opt_storage_live_after_comparison(a: bool) {
+    // CHECK-LABEL: fn dont_opt_storage_live_after_comparison(
+    // CHECK: [[b:_.*]] = copy _1 as u32 (IntToInt);
+    // CHECK: [[cmp:_.*]] = Eq(copy [[b]], const 42_u32);
+    // CHECK: StorageDead([[b]]);
+    // CHECK: StorageLive([[b]]);
+    // CHECK: switchInt(move [[cmp]]) -> [1: {{bb.*}}, otherwise: {{bb.*}}];
+    mir! {
+        let b: u32;
+        let c: bool;
+        {
+            StorageLive(b);
+            Goto(bb1)
+        }
+        bb1 = {
+            b = a as u32;
+            c = b == 42;
+            StorageDead(b);
+            StorageLive(b);
+            match Move(c) {
+                true => bb1,
+                _ => bb2,
+            }
+        }
+        bb2 = {
+            StorageDead(b);
+            Return()
+        }
+    }
+}
+
 // EMIT_MIR if_condition_int.dont_opt_floats.SimplifyComparisonIntegral.diff
 // test that we do not optimize on floats
 fn dont_opt_floats(a: f32) -> i32 {

tests/mir-opt/if_condition_int_storage.dont_opt_storage_dead_in_loop.SimplifyComparisonIntegral.diff

@@ -0,0 +1,29 @@
+- // MIR for `dont_opt_storage_dead_in_loop` before SimplifyComparisonIntegral
++ // MIR for `dont_opt_storage_dead_in_loop` after SimplifyComparisonIntegral
+
+  fn dont_opt_storage_dead_in_loop(_1: u32) -> () {
+      let mut _0: ();
+      let mut _2: u32;
+      let mut _3: bool;
+
+      bb0: {
+          StorageLive(_2);
+          _2 = copy _1;
+          goto -> bb1;
+      }
+
+      bb1: {
+          goto -> bb2;
+      }
+
+      bb2: {
+          _3 = Eq(copy _2, const 123_u32);
+          StorageDead(_2);
+          switchInt(move _3) -> [1: bb1, otherwise: bb3];
+      }
+
+      bb3: {
+          return;
+      }
+  }
+

tests/mir-opt/if_condition_int_storage.rs

@@ -0,0 +1,45 @@
+//@ test-mir-pass: SimplifyComparisonIntegral
+// MIR lint assumes all paths can execute. The backedge is undefined, but the first traversal is
+// defined and must stay defined.
+//@ compile-flags: -Zlint-mir=false
+
+#![feature(custom_mir, core_intrinsics)]
+
+extern crate core;
+use core::intrinsics::mir::*;
+
+// EMIT_MIR if_condition_int_storage.dont_opt_storage_dead_in_loop.SimplifyComparisonIntegral.diff
+// Regression test for https://github.com/rust-lang/rust/issues/158231.
+#[custom_mir(dialect = "runtime", phase = "post-cleanup")]
+fn dont_opt_storage_dead_in_loop(input: u32) {
+    // CHECK-LABEL: fn dont_opt_storage_dead_in_loop(
+    // CHECK: [[a:_.*]] = copy _1;
+    // CHECK: [[cmp:_.*]] = Eq(copy [[a]], const 123_u32);
+    // CHECK: StorageDead([[a]]);
+    // CHECK: switchInt(move [[cmp]]) -> [1: {{bb.*}}, otherwise: {{bb.*}}];
+    mir! {
+        let a: u32;
+        let cmp: bool;
+        {
+            StorageLive(a);
+            a = input;
+            Goto(bb1)
+        }
+        bb1 = {
+            Goto(bb2)
+        }
+        bb2 = {
+            cmp = a == 123;
+            StorageDead(a);
+            match Move(cmp) {
+                true => bb1,
+                _ => bb3,
+            }
+        }
+        bb3 = {
+            Return()
+        }
+    }
+}
+
+fn main() {}

tests/mir-opt/issues/issue_59352.num_to_digit.runtime-optimized.after.32bit.panic-abort.mir

@@ -3,7 +3,8 @@
 fn num_to_digit(_1: char) -> u32 {
     debug num => _1;
     let mut _0: u32;
-    let mut _4: std::option::Option<u32>;
+    let mut _4: bool;
+    let mut _5: std::option::Option<u32>;
     scope 1 (inlined char::methods::<impl char>::is_digit) {
         let _2: std::option::Option<u32>;
         scope 2 (inlined Option::<u32>::is_some) {
@@ -13,58 +14,60 @@ fn num_to_digit(_1: char) -> u32 {
         }
     }
     scope 4 (inlined #[track_caller] Option::<u32>::unwrap) {
-        let mut _5: isize;
-        let mut _6: !;
+        let mut _6: isize;
+        let mut _7: !;
         scope 5 {
         }
     }
 
     bb0: {
+        StorageLive(_4);
         StorageLive(_2);
         _2 = char::methods::<impl char>::to_digit(copy _1, const 8_u32) -> [return: bb1, unwind unreachable];
     }
 
     bb1: {
         StorageLive(_3);
         _3 = discriminant(_2);
+        _4 = Eq(copy _3, const 1_isize);
+        StorageDead(_3);
         StorageDead(_2);
-        switchInt(copy _3) -> [1: bb2, otherwise: bb7];
+        switchInt(move _4) -> [0: bb2, otherwise: bb3];
     }
 
     bb2: {
-        StorageDead(_3);
-        StorageLive(_4);
-        _4 = char::methods::<impl char>::to_digit(move _1, const 8_u32) -> [return: bb3, unwind unreachable];
+        _0 = const 0_u32;
+        goto -> bb7;
     }
 
     bb3: {
         StorageLive(_5);
-        _5 = discriminant(_4);
-        switchInt(move _5) -> [0: bb4, 1: bb5, otherwise: bb6];
+        _5 = char::methods::<impl char>::to_digit(move _1, const 8_u32) -> [return: bb4, unwind unreachable];
     }
 
     bb4: {
-        _6 = option::unwrap_failed() -> unwind unreachable;
+        StorageLive(_6);
+        _6 = discriminant(_5);
+        switchInt(move _6) -> [0: bb5, 1: bb6, otherwise: bb8];
     }
 
     bb5: {
-        _0 = move ((_4 as Some).0: u32);
-        StorageDead(_5);
-        StorageDead(_4);
-        goto -> bb8;
+        _7 = option::unwrap_failed() -> unwind unreachable;
     }
 
     bb6: {
-        unreachable;
+        _0 = move ((_5 as Some).0: u32);
+        StorageDead(_6);
+        StorageDead(_5);
+        goto -> bb7;
     }
 
     bb7: {
-        StorageDead(_3);
-        _0 = const 0_u32;
-        goto -> bb8;
+        StorageDead(_4);
+        return;
     }
 
     bb8: {
-        return;
+        unreachable;
     }
 }