mir-opt: Eliminate trivial unnecessary storage annotations

Author: dianqk

compiler/rustc_middle/src/mir/mod.rs

@@ -1429,6 +1429,17 @@ impl<'tcx> BasicBlockData<'tcx> {
         });
         self.after_last_stmt_debuginfos.extend_from_slice(&debuginfos);
     }
+
+    pub fn strip_nops(&mut self) {
+        self.retain_statements(|stmt| !matches!(stmt.kind, StatementKind::Nop))
+    }
+
+    pub fn drop_debuginfo(&mut self) {
+        self.after_last_stmt_debuginfos = Vec::new();
+        for stmt in self.statements.iter_mut() {
+            stmt.debuginfos = Vec::new();
+        }
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////

compiler/rustc_middle/src/mir/statement.rs

@@ -27,16 +27,10 @@ impl<'tcx> Statement<'tcx> {
         }
         let replaced_stmt = std::mem::replace(&mut self.kind, StatementKind::Nop);
         if !drop_debuginfo {
-            match replaced_stmt {
-                StatementKind::Assign(box (place, Rvalue::Ref(_, _, ref_place)))
-                    if let Some(local) = place.as_local() =>
-                {
-                    self.debuginfos.push(StmtDebugInfo::AssignRef(local, ref_place));
-                }
-                _ => {
-                    bug!("debuginfo is not yet supported.")
-                }
-            }
+            let Some(debuginfo) = replaced_stmt.as_debuginfo() else {
+                bug!("debuginfo is not yet supported.")
+            };
+            self.debuginfos.push(debuginfo);
         }
     }
 
@@ -79,6 +73,17 @@ impl<'tcx> StatementKind<'tcx> {
             _ => None,
         }
     }
+
+    pub fn as_debuginfo(&self) -> Option<StmtDebugInfo<'tcx>> {
+        match self {
+            StatementKind::Assign(box (place, Rvalue::Ref(_, _, ref_place)))
+                if let Some(local) = place.as_local() =>
+            {
+                Some(StmtDebugInfo::AssignRef(local, *ref_place))
+            }
+            _ => None,
+        }
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////

compiler/rustc_mir_dataflow/src/impls/liveness.rs

@@ -228,8 +228,10 @@ impl<'a> MaybeTransitiveLiveLocals<'a> {
         // Compute the place that we are storing to, if any
         let destination = match stmt_kind {
             StatementKind::Assign(box (place, rvalue)) => (rvalue.is_safe_to_remove()
+                // FIXME: We are not sure how we should represent this debugging information for some statements,
+                // keep it for now.
                 && (!debuginfo_locals.contains(place.local)
-                    || (place.as_local().is_some() && matches!(rvalue, mir::Rvalue::Ref(..)))))
+                    || (place.as_local().is_some() && stmt_kind.as_debuginfo().is_some())))
             .then_some(*place),
             StatementKind::SetDiscriminant { place, .. } | StatementKind::Deinit(place) => {
                 (!debuginfo_locals.contains(place.local)).then_some(**place)

compiler/rustc_mir_transform/src/dead_store_elimination.rs

@@ -22,13 +22,14 @@ use rustc_mir_dataflow::impls::{
     LivenessTransferFunction, MaybeTransitiveLiveLocals, borrowed_locals,
 };
 
+use crate::simplify::UsedInStmtLocals;
 use crate::util::is_within_packed;
 
 /// 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>) {
+fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) -> bool {
     let borrowed_locals = borrowed_locals(body);
 
     // If the user requests complete debuginfo, mark the locals that appear in it as live, so
@@ -89,8 +90,9 @@ fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     }
 
     if patch.is_empty() && call_operands_to_move.is_empty() {
-        return;
+        return false;
     }
+    let eliminated = !patch.is_empty();
 
     let bbs = body.basic_blocks.as_mut_preserves_cfg();
     for (Location { block, statement_index }, drop_debuginfo) in patch {
@@ -104,6 +106,8 @@ fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         let Operand::Copy(place) = *arg else { bug!() };
         *arg = Operand::Move(place);
     }
+
+    eliminated
 }
 
 pub(super) enum DeadStoreElimination {
@@ -124,7 +128,12 @@ impl<'tcx> crate::MirPass<'tcx> for DeadStoreElimination {
     }
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-        eliminate(tcx, body);
+        if eliminate(tcx, body) {
+            UsedInStmtLocals::new(body).remove_unused_storage_annotations(body);
+            for data in body.basic_blocks.as_mut_preserves_cfg() {
+                data.strip_nops();
+            }
+        }
     }
 
     fn is_required(&self) -> bool {

compiler/rustc_mir_transform/src/inline.rs

@@ -21,7 +21,7 @@ use tracing::{debug, instrument, trace, trace_span};
 
 use crate::cost_checker::{CostChecker, is_call_like};
 use crate::deref_separator::deref_finder;
-use crate::simplify::simplify_cfg;
+use crate::simplify::{UsedInStmtLocals, simplify_cfg};
 use crate::validate::validate_types;
 use crate::{check_inline, util};
 
@@ -935,7 +935,7 @@ fn inline_call<'tcx, I: Inliner<'tcx>>(
         in_cleanup_block: false,
         return_block,
         tcx,
-        always_live_locals: DenseBitSet::new_filled(callee_body.local_decls.len()),
+        always_live_locals: UsedInStmtLocals::new(&callee_body).locals,
     };
 
     // Map all `Local`s, `SourceScope`s and `BasicBlock`s to new ones
@@ -995,6 +995,10 @@ fn inline_call<'tcx, I: Inliner<'tcx>>(
         // people working on rust can build with or without debuginfo while
         // still getting consistent results from the mir-opt tests.
         caller_body.var_debug_info.append(&mut callee_body.var_debug_info);
+    } else {
+        for bb in callee_body.basic_blocks_mut() {
+            bb.drop_debuginfo();
+        }
     }
     caller_body.basic_blocks_mut().append(callee_body.basic_blocks_mut());
 

compiler/rustc_mir_transform/src/simplify.rs

@@ -34,12 +34,12 @@
 //! The normal logic that a program with UB can be changed to do anything does not apply to
 //! pre-"runtime" MIR!
 
+use rustc_index::bit_set::DenseBitSet;
 use rustc_index::{Idx, IndexSlice, IndexVec};
-use rustc_middle::mir::visit::{
-    MutVisitor, MutatingUseContext, NonUseContext, PlaceContext, Visitor,
-};
+use rustc_middle::mir::visit::{MutVisitor, MutatingUseContext, PlaceContext, Visitor};
 use rustc_middle::mir::*;
 use rustc_middle::ty::TyCtxt;
+use rustc_mir_dataflow::debuginfo::debuginfo_locals;
 use rustc_span::DUMMY_SP;
 use smallvec::SmallVec;
 use tracing::{debug, trace};
@@ -344,7 +344,7 @@ impl<'a, 'tcx> CfgSimplifier<'a, 'tcx> {
 
     fn strip_nops(&mut self) {
         for blk in self.basic_blocks.iter_mut() {
-            blk.retain_statements(|stmt| !matches!(stmt.kind, StatementKind::Nop))
+            blk.strip_nops();
         }
     }
 }
@@ -517,28 +517,39 @@ fn make_local_map<V>(
 /// Keeps track of used & unused locals.
 struct UsedLocals {
     increment: bool,
-    arg_count: u32,
     use_count: IndexVec<Local, u32>,
+    always_used: DenseBitSet<Local>,
 }
 
 impl UsedLocals {
     /// Determines which locals are used & unused in the given body.
     fn new(body: &Body<'_>) -> Self {
+        let mut always_used = debuginfo_locals(body);
+        always_used.insert(RETURN_PLACE);
+        for arg in body.args_iter() {
+            always_used.insert(arg);
+        }
         let mut this = Self {
             increment: true,
-            arg_count: body.arg_count.try_into().unwrap(),
             use_count: IndexVec::from_elem(0, &body.local_decls),
+            always_used,
         };
         this.visit_body(body);
         this
     }
 
     /// Checks if local is used.
     ///
-    /// Return place and arguments are always considered used.
+    /// Return place, arguments, var debuginfo are always considered used.
     fn is_used(&self, local: Local) -> bool {
-        trace!("is_used({:?}): use_count: {:?}", local, self.use_count[local]);
-        local.as_u32() <= self.arg_count || self.use_count[local] != 0
+        trace!(
+            "is_used({:?}): use_count: {:?}, always_used: {}",
+            local,
+            self.use_count[local],
+            self.always_used.contains(local)
+        );
+        // To keep things simple, we don't handle debugging information here, these are in DSE.
+        self.always_used.contains(local) || self.use_count[local] != 0
     }
 
     /// Updates the use counts to reflect the removal of given statement.
@@ -583,13 +594,10 @@ impl<'tcx> Visitor<'tcx> for UsedLocals {
             StatementKind::ConstEvalCounter
             | StatementKind::Nop
             | StatementKind::StorageLive(..)
-            | StatementKind::StorageDead(..) => {
-                self.visit_statement_debuginfos(&statement.debuginfos, location);
-            }
+            | StatementKind::StorageDead(..) => {}
 
             StatementKind::Assign(box (ref place, ref rvalue)) => {
                 if rvalue.is_safe_to_remove() {
-                    self.visit_statement_debuginfos(&statement.debuginfos, location);
                     self.visit_lhs(place, location);
                     self.visit_rvalue(rvalue, location);
                 } else {
@@ -600,16 +608,18 @@ impl<'tcx> Visitor<'tcx> for UsedLocals {
             StatementKind::SetDiscriminant { ref place, variant_index: _ }
             | StatementKind::Deinit(ref place)
             | StatementKind::BackwardIncompatibleDropHint { ref place, reason: _ } => {
-                self.visit_statement_debuginfos(&statement.debuginfos, location);
                 self.visit_lhs(place, location);
             }
         }
     }
 
     fn visit_local(&mut self, local: Local, ctx: PlaceContext, _location: Location) {
+        if matches!(ctx, PlaceContext::NonUse(_)) {
+            return;
+        }
         if self.increment {
             self.use_count[local] += 1;
-        } else if ctx != PlaceContext::NonUse(NonUseContext::VarDebugInfo) {
+        } else {
             assert_ne!(self.use_count[local], 0);
             self.use_count[local] -= 1;
         }
@@ -629,28 +639,26 @@ fn remove_unused_definitions_helper(used_locals: &mut UsedLocals, body: &mut Bod
 
         for data in body.basic_blocks.as_mut_preserves_cfg() {
             // Remove unnecessary StorageLive and StorageDead annotations.
-            data.retain_statements(|statement| {
-                let keep = match &statement.kind {
+            for statement in data.statements.iter_mut() {
+                let keep_statement = match &statement.kind {
                     StatementKind::StorageLive(local) | StatementKind::StorageDead(local) => {
                         used_locals.is_used(*local)
                     }
-                    StatementKind::Assign(box (place, _)) => used_locals.is_used(place.local),
-
-                    StatementKind::SetDiscriminant { place, .. }
-                    | StatementKind::BackwardIncompatibleDropHint { place, reason: _ }
-                    | StatementKind::Deinit(place) => used_locals.is_used(place.local),
-                    StatementKind::Nop => false,
-                    _ => true,
+                    StatementKind::Assign(box (place, _))
+                    | StatementKind::SetDiscriminant { box place, .. }
+                    | StatementKind::BackwardIncompatibleDropHint { box place, .. }
+                    | StatementKind::Deinit(box place) => used_locals.is_used(place.local),
+                    _ => continue,
                 };
-
-                if !keep {
-                    trace!("removing statement {:?}", statement);
-                    modified = true;
-                    used_locals.statement_removed(statement);
+                if keep_statement {
+                    continue;
                 }
-
-                keep
-            });
+                trace!("removing statement {:?}", statement);
+                modified = true;
+                used_locals.statement_removed(statement);
+                statement.make_nop(true);
+            }
+            data.strip_nops();
         }
     }
 }
@@ -669,3 +677,42 @@ impl<'tcx> MutVisitor<'tcx> for LocalUpdater<'tcx> {
         *l = self.map[*l].unwrap();
     }
 }
+
+pub(crate) struct UsedInStmtLocals {
+    pub(crate) locals: DenseBitSet<Local>,
+}
+
+impl UsedInStmtLocals {
+    pub(crate) fn new(body: &Body<'_>) -> Self {
+        let mut this = Self { locals: DenseBitSet::new_empty(body.local_decls.len()) };
+        this.visit_body(body);
+        this
+    }
+
+    pub(crate) fn remove_unused_storage_annotations<'tcx>(&self, body: &mut Body<'tcx>) {
+        for data in body.basic_blocks.as_mut_preserves_cfg() {
+            // Remove unnecessary StorageLive and StorageDead annotations.
+            for statement in data.statements.iter_mut() {
+                let keep_statement = match &statement.kind {
+                    StatementKind::StorageLive(local) | StatementKind::StorageDead(local) => {
+                        self.locals.contains(*local)
+                    }
+                    _ => continue,
+                };
+                if keep_statement {
+                    continue;
+                }
+                statement.make_nop(true);
+            }
+        }
+    }
+}
+
+impl<'tcx> Visitor<'tcx> for UsedInStmtLocals {
+    fn visit_local(&mut self, local: Local, context: PlaceContext, _: Location) {
+        if matches!(context, PlaceContext::NonUse(_)) {
+            return;
+        }
+        self.locals.insert(local);
+    }
+}

tests/mir-opt/dead-store-elimination/cycle.cycle.DeadStoreElimination-initial.diff

@@ -19,10 +19,6 @@
 -         _3 = copy _2;
 -         _2 = copy _1;
 -         _1 = copy _5;
-+         nop;
-+         nop;
-+         nop;
-+         nop;
           _4 = cond() -> [return: bb1, unwind continue];
       }
 

tests/mir-opt/dead-store-elimination/ref.rs

@@ -12,7 +12,7 @@ pub fn tuple(v: (i32, &Foo)) -> i32 {
     // CHECK: debug _dead => [[dead:_[0-9]+]];
     // CHECK: bb0:
     // FIXME: Preserve `tmp` for debuginfo, but we can merge it into the debuginfo.
-    // CHECK: [[tmp:_[0-9]+]] = deref_copy (_1.1: &Foo);
+    // CHECK-NEXT: [[tmp:_[0-9]+]] = deref_copy (_1.1: &Foo);
     // CHECK-NEXT: DBG: [[dead]] = &((*[[tmp]]).2: i32)
     let _dead = &v.1.c;
     v.1.a

tests/mir-opt/dead-store-elimination/ref.tuple.DeadStoreElimination-initial.diff

@@ -12,14 +12,13 @@
       }
 
       bb0: {
-          StorageLive(_2);
+-         StorageLive(_2);
           _3 = deref_copy (_1.1: &Foo);
 -         _2 = &((*_3).2: i32);
 +         // DBG: _2 = &((*_3).2: i32)
-+         nop;
           _4 = deref_copy (_1.1: &Foo);
           _0 = copy ((*_4).0: i32);
-          StorageDead(_2);
+-         StorageDead(_2);
           return;
       }
   }

tests/mir-opt/debuginfo/simplifycfg.drop_debuginfo.SimplifyCfg-final.diff

@@ -14,13 +14,11 @@
 - 
 -     bb1: {
 -         // DBG: _3 = &((*_1).0: i32)
--         nop;
 -         goto -> bb2;
 -     }
 - 
 -     bb2: {
           // DBG: _4 = &((*_1).1: i64)
--         nop;
           _0 = copy ((*_1).2: i32);
           return;
       }