mir-opt: Eliminate dead ref statements

Author: dianqk

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;
@@ -1341,6 +1341,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.
@@ -1368,7 +1372,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.
@@ -1403,6 +1412,23 @@ impl<'tcx> BasicBlockData<'tcx> {
             self.terminator().successors()
         }
     }
+
+    pub fn retain_statements<F>(&mut self, mut f: F)
+    where
+        F: FnMut(&Statement<'tcx>) -> bool,
+    {
+        let mut debuginfos = Vec::new();
+        self.statements.retain_mut(|stmt| {
+            let retain = f(stmt);
+            if retain {
+                stmt.debuginfos.splice(0..0, mem::take(&mut debuginfos));
+            } else {
+                debuginfos.extend_from_slice(&stmt.debuginfos);
+            }
+            retain
+        });
+        self.after_last_stmt_debuginfos.extend_from_slice(&debuginfos);
+    }
 }
 
 ///////////////////////////////////////////////////////////////////////////
@@ -1707,7 +1733,7 @@ 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<'_>, 56);

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 {

compiler/rustc_middle/src/mir/statement.rs

@@ -15,19 +15,33 @@ use crate::ty::CoroutineArgsExt;
 pub struct Statement<'tcx> {
     pub source_info: SourceInfo,
     pub kind: StatementKind<'tcx>,
-    pub debug_info: Vec<StmtDebugInfo>,
+    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) {
-        // TODO: convert the stmt to debuginfo.
-        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, debug_info: Vec::new() }
+        Statement { source_info, kind, debuginfos: Vec::new() }
     }
 }
 
@@ -943,4 +957,6 @@ impl RawPtrKind {
 }
 
 #[derive(Debug, Clone, TyEncodable, TyDecodable, HashStable, TypeFoldable, TypeVisitable)]
-pub struct StmtDebugInfo {}
+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,15 +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
             ) {
-                // TODO: visit debuginfo
-                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/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,12 @@ 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)
+        {
+            // 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(

compiler/rustc_mir_transform/src/copy_prop.rs

@@ -138,7 +138,7 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
         if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = stmt.kind
             && self.storage_to_remove.contains(l)
         {
-            stmt.make_nop();
+            stmt.make_nop(true);
             return;
         }
 
@@ -150,7 +150,7 @@ impl<'tcx> MutVisitor<'tcx> for Replacer<'_, 'tcx> {
                 *rhs
             && lhs == rhs
         {
-            stmt.make_nop();
+            stmt.make_nop(true);
         }
     }
 }

compiler/rustc_mir_transform/src/coroutine.rs

@@ -412,7 +412,7 @@ impl<'tcx> MutVisitor<'tcx> for TransformVisitor<'tcx> {
             if let StatementKind::StorageLive(l) | StatementKind::StorageDead(l) = s.kind
                 && self.remap.contains(l)
             {
-                s.make_nop();
+                s.make_nop(true);
             }
         }