Auto merge of #97158 - JakobDegen:dse, r=tmiasko,cjgillot

Split dead store elimination off dest prop

This splits off a part of #96451 . I've added this in as its own pass for now, so that it actually runs, can be tested, etc. In the dest prop PR, I'll stop invoking this as its own pass, so that it doesn't get invoked twice.

r? `@tmiasko`

2 files changed, 151 insertions, 1 deletions

diff --git a/compiler/rustc_mir_transform/src/dead_store_elimination.rs b/compiler/rustc_mir_transform/src/dead_store_elimination.rs
new file mode 100644
index 00000000000..84f2ee639e4
--- /dev/null
+++ b/compiler/rustc_mir_transform/src/dead_store_elimination.rs
@@ -0,0 +1,148 @@
+//! This module implements a dead store elimination (DSE) routine.
+//!
+//! This transformation was written specifically for the needs of dest prop. Although it is
+//! perfectly sound to use it in any context that might need it, its behavior should not be changed
+//! without analyzing the interaction this will have with dest prop. Specifically, in addition to
+//! the soundness of this pass in general, dest prop needs it to satisfy two additional conditions:
+//!
+//!  1. It's idempotent, meaning that running this pass a second time immediately after running it a
+//!     first time will not cause any further changes.
+//!  2. This idempotence persists across dest prop's main transform, in other words inserting any
+//!     number of iterations of dest prop between the first and second application of this transform
+//!     will still not cause any further changes.
+//!
+
+use rustc_index::bit_set::BitSet;
+use rustc_middle::{
+    mir::{visit::Visitor, *},
+    ty::TyCtxt,
+};
+use rustc_mir_dataflow::{impls::MaybeTransitiveLiveLocals, Analysis};
+
+/// Performs the optimization on the body
+///
+/// The `borrowed` set must be a `BitSet` of all the locals that are ever borrowed in this body. It
+/// can be generated via the [`get_borrowed_locals`] function.
+pub fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, borrowed: &BitSet<Local>) {
+    let mut live = MaybeTransitiveLiveLocals::new(borrowed, &body.local_decls, tcx)
+        .into_engine(tcx, body)
+        .iterate_to_fixpoint()
+        .into_results_cursor(body);
+
+    let mut patch = Vec::new();
+    for (bb, bb_data) in traversal::preorder(body) {
+        for (statement_index, statement) in bb_data.statements.iter().enumerate().rev() {
+            let loc = Location { block: bb, statement_index };
+            if let StatementKind::Assign(assign) = &statement.kind {
+                if assign.1.is_pointer_int_cast(&body.local_decls, tcx) {
+                    continue;
+                }
+            }
+            match &statement.kind {
+                StatementKind::Assign(box (place, _))
+                | StatementKind::SetDiscriminant { place: box place, .. }
+                | StatementKind::Deinit(box place) => {
+                    if !place.is_indirect() && !borrowed.contains(place.local) {
+                        live.seek_before_primary_effect(loc);
+                        if !live.get().contains(place.local) {
+                            patch.push(loc);
+                        }
+                    }
+                }
+                StatementKind::Retag(_, _)
+                | StatementKind::StorageLive(_)
+                | StatementKind::StorageDead(_)
+                | StatementKind::Coverage(_)
+                | StatementKind::CopyNonOverlapping(_)
+                | StatementKind::Nop => (),
+
+                StatementKind::FakeRead(_) | StatementKind::AscribeUserType(_, _) => {
+                    bug!("{:?} not found in this MIR phase!", &statement.kind)
+                }
+            }
+        }
+    }
+
+    if patch.is_empty() {
+        return;
+    }
+
+    let bbs = body.basic_blocks_mut();
+    for Location { block, statement_index } in patch {
+        bbs[block].statements[statement_index].make_nop();
+    }
+}
+
+pub fn get_borrowed_locals(body: &Body<'_>) -> BitSet<Local> {
+    let mut b = BorrowedLocals(BitSet::new_empty(body.local_decls.len()));
+    b.visit_body(body);
+    b.0
+}
+
+struct BorrowedLocals(BitSet<Local>);
+
+impl<'tcx> Visitor<'tcx> for BorrowedLocals {
+    fn visit_rvalue(&mut self, rvalue: &Rvalue<'tcx>, loc: Location) {
+        self.super_rvalue(rvalue, loc);
+        match rvalue {
+            Rvalue::AddressOf(_, borrowed_place) | Rvalue::Ref(_, _, borrowed_place) => {
+                if !borrowed_place.is_indirect() {
+                    self.0.insert(borrowed_place.local);
+                }
+            }
+
+            Rvalue::Cast(..)
+            | Rvalue::ShallowInitBox(..)
+            | Rvalue::Use(..)
+            | Rvalue::Repeat(..)
+            | Rvalue::Len(..)
+            | Rvalue::BinaryOp(..)
+            | Rvalue::CheckedBinaryOp(..)
+            | Rvalue::NullaryOp(..)
+            | Rvalue::UnaryOp(..)
+            | Rvalue::Discriminant(..)
+            | Rvalue::Aggregate(..)
+            | Rvalue::ThreadLocalRef(..) => {}
+        }
+    }
+
+    fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
+        self.super_terminator(terminator, location);
+
+        match terminator.kind {
+            TerminatorKind::Drop { place: dropped_place, .. } => {
+                if !dropped_place.is_indirect() {
+                    self.0.insert(dropped_place.local);
+                }
+            }
+
+            TerminatorKind::Abort
+            | TerminatorKind::DropAndReplace { .. }
+            | TerminatorKind::Assert { .. }
+            | TerminatorKind::Call { .. }
+            | TerminatorKind::FalseEdge { .. }
+            | TerminatorKind::FalseUnwind { .. }
+            | TerminatorKind::GeneratorDrop
+            | TerminatorKind::Goto { .. }
+            | TerminatorKind::Resume
+            | TerminatorKind::Return
+            | TerminatorKind::SwitchInt { .. }
+            | TerminatorKind::Unreachable
+            | TerminatorKind::Yield { .. }
+            | TerminatorKind::InlineAsm { .. } => {}
+        }
+    }
+}
+
+pub struct DeadStoreElimination;
+
+impl<'tcx> MirPass<'tcx> for DeadStoreElimination {
+    fn is_enabled(&self, sess: &rustc_session::Session) -> bool {
+        sess.mir_opt_level() >= 2
+    }
+
+    fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
+        let borrowed = get_borrowed_locals(body);
+        eliminate(tcx, body, &borrowed);
+    }
+}
diff --git a/compiler/rustc_mir_transform/src/lib.rs b/compiler/rustc_mir_transform/src/lib.rs
index 571f541072a..0e57c3c6979 100644
--- a/compiler/rustc_mir_transform/src/lib.rs
+++ b/compiler/rustc_mir_transform/src/lib.rs
@@ -49,6 +49,7 @@ mod const_goto;
 mod const_prop;
 mod const_prop_lint;
 mod coverage;
+mod dead_store_elimination;
 mod deaggregator;
 mod deduplicate_blocks;
 mod deref_separator;
@@ -481,17 +482,18 @@ fn run_optimization_passes<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
             &const_prop::ConstProp,
             //
             // Const-prop runs unconditionally, but doesn't mutate the MIR at mir-opt-level=0.
+            &const_debuginfo::ConstDebugInfo,
             &o1(simplify_branches::SimplifyConstCondition::new("after-const-prop")),
             &early_otherwise_branch::EarlyOtherwiseBranch,
             &simplify_comparison_integral::SimplifyComparisonIntegral,
             &simplify_try::SimplifyArmIdentity,
             &simplify_try::SimplifyBranchSame,
+            &dead_store_elimination::DeadStoreElimination,
             &dest_prop::DestinationPropagation,
             &o1(simplify_branches::SimplifyConstCondition::new("final")),
             &o1(remove_noop_landing_pads::RemoveNoopLandingPads),
             &o1(simplify::SimplifyCfg::new("final")),
             &nrvo::RenameReturnPlace,
-            &const_debuginfo::ConstDebugInfo,
             &simplify::SimplifyLocals,
             &multiple_return_terminators::MultipleReturnTerminators,
             &deduplicate_blocks::DeduplicateBlocks,