mir-opt: Do not eliminate dead statements that are being used by other BBs during merging

Since we will not eliminate dead statements in other BBs,
they may still use the dead statements within the merged BB.

Author: dianqk

compiler/rustc_mir_dataflow/src/impls/liveness.rs

@@ -3,8 +3,9 @@ use rustc_middle::mir::visit::{MutatingUseContext, NonMutatingUseContext, PlaceC
 use rustc_middle::mir::{
     self, CallReturnPlaces, Local, Location, Place, StatementKind, TerminatorEdges,
 };
+use rustc_middle::ty::TyCtxt;
 
-use crate::{Analysis, Backward, GenKill};
+use crate::{Analysis, Backward, GenKill, ResultsCursor};
 
 /// A [live-variable dataflow analysis][liveness].
 ///
@@ -203,21 +204,42 @@ impl DefUse {
 /// This is basically written for dead store elimination and nothing else.
 ///
 /// All of the caveats of `MaybeLiveLocals` apply.
-pub struct MaybeTransitiveLiveLocals<'a> {
+pub struct MaybeTransitiveLiveLocals<'tcx, 'mir, 'a> {
     always_live: &'a DenseBitSet<Local>,
+    live: Option<ResultsCursor<'mir, 'tcx, MaybeLiveLocals>>,
 }
 
-impl<'a> MaybeTransitiveLiveLocals<'a> {
+impl<'tcx, 'mir, 'a> MaybeTransitiveLiveLocals<'tcx, 'mir, 'a> {
     /// The `always_alive` set is the set of locals to which all stores should unconditionally be
     /// considered live.
+    /// The `may_live_in_other_bbs` indicates that, when analyzing the current statement,
+    /// statements in other basic blocks are assumed to be live.
     ///
     /// This should include at least all locals that are ever borrowed.
-    pub fn new(always_live: &'a DenseBitSet<Local>) -> Self {
-        MaybeTransitiveLiveLocals { always_live }
+    pub fn new(
+        always_live: &'a DenseBitSet<Local>,
+        may_live_in_other_bbs: bool,
+        tcx: TyCtxt<'tcx>,
+        body: &'mir mir::Body<'tcx>,
+    ) -> Self {
+        let live = if may_live_in_other_bbs {
+            Some(MaybeLiveLocals.iterate_to_fixpoint(tcx, body, None).into_results_cursor(body))
+        } else {
+            None
+        };
+        MaybeTransitiveLiveLocals { always_live, live }
+    }
+
+    fn live_on(&mut self, place: Place<'_>, location: Location) -> bool {
+        let Some(live) = &mut self.live else {
+            return false;
+        };
+        live.seek_before_primary_effect(location);
+        live.get().contains(place.local)
     }
 }
 
-impl<'a, 'tcx> Analysis<'tcx> for MaybeTransitiveLiveLocals<'a> {
+impl<'a, 'tcx> Analysis<'tcx> for MaybeTransitiveLiveLocals<'tcx, '_, 'a> {
     type Domain = DenseBitSet<Local>;
     type Direction = Backward;
 
@@ -256,14 +278,14 @@ impl<'a, 'tcx> Analysis<'tcx> for MaybeTransitiveLiveLocals<'a> {
             | StatementKind::BackwardIncompatibleDropHint { .. }
             | StatementKind::Nop => None,
         };
-        if let Some(destination) = destination {
-            if !destination.is_indirect()
-                && !state.contains(destination.local)
-                && !self.always_live.contains(destination.local)
-            {
-                // This store is dead
-                return;
-            }
+        if let Some(destination) = destination
+            && !destination.is_indirect()
+            && !state.contains(destination.local)
+            && !self.always_live.contains(destination.local)
+            && !self.live_on(destination, location)
+        {
+            // This store is dead
+            return;
         }
         TransferFunction(state).visit_statement(statement, location);
     }

compiler/rustc_mir_transform/src/dead_store_elimination.rs

@@ -55,6 +55,7 @@ pub(super) fn eliminate<'tcx>(
     mut modify_basic_blocks: ModifyBasicBlocks<'tcx, '_>,
     ignore_debuginfo: bool,
     arg_copy_to_move: bool,
+    may_live_in_other_bbs: bool,
 ) {
     let body = modify_basic_blocks.body();
     let borrowed_locals = borrowed_locals(body);
@@ -69,7 +70,7 @@ pub(super) fn eliminate<'tcx>(
         always_live
     };
 
-    let mut live = MaybeTransitiveLiveLocals::new(&always_live)
+    let mut live = MaybeTransitiveLiveLocals::new(&always_live, may_live_in_other_bbs, tcx, body)
         .iterate_to_fixpoint(tcx, body, None)
         .into_results_cursor(body);
 
@@ -179,7 +180,7 @@ impl<'tcx> crate::MirPass<'tcx> for DeadStoreElimination {
     }
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-        eliminate(tcx, ModifyBasicBlocks::Direct(body), false, true);
+        eliminate(tcx, ModifyBasicBlocks::Direct(body), false, true, false);
     }
 
     fn is_required(&self) -> bool {

compiler/rustc_mir_transform/src/match_branches.rs

@@ -31,6 +31,7 @@ impl<'tcx> crate::MirPass<'tcx> for MatchBranchSimplification {
             ModifyBasicBlocks::BasicBlocks(body, bbs),
             true,
             false,
+            true,
         );
         eliminate_unused_storage_mark(body, bbs);
         strip_nops(bbs.as_mut_slice());

tests/mir-opt/matches_reduce_branches.match_dead_store_failed.MatchBranchSimplification.diff

@@ -0,0 +1,45 @@
+- // MIR for `match_dead_store_failed` before MatchBranchSimplification
++ // MIR for `match_dead_store_failed` after MatchBranchSimplification
+  
+  fn match_dead_store_failed(_1: bool) -> () {
+      debug b => _1;
+      let mut _0: ();
+      let _2: (&str,);
+      let mut _3: &str;
+      let mut _4: bool;
+      let _5: &str;
+      scope 1 {
+          debug _b => _2;
+      }
+  
+      bb0: {
+          StorageLive(_2);
+          StorageLive(_3);
+          StorageLive(_4);
+          _4 = copy _1;
+          switchInt(move _4) -> [0: bb2, otherwise: bb1];
+      }
+  
+      bb1: {
+          _3 = const " ";
+          goto -> bb3;
+      }
+  
+      bb2: {
+          StorageLive(_5);
+          _5 = const "";
+          _3 = &(*_5);
+          StorageDead(_5);
+          goto -> bb3;
+      }
+  
+      bb3: {
+          StorageDead(_4);
+          _2 = (move _3,);
+          StorageDead(_3);
+          _0 = const ();
+          StorageDead(_2);
+          return;
+      }
+  }
+  

tests/mir-opt/matches_reduce_branches.rs

@@ -672,6 +672,15 @@ fn match_i128_u128(i: EnumAi128) -> u128 {
     }
 }
 
+// We cannot merge these branches unless all the dead statements are eliminated.
+// EMIT_MIR matches_reduce_branches.match_dead_store_failed.MatchBranchSimplification.diff
+pub fn match_dead_store_failed(b: bool) {
+    // CHECK-LABEL: fn match_dead_store_failed(
+    // CHECK: switchInt
+    // CHECK: return
+    let _b = (if b { " " } else { "" },);
+}
+
 fn main() {
     let _ = foo(None);
     let _ = foo(Some(()));