mir-opt: Eliminate dead ref statements

Author: dianqk

compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs

@@ -1146,7 +1146,6 @@ impl<'infcx, 'tcx> MirBorrowckCtxt<'_, 'infcx, 'tcx> {
                 let opt_assignment_rhs_span =
                     self.find_assignments(local).first().map(|&location| {
                         if let Some(mir::Statement {
-                            source_info: _,
                             kind:
                                 mir::StatementKind::Assign(box (
                                     _,

compiler/rustc_middle/src/mir/mod.rs

@@ -4,8 +4,8 @@
 
 use std::borrow::Cow;
 use std::fmt::{self, Debug, Formatter};
-use std::iter;
 use std::ops::{Index, IndexMut};
+use std::{iter, mem};
 
 pub use basic_blocks::{BasicBlocks, SwitchTargetValue};
 use either::Either;
@@ -1342,6 +1342,10 @@ pub struct BasicBlockData<'tcx> {
     /// List of statements in this block.
     pub statements: Vec<Statement<'tcx>>,
 
+    /// All debuginfos happen before the statement.
+    /// Put debuginfos here when the last statement is eliminated.
+    pub after_last_stmt_debuginfos: Vec<StmtDebugInfo<'tcx>>,
+
     /// Terminator for this block.
     ///
     /// N.B., this should generally ONLY be `None` during construction.
@@ -1369,7 +1373,12 @@ impl<'tcx> BasicBlockData<'tcx> {
         terminator: Option<Terminator<'tcx>>,
         is_cleanup: bool,
     ) -> BasicBlockData<'tcx> {
-        BasicBlockData { statements, terminator, is_cleanup }
+        BasicBlockData {
+            statements,
+            after_last_stmt_debuginfos: Vec::new(),
+            terminator,
+            is_cleanup,
+        }
     }
 
     /// Accessor for terminator.
@@ -1404,6 +1413,37 @@ impl<'tcx> BasicBlockData<'tcx> {
             self.terminator().successors()
         }
     }
+
+    pub fn retain_statements<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&Statement<'tcx>) -> bool,
+    {
+        // Fast path without debugging information
+        if self.statements.iter().all(|stmt| stmt.debuginfos.is_empty()) {
+            self.statements.retain_mut(|stmt| f(stmt));
+            return;
+        }
+        let mut debuginfos = Vec::new();
+        self.statements.retain_mut(|stmt| {
+            let retain = f(stmt);
+            if retain {
+                if !debuginfos.is_empty() {
+                    if !stmt.debuginfos.is_empty() {
+                        debuginfos.append(&mut stmt.debuginfos);
+                    }
+                    mem::swap(&mut debuginfos, &mut stmt.debuginfos);
+                }
+            } else {
+                if !stmt.debuginfos.is_empty() {
+                    debuginfos.append(&mut stmt.debuginfos);
+                }
+            }
+            retain
+        });
+        if !debuginfos.is_empty() {
+            self.after_last_stmt_debuginfos.append(&mut debuginfos);
+        }
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////
@@ -1708,10 +1748,10 @@ mod size_asserts {
 
     use super::*;
     // tidy-alphabetical-start
-    static_assert_size!(BasicBlockData<'_>, 128);
+    static_assert_size!(BasicBlockData<'_>, 152);
     static_assert_size!(LocalDecl<'_>, 40);
     static_assert_size!(SourceScopeData<'_>, 64);
-    static_assert_size!(Statement<'_>, 32);
+    static_assert_size!(Statement<'_>, 56);
     static_assert_size!(Terminator<'_>, 96);
     static_assert_size!(VarDebugInfo<'_>, 88);
     // tidy-alphabetical-end

compiler/rustc_middle/src/mir/pretty.rs

@@ -778,6 +778,9 @@ where
     let mut current_location = Location { block, statement_index: 0 };
     for statement in &data.statements {
         extra_data(PassWhere::BeforeLocation(current_location), w)?;
+        for debuginfo in statement.debuginfos.iter() {
+            writeln!(w, "{INDENT}{INDENT}// DBG: {debuginfo:?}")?;
+        }
         let indented_body = format!("{INDENT}{INDENT}{statement:?};");
         if options.include_extra_comments {
             writeln!(
@@ -812,6 +815,9 @@ where
 
     // Terminator at the bottom.
     extra_data(PassWhere::BeforeLocation(current_location), w)?;
+    for debuginfo in data.after_last_stmt_debuginfos.iter() {
+        writeln!(w, "{INDENT}{INDENT}// DBG: {debuginfo:?}")?;
+    }
     if data.terminator.is_some() {
         let indented_terminator = format!("{0}{0}{1:?};", INDENT, data.terminator().kind);
         if options.include_extra_comments {
@@ -891,6 +897,16 @@ impl Debug for Statement<'_> {
     }
 }
 
+impl Debug for StmtDebugInfo<'_> {
+    fn fmt(&self, fmt: &mut Formatter<'_>) -> fmt::Result {
+        match self {
+            StmtDebugInfo::AssignRef(local, place) => {
+                write!(fmt, "{local:?} = &{place:?}")
+            }
+        }
+    }
+}
+
 impl Display for NonDivergingIntrinsic<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {

compiler/rustc_middle/src/mir/statement.rs

@@ -15,17 +15,33 @@ use crate::ty::CoroutineArgsExt;
 pub struct Statement<'tcx> {
     pub source_info: SourceInfo,
     pub kind: StatementKind<'tcx>,
+    pub debuginfos: Vec<StmtDebugInfo<'tcx>>,
 }
 
 impl<'tcx> Statement<'tcx> {
     /// Changes a statement to a nop. This is both faster than deleting instructions and avoids
     /// invalidating statement indices in `Location`s.
-    pub fn make_nop(&mut self) {
-        self.kind = StatementKind::Nop
+    pub fn make_nop(&mut self, drop_debuginfo: bool) {
+        if matches!(self.kind, StatementKind::Nop) {
+            return;
+        }
+        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.")
+                }
+            }
+        }
     }
 
     pub fn new(source_info: SourceInfo, kind: StatementKind<'tcx>) -> Self {
-        Statement { source_info, kind }
+        Statement { source_info, kind, debuginfos: Vec::new() }
     }
 }
 
@@ -939,3 +955,8 @@ impl RawPtrKind {
         }
     }
 }
+
+#[derive(Clone, TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable)]
+pub enum StmtDebugInfo<'tcx> {
+    AssignRef(Local, Place<'tcx>),
+}

compiler/rustc_middle/src/mir/visit.rs

@@ -95,6 +95,22 @@ macro_rules! make_mir_visitor {
                 self.super_source_scope_data(scope_data);
             }
 
+            fn visit_statement_debuginfos(
+                &mut self,
+                stmt_debuginfo: & $($mutability)? [StmtDebugInfo<'tcx>],
+                location: Location
+            ) {
+                self.super_statement_debuginfos(stmt_debuginfo, location);
+            }
+
+            fn visit_statement_debuginfo(
+                &mut self,
+                stmt_debuginfo: & $($mutability)? StmtDebugInfo<'tcx>,
+                location: Location
+            ) {
+                self.super_statement_debuginfo(stmt_debuginfo, location);
+            }
+
             fn visit_statement(
                 &mut self,
                 statement: & $($mutability)? Statement<'tcx>,
@@ -301,6 +317,7 @@ macro_rules! make_mir_visitor {
             {
                 let BasicBlockData {
                     statements,
+                    after_last_stmt_debuginfos,
                     terminator,
                     is_cleanup: _
                 } = data;
@@ -312,8 +329,9 @@ macro_rules! make_mir_visitor {
                     index += 1;
                 }
 
+                let location = Location { block, statement_index: index };
+                self.visit_statement_debuginfos(after_last_stmt_debuginfos, location);
                 if let Some(terminator) = terminator {
-                    let location = Location { block, statement_index: index };
                     self.visit_terminator(terminator, location);
                 }
             }
@@ -376,14 +394,46 @@ macro_rules! make_mir_visitor {
                 }
             }
 
+            fn super_statement_debuginfos(
+                &mut self,
+                stmt_debuginfo: & $($mutability)? [StmtDebugInfo<'tcx>],
+                location: Location
+            ) {
+                for debuginfo in stmt_debuginfo {
+                    self.visit_statement_debuginfo(debuginfo, location);
+                }
+            }
+
+            fn super_statement_debuginfo(
+                &mut self,
+                stmt_debuginfo: & $($mutability)? StmtDebugInfo<'tcx>,
+                location: Location
+            ) {
+                match stmt_debuginfo {
+                    StmtDebugInfo::AssignRef(local, place) => {
+                        self.visit_local(
+                            $(& $mutability)? *local,
+                            PlaceContext::NonUse(NonUseContext::VarDebugInfo),
+                            location
+                        );
+                        self.visit_place(
+                            place,
+                            PlaceContext::NonUse(NonUseContext::VarDebugInfo),
+                            location
+                        );
+                    },
+                }
+            }
+
             fn super_statement(
                 &mut self,
                 statement: & $($mutability)? Statement<'tcx>,
                 location: Location
             ) {
-                let Statement { source_info, kind } = statement;
+                let Statement { source_info, kind, debuginfos } = statement;
 
                 self.visit_source_info(source_info);
+                self.visit_statement_debuginfos(debuginfos, location);
                 match kind {
                     StatementKind::Assign(box (place, rvalue)) => {
                         self.visit_assign(place, rvalue, location);

compiler/rustc_mir_dataflow/src/debuginfo.rs

@@ -5,16 +5,16 @@ use rustc_middle::mir::*;
 /// Return the set of locals that appear in debuginfo.
 pub fn debuginfo_locals(body: &Body<'_>) -> DenseBitSet<Local> {
     let mut visitor = DebuginfoLocals(DenseBitSet::new_empty(body.local_decls.len()));
-    for debuginfo in body.var_debug_info.iter() {
-        visitor.visit_var_debug_info(debuginfo);
-    }
+    visitor.visit_body(body);
     visitor.0
 }
 
 struct DebuginfoLocals(DenseBitSet<Local>);
 
 impl Visitor<'_> for DebuginfoLocals {
-    fn visit_local(&mut self, local: Local, _: PlaceContext, _: Location) {
-        self.0.insert(local);
+    fn visit_local(&mut self, local: Local, place_context: PlaceContext, _: Location) {
+        if place_context == PlaceContext::NonUse(NonUseContext::VarDebugInfo) {
+            self.0.insert(local);
+        }
     }
 }

compiler/rustc_mir_dataflow/src/impls/liveness.rs

@@ -205,15 +205,54 @@ impl DefUse {
 /// All of the caveats of `MaybeLiveLocals` apply.
 pub struct MaybeTransitiveLiveLocals<'a> {
     always_live: &'a DenseBitSet<Local>,
+    debuginfo_locals: &'a DenseBitSet<Local>,
 }
 
 impl<'a> MaybeTransitiveLiveLocals<'a> {
     /// The `always_alive` set is the set of locals to which all stores should unconditionally be
     /// considered 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>,
+        debuginfo_locals: &'a DenseBitSet<Local>,
+    ) -> Self {
+        MaybeTransitiveLiveLocals { always_live, debuginfo_locals }
+    }
+
+    pub fn can_be_removed_if_dead<'tcx>(
+        stmt_kind: &StatementKind<'tcx>,
+        always_live: &DenseBitSet<Local>,
+        debuginfo_locals: &'a DenseBitSet<Local>,
+    ) -> Option<Place<'tcx>> {
+        // 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()
+                && (!debuginfo_locals.contains(place.local)
+                    || (place.as_local().is_some() && matches!(rvalue, mir::Rvalue::Ref(..)))))
+            .then_some(*place),
+            StatementKind::SetDiscriminant { place, .. } | StatementKind::Deinit(place) => {
+                (!debuginfo_locals.contains(place.local)).then_some(**place)
+            }
+            StatementKind::FakeRead(_)
+            | StatementKind::StorageLive(_)
+            | StatementKind::StorageDead(_)
+            | StatementKind::Retag(..)
+            | StatementKind::AscribeUserType(..)
+            | StatementKind::PlaceMention(..)
+            | StatementKind::Coverage(..)
+            | StatementKind::Intrinsic(..)
+            | StatementKind::ConstEvalCounter
+            | StatementKind::BackwardIncompatibleDropHint { .. }
+            | StatementKind::Nop => None,
+        };
+        if let Some(destination) = destination
+            && !destination.is_indirect()
+            && !always_live.contains(destination.local)
+        {
+            return Some(destination);
+        }
+        None
     }
 }
 
@@ -238,32 +277,15 @@ impl<'a, 'tcx> Analysis<'tcx> for MaybeTransitiveLiveLocals<'a> {
         statement: &mir::Statement<'tcx>,
         location: Location,
     ) {
-        // Compute the place that we are storing to, if any
-        let destination = match &statement.kind {
-            StatementKind::Assign(assign) => assign.1.is_safe_to_remove().then_some(assign.0),
-            StatementKind::SetDiscriminant { place, .. } | StatementKind::Deinit(place) => {
-                Some(**place)
-            }
-            StatementKind::FakeRead(_)
-            | StatementKind::StorageLive(_)
-            | StatementKind::StorageDead(_)
-            | StatementKind::Retag(..)
-            | StatementKind::AscribeUserType(..)
-            | StatementKind::PlaceMention(..)
-            | StatementKind::Coverage(..)
-            | StatementKind::Intrinsic(..)
-            | StatementKind::ConstEvalCounter
-            | 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) =
+            Self::can_be_removed_if_dead(&statement.kind, &self.always_live, &self.debuginfo_locals)
+            && !state.contains(destination.local)
+            // FIXME: We can eliminate the statement, but we'll need the statements it depends on
+            // for debuginfos. We need a way to handle this.
+            && !self.debuginfo_locals.contains(destination.local)
+        {
+            // This store is dead
+            return;
         }
         TransferFunction(state).visit_statement(statement, location);
     }

compiler/rustc_mir_transform/src/cleanup_post_borrowck.rs

@@ -36,7 +36,9 @@ impl<'tcx> crate::MirPass<'tcx> for CleanupPostBorrowck {
                         CoverageKind::BlockMarker { .. } | CoverageKind::SpanMarker { .. },
                     )
                     | StatementKind::FakeRead(..)
-                    | StatementKind::BackwardIncompatibleDropHint { .. } => statement.make_nop(),
+                    | StatementKind::BackwardIncompatibleDropHint { .. } => {
+                        statement.make_nop(true)
+                    }
                     StatementKind::Assign(box (
                         _,
                         Rvalue::Cast(