Auto merge of #142771 - dianqk:mir-stmt-debuginfo, r=cjgillot

Introduce debuginfo to statements in MIR

The PR introduces support for debug information within dead statements. Currently, only the reference statement is supported, which is sufficient to fix rust-lang/rust#128081.

I don't modify Stable MIR, as I don't think we need debug information when using it.

This PR represents the debug information for the dead reference statement via `#dbg_value`. For example, `let _foo_b = &foo.b` becomes `#dbg_value(ptr %foo, !22, !DIExpression(DW_OP_plus_uconst, 4, DW_OP_stack_value), !26)`. You can see this here: https://rust.godbolt.org/z/d43js6adv.

The general principle for handling debug information is to never provide less debug information than the optimized LLVM IR.

The current rules for dropping debug information in this PR are:

- If the LLVM IR cannot represent a reference address, it's replaced with poison or simply dropped. For example, see: https://rust.godbolt.org/z/shGqPec8W. I'm using poison in all such cases now.
- All debuginfos is dropped when merging multiple successor BBs. An example is available here: https://rust.godbolt.org/z/TE1q3Wq6M.

I doesn't drop debuginfos in `MatchBranchSimplification`, because LLVM also pick one branch for it.

5 files changed, 226 insertions, 8 deletions

diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index 28142382b13..8eb7aa71fcd 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/compiler/rustc_middle/src/mir/mod.rs
@@ -1298,6 +1298,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: StmtDebugInfos<'tcx>,
+
     /// Terminator for this block.
     ///
     /// N.B., this should generally ONLY be `None` during construction.
@@ -1325,7 +1329,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: StmtDebugInfos::default(),
+            terminator,
+            is_cleanup,
+        }
     }
 
     /// Accessor for terminator.
@@ -1360,6 +1369,36 @@ impl<'tcx> BasicBlockData<'tcx> {
             self.terminator().successors()
         }
     }
+
+    pub fn retain_statements<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&Statement<'tcx>) -> bool,
+    {
+        // Place debuginfos into the next retained statement,
+        // this `debuginfos` variable is used to cache debuginfos between two retained statements.
+        let mut debuginfos = StmtDebugInfos::default();
+        self.statements.retain_mut(|stmt| {
+            let retain = f(stmt);
+            if retain {
+                stmt.debuginfos.prepend(&mut debuginfos);
+            } else {
+                debuginfos.append(&mut stmt.debuginfos);
+            }
+            retain
+        });
+        self.after_last_stmt_debuginfos.prepend(&mut 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.drop_debuginfo();
+        for stmt in self.statements.iter_mut() {
+            stmt.debuginfos.drop_debuginfo();
+        }
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////
@@ -1664,10 +1703,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
diff --git a/compiler/rustc_middle/src/mir/pretty.rs b/compiler/rustc_middle/src/mir/pretty.rs
index 350d75c2ee7..d87e3abe3b2 100644
--- a/compiler/rustc_middle/src/mir/pretty.rs
+++ b/compiler/rustc_middle/src/mir/pretty.rs
@@ -719,6 +719,11 @@ impl<'de, 'tcx> MirWriter<'de, 'tcx> {
         let mut current_location = Location { block, statement_index: 0 };
         for statement in &data.statements {
             (self.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 self.options.include_extra_comments {
                 writeln!(
@@ -749,6 +754,10 @@ impl<'de, 'tcx> MirWriter<'de, 'tcx> {
             current_location.statement_index += 1;
         }
 
+        for debuginfo in data.after_last_stmt_debuginfos.iter() {
+            writeln!(w, "{INDENT}{INDENT}// DBG: {debuginfo:?};")?;
+        }
+
         // Terminator at the bottom.
         (self.extra_data)(PassWhere::BeforeLocation(current_location), w)?;
         if data.terminator.is_some() {
@@ -829,6 +838,19 @@ 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:?}")
+            }
+            StmtDebugInfo::InvalidAssign(local) => {
+                write!(fmt, "{local:?} = &?")
+            }
+        }
+    }
+}
+
 impl Display for NonDivergingIntrinsic<'_> {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         match self {
diff --git a/compiler/rustc_middle/src/mir/statement.rs b/compiler/rustc_middle/src/mir/statement.rs
index e009fe05b53..f310e1e5762 100644
--- a/compiler/rustc_middle/src/mir/statement.rs
+++ b/compiler/rustc_middle/src/mir/statement.rs
@@ -1,5 +1,7 @@
 //! Functionality for statements, operands, places, and things that appear in them.
 
+use std::ops;
+
 use tracing::{debug, instrument};
 
 use super::interpret::GlobalAlloc;
@@ -15,17 +17,28 @@ use crate::ty::CoroutineArgsExt;
 pub struct Statement<'tcx> {
     pub source_info: SourceInfo,
     pub kind: StatementKind<'tcx>,
+    /// Some debuginfos appearing before the primary statement.
+    pub debuginfos: StmtDebugInfos<'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 {
+            let Some(debuginfo) = replaced_stmt.as_debuginfo() else {
+                bug!("debuginfo is not yet supported.")
+            };
+            self.debuginfos.push(debuginfo);
+        }
     }
 
     pub fn new(source_info: SourceInfo, kind: StatementKind<'tcx>) -> Self {
-        Statement { source_info, kind }
+        Statement { source_info, kind, debuginfos: StmtDebugInfos::default() }
     }
 }
 
@@ -63,6 +76,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,
+        }
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////
@@ -503,6 +527,20 @@ impl<'tcx> PlaceRef<'tcx> {
         })
     }
 
+    /// Return the place accessed locals that include the base local.
+    pub fn accessed_locals(self) -> impl Iterator<Item = Local> {
+        std::iter::once(self.local).chain(self.projection.iter().filter_map(|proj| match proj {
+            ProjectionElem::Index(local) => Some(*local),
+            ProjectionElem::Deref
+            | ProjectionElem::Field(_, _)
+            | ProjectionElem::ConstantIndex { .. }
+            | ProjectionElem::Subslice { .. }
+            | ProjectionElem::Downcast(_, _)
+            | ProjectionElem::OpaqueCast(_)
+            | ProjectionElem::UnwrapUnsafeBinder(_) => None,
+        }))
+    }
+
     /// Generates a new place by appending `more_projections` to the existing ones
     /// and interning the result.
     pub fn project_deeper(
@@ -967,3 +1005,71 @@ impl RawPtrKind {
         }
     }
 }
+
+#[derive(Default, Debug, Clone, TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable)]
+pub struct StmtDebugInfos<'tcx>(Vec<StmtDebugInfo<'tcx>>);
+
+impl<'tcx> StmtDebugInfos<'tcx> {
+    pub fn push(&mut self, debuginfo: StmtDebugInfo<'tcx>) {
+        self.0.push(debuginfo);
+    }
+
+    pub fn drop_debuginfo(&mut self) {
+        self.0.clear();
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.0.is_empty()
+    }
+
+    pub fn prepend(&mut self, debuginfos: &mut Self) {
+        if debuginfos.is_empty() {
+            return;
+        };
+        debuginfos.0.append(self);
+        std::mem::swap(debuginfos, self);
+    }
+
+    pub fn append(&mut self, debuginfos: &mut Self) {
+        if debuginfos.is_empty() {
+            return;
+        };
+        self.0.append(debuginfos);
+    }
+
+    pub fn extend(&mut self, debuginfos: &Self) {
+        if debuginfos.is_empty() {
+            return;
+        };
+        self.0.extend_from_slice(debuginfos);
+    }
+
+    pub fn retain<F>(&mut self, f: F)
+    where
+        F: FnMut(&StmtDebugInfo<'tcx>) -> bool,
+    {
+        self.0.retain(f);
+    }
+}
+
+impl<'tcx> ops::Deref for StmtDebugInfos<'tcx> {
+    type Target = Vec<StmtDebugInfo<'tcx>>;
+
+    #[inline]
+    fn deref(&self) -> &Vec<StmtDebugInfo<'tcx>> {
+        &self.0
+    }
+}
+
+impl<'tcx> ops::DerefMut for StmtDebugInfos<'tcx> {
+    #[inline]
+    fn deref_mut(&mut self) -> &mut Vec<StmtDebugInfo<'tcx>> {
+        &mut self.0
+    }
+}
+
+#[derive(Clone, TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable)]
+pub enum StmtDebugInfo<'tcx> {
+    AssignRef(Local, Place<'tcx>),
+    InvalidAssign(Local),
+}
diff --git a/compiler/rustc_middle/src/mir/terminator.rs b/compiler/rustc_middle/src/mir/terminator.rs
index 4034a3a06e9..4249914346c 100644
--- a/compiler/rustc_middle/src/mir/terminator.rs
+++ b/compiler/rustc_middle/src/mir/terminator.rs
@@ -444,6 +444,14 @@ impl<'tcx> Terminator<'tcx> {
         self.kind.successors()
     }
 
+    /// Return `Some` if all successors are identical.
+    #[inline]
+    pub fn identical_successor(&self) -> Option<BasicBlock> {
+        let mut successors = self.successors();
+        let first_succ = successors.next()?;
+        if successors.all(|succ| first_succ == succ) { Some(first_succ) } else { None }
+    }
+
     #[inline]
     pub fn successors_mut<'a>(&'a mut self, f: impl FnMut(&'a mut BasicBlock)) {
         self.kind.successors_mut(f)
diff --git a/compiler/rustc_middle/src/mir/visit.rs b/compiler/rustc_middle/src/mir/visit.rs
index f3923477800..9654e189f2e 100644
--- a/compiler/rustc_middle/src/mir/visit.rs
+++ b/compiler/rustc_middle/src/mir/visit.rs
@@ -95,6 +95,14 @@ macro_rules! make_mir_visitor {
                 self.super_source_scope_data(scope_data);
             }
 
+            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 +309,7 @@ macro_rules! make_mir_visitor {
             {
                 let BasicBlockData {
                     statements,
+                    after_last_stmt_debuginfos,
                     terminator,
                     is_cleanup: _
                 } = data;
@@ -312,8 +321,11 @@ macro_rules! make_mir_visitor {
                     index += 1;
                 }
 
+                let location = Location { block, statement_index: index };
+                for debuginfo in after_last_stmt_debuginfos as & $($mutability)? [_] {
+                    self.visit_statement_debuginfo(debuginfo, location);
+                }
                 if let Some(terminator) = terminator {
-                    let location = Location { block, statement_index: index };
                     self.visit_terminator(terminator, location);
                 }
             }
@@ -376,14 +388,45 @@ macro_rules! make_mir_visitor {
                 }
             }
 
+            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
+                        );
+                    },
+                    StmtDebugInfo::InvalidAssign(local) => {
+                        self.visit_local(
+                            $(& $mutability)? *local,
+                            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);
+                for debuginfo in debuginfos as & $($mutability)? [_] {
+                    self.visit_statement_debuginfo(debuginfo, location);
+                }
                 match kind {
                     StatementKind::Assign(box (place, rvalue)) => {
                         self.visit_assign(place, rvalue, location);