mir-opt: Ignore the dead store statements in MatchBranchSimplification

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

@@ -12,6 +12,7 @@
 //!     will still not cause any further changes.
 //!
 
+use rustc_index::IndexVec;
 use rustc_middle::bug;
 use rustc_middle::mir::visit::Visitor;
 use rustc_middle::mir::*;
@@ -24,19 +25,52 @@ use rustc_mir_dataflow::impls::{
 
 use crate::util::is_within_packed;
 
+pub(super) enum ModifyBasicBlocks<'tcx, 'a> {
+    Direct(&'a mut Body<'tcx>),
+    BasicBlocks(&'a Body<'tcx>, &'a mut IndexVec<BasicBlock, BasicBlockData<'tcx>>),
+}
+
+impl<'tcx, 'a> ModifyBasicBlocks<'tcx, 'a> {
+    pub(super) fn body(&self) -> &Body<'tcx> {
+        match self {
+            ModifyBasicBlocks::Direct(body) => body,
+            ModifyBasicBlocks::BasicBlocks(body, _) => body,
+        }
+    }
+
+    pub(super) fn bbs(&mut self) -> &mut IndexVec<BasicBlock, BasicBlockData<'tcx>> {
+        match self {
+            ModifyBasicBlocks::Direct(body) => body.basic_blocks.as_mut_preserves_cfg(),
+            ModifyBasicBlocks::BasicBlocks(_, bbs) => bbs,
+        }
+    }
+}
+
 /// Performs the optimization on the body
 ///
 /// The `borrowed` set must be a `DenseBitSet` of all the locals that are ever borrowed in this
 /// body. It can be generated via the [`borrowed_locals`] function.
-fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+pub(super) fn eliminate<'tcx>(
+    tcx: TyCtxt<'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);
 
     // If the user requests complete debuginfo, mark the locals that appear in it as live, so
     // we don't remove assignments to them.
-    let mut always_live = debuginfo_locals(body);
-    always_live.union(&borrowed_locals);
-
-    let mut live = MaybeTransitiveLiveLocals::new(&always_live)
+    let always_live = if ignore_debuginfo {
+        borrowed_locals.clone()
+    } else {
+        let mut always_live = debuginfo_locals(body);
+        always_live.union(&borrowed_locals);
+        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);
 
@@ -46,7 +80,8 @@ fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     let mut patch = Vec::new();
 
     for (bb, bb_data) in traversal::preorder(body) {
-        if let TerminatorKind::Call { ref args, .. } = bb_data.terminator().kind {
+        if arg_copy_to_move && let TerminatorKind::Call { ref args, .. } = bb_data.terminator().kind
+        {
             let loc = Location { block: bb, statement_index: bb_data.statements.len() };
 
             // Position ourselves between the evaluation of `args` and the write to `destination`.
@@ -113,7 +148,7 @@ fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         return;
     }
 
-    let bbs = body.basic_blocks.as_mut_preserves_cfg();
+    let bbs = modify_basic_blocks.bbs();
     for Location { block, statement_index } in patch {
         bbs[block].statements[statement_index].make_nop();
     }
@@ -145,7 +180,7 @@ impl<'tcx> crate::MirPass<'tcx> for DeadStoreElimination {
     }
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-        eliminate(tcx, body);
+        eliminate(tcx, ModifyBasicBlocks::Direct(body), false, true, false);
     }
 
     fn is_required(&self) -> bool {

compiler/rustc_mir_transform/src/lib.rs

@@ -700,8 +700,6 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             // Now, we need to shrink the generated MIR.
             &ref_prop::ReferencePropagation,
             &sroa::ScalarReplacementOfAggregates,
-            &match_branches::MatchBranchSimplification,
-            // inst combine is after MatchBranchSimplification to clean up Ne(_1, false)
             &multiple_return_terminators::MultipleReturnTerminators,
             // After simplifycfg, it allows us to discover new opportunities for peephole
             // optimizations.
@@ -710,6 +708,7 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &dead_store_elimination::DeadStoreElimination::Initial,
             &gvn::GVN,
             &simplify::SimplifyLocals::AfterGVN,
+            &match_branches::MatchBranchSimplification,
             &dataflow_const_prop::DataflowConstProp,
             &single_use_consts::SingleUseConsts,
             &o1(simplify_branches::SimplifyConstCondition::AfterConstProp),

compiler/rustc_mir_transform/src/match_branches.rs

@@ -1,14 +1,17 @@
 use std::iter;
 
 use rustc_abi::Integer;
-use rustc_index::IndexSlice;
+use rustc_index::{IndexSlice, IndexVec};
+use rustc_middle::mir::visit::Visitor;
 use rustc_middle::mir::*;
 use rustc_middle::ty::layout::{IntegerExt, TyAndLayout};
 use rustc_middle::ty::{self, ScalarInt, Ty, TyCtxt};
 use tracing::instrument;
 
 use super::simplify::simplify_cfg;
+use crate::dead_store_elimination::{self, ModifyBasicBlocks};
 use crate::patch::MirPatch;
+use crate::simplify::strip_nops;
 
 pub(super) struct MatchBranchSimplification;
 
@@ -19,32 +22,50 @@ impl<'tcx> crate::MirPass<'tcx> for MatchBranchSimplification {
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         let typing_env = body.typing_env(tcx);
-        let mut should_cleanup = false;
-        for i in 0..body.basic_blocks.len() {
-            let bbs = &*body.basic_blocks;
-            let bb_idx = BasicBlock::from_usize(i);
-            match bbs[bb_idx].terminator().kind {
+        let mut apply_patch = false;
+        let mut bbs = body.basic_blocks.clone();
+        let bbs = bbs.as_mut_preserves_cfg();
+        // We can ignore the dead store statements when merging branches.
+        dead_store_elimination::eliminate(
+            tcx,
+            ModifyBasicBlocks::BasicBlocks(body, bbs),
+            true,
+            false,
+            true,
+        );
+        eliminate_unused_storage_mark(body, bbs);
+        strip_nops(bbs.as_mut_slice());
+        let mut patch = MirPatch::new(body);
+        for (bb, data) in body.basic_blocks.iter_enumerated() {
+            match data.terminator().kind {
                 TerminatorKind::SwitchInt {
                     discr: ref _discr @ (Operand::Copy(_) | Operand::Move(_)),
                     ref targets,
                     ..
                     // We require that the possible target blocks don't contain this block.
-                } if !targets.all_targets().contains(&bb_idx) => {}
+                } if !targets.all_targets().contains(&bb) => {}
                 // Only optimize switch int statements
                 _ => continue,
             };
 
-            if SimplifyToIf.simplify(tcx, body, bb_idx, typing_env).is_some() {
-                should_cleanup = true;
+            if SimplifyToIf
+                .simplify(tcx, body, bbs.as_slice(), &mut patch, bb, typing_env)
+                .is_some()
+            {
+                apply_patch = true;
                 continue;
             }
-            if SimplifyToExp::default().simplify(tcx, body, bb_idx, typing_env).is_some() {
-                should_cleanup = true;
+            if SimplifyToExp::default()
+                .simplify(tcx, body, bbs.as_slice(), &mut patch, bb, typing_env)
+                .is_some()
+            {
+                apply_patch = true;
                 continue;
             }
         }
 
-        if should_cleanup {
+        if apply_patch {
+            patch.apply(body);
             simplify_cfg(body);
         }
     }
@@ -61,11 +82,12 @@ trait SimplifyMatch<'tcx> {
     fn simplify(
         &mut self,
         tcx: TyCtxt<'tcx>,
-        body: &mut Body<'tcx>,
+        body: &Body<'tcx>,
+        bbs: &IndexSlice<BasicBlock, BasicBlockData<'tcx>>,
+        patch: &mut MirPatch<'tcx>,
         switch_bb_idx: BasicBlock,
         typing_env: ty::TypingEnv<'tcx>,
     ) -> Option<()> {
-        let bbs = &body.basic_blocks;
         let (discr, targets) = match bbs[switch_bb_idx].terminator().kind {
             TerminatorKind::SwitchInt { ref discr, ref targets, .. } => (discr, targets),
             _ => unreachable!(),
@@ -74,8 +96,6 @@ trait SimplifyMatch<'tcx> {
         let discr_ty = discr.ty(body.local_decls(), tcx);
         self.can_simplify(tcx, targets, typing_env, bbs, discr_ty)?;
 
-        let mut patch = MirPatch::new(body);
-
         // Take ownership of items now that we know we can optimize.
         let discr = discr.clone();
 
@@ -84,23 +104,13 @@ trait SimplifyMatch<'tcx> {
         let discr_local = patch.new_temp(discr_ty, source_info.span);
 
         let (_, first) = targets.iter().next().unwrap();
-        let statement_index = bbs[switch_bb_idx].statements.len();
+        let statement_index = body.basic_blocks[switch_bb_idx].statements.len();
         let parent_end = Location { block: switch_bb_idx, statement_index };
         patch.add_statement(parent_end, StatementKind::StorageLive(discr_local));
         patch.add_assign(parent_end, Place::from(discr_local), Rvalue::Use(discr));
-        self.new_stmts(
-            tcx,
-            targets,
-            typing_env,
-            &mut patch,
-            parent_end,
-            bbs,
-            discr_local,
-            discr_ty,
-        );
+        self.new_stmts(tcx, targets, typing_env, patch, parent_end, bbs, discr_local, discr_ty);
         patch.add_statement(parent_end, StatementKind::StorageDead(discr_local));
         patch.patch_terminator(switch_bb_idx, bbs[first].terminator().kind.clone());
-        patch.apply(body);
         Some(())
     }
 
@@ -537,3 +547,43 @@ impl<'tcx> SimplifyMatch<'tcx> for SimplifyToExp {
         }
     }
 }
+
+struct EliminateUnusedStorageMark {
+    storage_live_locals: IndexVec<Local, Option<usize>>,
+}
+
+impl<'tcx> Visitor<'tcx> for EliminateUnusedStorageMark {
+    fn visit_local(&mut self, local: Local, _: visit::PlaceContext, _: Location) {
+        self.storage_live_locals[local] = None;
+    }
+}
+
+fn eliminate_unused_storage_mark<'tcx>(
+    body: &Body<'tcx>,
+    basic_blocks: &mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+) {
+    for (bb, data) in basic_blocks.iter_enumerated_mut() {
+        let mut unused_storage_mark = EliminateUnusedStorageMark {
+            storage_live_locals: IndexVec::from_elem_n(None, body.local_decls.len()),
+        };
+        for stmt_index in 0..data.statements.len() {
+            let loc = Location { block: bb, statement_index: stmt_index };
+            match data.statements[stmt_index].kind {
+                StatementKind::StorageLive(local) => {
+                    unused_storage_mark.storage_live_locals[local] = Some(stmt_index);
+                }
+                StatementKind::StorageDead(local)
+                    if let Some(live_stmt_index) =
+                        unused_storage_mark.storage_live_locals[local] =>
+                {
+                    data.statements[live_stmt_index].make_nop();
+                    data.statements[stmt_index].make_nop();
+                    unused_storage_mark.storage_live_locals[local] = None;
+                }
+                _ => {
+                    unused_storage_mark.visit_statement(&data.statements[stmt_index], loc);
+                }
+            }
+        }
+    }
+}

compiler/rustc_mir_transform/src/simplify.rs

@@ -116,7 +116,7 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
     }
 
     fn simplify(mut self) {
-        self.strip_nops();
+        strip_nops(self.basic_blocks);
 
         // Vec of the blocks that should be merged. We store the indices here, instead of the
         // statements itself to avoid moving the (relatively) large statements twice.
@@ -276,11 +276,11 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
         terminator.kind = TerminatorKind::Goto { target: first_succ };
         true
     }
+}
 
-    fn strip_nops(&mut self) {
-        for blk in self.basic_blocks.iter_mut() {
-            blk.statements.retain(|stmt| !matches!(stmt.kind, StatementKind::Nop))
-        }
+pub(super) fn strip_nops(basic_blocks: &mut IndexSlice<BasicBlock, BasicBlockData<'_>>) {
+    for blk in basic_blocks.iter_mut() {
+        blk.statements.retain(|stmt| !matches!(stmt.kind, StatementKind::Nop))
     }
 }