Auto merge of #109087 - cjgillot:sparse-bb-clear, r=davidtwco

Only clear written-to locals in ConstProp

This aims to revert the regression in https://github.com/rust-lang/rust/pull/108872

Clearing all locals at the end of each bb is very costly.
This PR goes back to the original implementation of recording which locals are modified.

2 files changed, 61 insertions, 22 deletions

diff --git a/compiler/rustc_mir_transform/src/const_prop.rs b/compiler/rustc_mir_transform/src/const_prop.rs
index de7b8c63fc8..bebd9723740 100644
--- a/compiler/rustc_mir_transform/src/const_prop.rs
+++ b/compiler/rustc_mir_transform/src/const_prop.rs
@@ -4,6 +4,7 @@
 use either::Right;
 
 use rustc_const_eval::const_eval::CheckAlignment;
+use rustc_data_structures::fx::FxHashSet;
 use rustc_hir::def::DefKind;
 use rustc_index::bit_set::BitSet;
 use rustc_index::vec::IndexVec;
@@ -151,12 +152,17 @@ impl<'tcx> MirPass<'tcx> for ConstProp {
 pub struct ConstPropMachine<'mir, 'tcx> {
     /// The virtual call stack.
     stack: Vec<Frame<'mir, 'tcx>>,
+    pub written_only_inside_own_block_locals: FxHashSet<Local>,
     pub can_const_prop: IndexVec<Local, ConstPropMode>,
 }
 
 impl ConstPropMachine<'_, '_> {
     pub fn new(can_const_prop: IndexVec<Local, ConstPropMode>) -> Self {
-        Self { stack: Vec::new(), can_const_prop }
+        Self {
+            stack: Vec::new(),
+            written_only_inside_own_block_locals: Default::default(),
+            can_const_prop,
+        }
     }
 }
 
@@ -249,7 +255,10 @@ impl<'mir, 'tcx> interpret::Machine<'mir, 'tcx> for ConstPropMachine<'mir, 'tcx>
                     "tried to write to a local that is marked as not propagatable"
                 )
             }
-            ConstPropMode::OnlyInsideOwnBlock | ConstPropMode::FullConstProp => {}
+            ConstPropMode::OnlyInsideOwnBlock => {
+                ecx.machine.written_only_inside_own_block_locals.insert(local);
+            }
+            ConstPropMode::FullConstProp => {}
         }
         ecx.machine.stack[frame].locals[local].access_mut()
     }
@@ -416,6 +425,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
     fn remove_const(ecx: &mut InterpCx<'mir, 'tcx, ConstPropMachine<'mir, 'tcx>>, local: Local) {
         ecx.frame_mut().locals[local].value =
             LocalValue::Live(interpret::Operand::Immediate(interpret::Immediate::Uninit));
+        ecx.machine.written_only_inside_own_block_locals.remove(&local);
     }
 
     /// Returns the value, if any, of evaluating `c`.
@@ -693,7 +703,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         }
     }
 
-    fn ensure_not_propagated(&mut self, local: Local) {
+    fn ensure_not_propagated(&self, local: Local) {
         if cfg!(debug_assertions) {
             assert!(
                 self.get_const(local.into()).is_none()
@@ -963,17 +973,31 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
         // We remove all Locals which are restricted in propagation to their containing blocks and
         // which were modified in the current block.
         // Take it out of the ecx so we can get a mutable reference to the ecx for `remove_const`.
-        let can_const_prop = std::mem::take(&mut self.ecx.machine.can_const_prop);
-        for (local, &mode) in can_const_prop.iter_enumerated() {
-            match mode {
-                ConstPropMode::FullConstProp => {}
-                ConstPropMode::NoPropagation => self.ensure_not_propagated(local),
-                ConstPropMode::OnlyInsideOwnBlock => {
-                    Self::remove_const(&mut self.ecx, local);
-                    self.ensure_not_propagated(local);
+        let mut written_only_inside_own_block_locals =
+            std::mem::take(&mut self.ecx.machine.written_only_inside_own_block_locals);
+
+        // This loop can get very hot for some bodies: it check each local in each bb.
+        // To avoid this quadratic behaviour, we only clear the locals that were modified inside
+        // the current block.
+        for local in written_only_inside_own_block_locals.drain() {
+            debug_assert_eq!(
+                self.ecx.machine.can_const_prop[local],
+                ConstPropMode::OnlyInsideOwnBlock
+            );
+            Self::remove_const(&mut self.ecx, local);
+        }
+        self.ecx.machine.written_only_inside_own_block_locals =
+            written_only_inside_own_block_locals;
+
+        if cfg!(debug_assertions) {
+            for (local, &mode) in self.ecx.machine.can_const_prop.iter_enumerated() {
+                match mode {
+                    ConstPropMode::FullConstProp => {}
+                    ConstPropMode::NoPropagation | ConstPropMode::OnlyInsideOwnBlock => {
+                        self.ensure_not_propagated(local);
+                    }
                 }
             }
         }
-        self.ecx.machine.can_const_prop = can_const_prop;
     }
 }
diff --git a/compiler/rustc_mir_transform/src/const_prop_lint.rs b/compiler/rustc_mir_transform/src/const_prop_lint.rs
index 68e50070e56..45bd98f39d2 100644
--- a/compiler/rustc_mir_transform/src/const_prop_lint.rs
+++ b/compiler/rustc_mir_transform/src/const_prop_lint.rs
@@ -247,6 +247,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
     fn remove_const(ecx: &mut InterpCx<'mir, 'tcx, ConstPropMachine<'mir, 'tcx>>, local: Local) {
         ecx.frame_mut().locals[local].value =
             LocalValue::Live(interpret::Operand::Immediate(interpret::Immediate::Uninit));
+        ecx.machine.written_only_inside_own_block_locals.remove(&local);
     }
 
     fn lint_root(&self, source_info: SourceInfo) -> Option<HirId> {
@@ -484,7 +485,7 @@ impl<'mir, 'tcx> ConstPropagator<'mir, 'tcx> {
         Some(())
     }
 
-    fn ensure_not_propagated(&mut self, local: Local) {
+    fn ensure_not_propagated(&self, local: Local) {
         if cfg!(debug_assertions) {
             assert!(
                 self.get_const(local.into()).is_none()
@@ -691,17 +692,31 @@ impl<'tcx> Visitor<'tcx> for ConstPropagator<'_, 'tcx> {
         // We remove all Locals which are restricted in propagation to their containing blocks and
         // which were modified in the current block.
         // Take it out of the ecx so we can get a mutable reference to the ecx for `remove_const`.
-        let can_const_prop = std::mem::take(&mut self.ecx.machine.can_const_prop);
-        for (local, &mode) in can_const_prop.iter_enumerated() {
-            match mode {
-                ConstPropMode::FullConstProp => {}
-                ConstPropMode::NoPropagation => self.ensure_not_propagated(local),
-                ConstPropMode::OnlyInsideOwnBlock => {
-                    Self::remove_const(&mut self.ecx, local);
-                    self.ensure_not_propagated(local);
+        let mut written_only_inside_own_block_locals =
+            std::mem::take(&mut self.ecx.machine.written_only_inside_own_block_locals);
+
+        // This loop can get very hot for some bodies: it check each local in each bb.
+        // To avoid this quadratic behaviour, we only clear the locals that were modified inside
+        // the current block.
+        for local in written_only_inside_own_block_locals.drain() {
+            debug_assert_eq!(
+                self.ecx.machine.can_const_prop[local],
+                ConstPropMode::OnlyInsideOwnBlock
+            );
+            Self::remove_const(&mut self.ecx, local);
+        }
+        self.ecx.machine.written_only_inside_own_block_locals =
+            written_only_inside_own_block_locals;
+
+        if cfg!(debug_assertions) {
+            for (local, &mode) in self.ecx.machine.can_const_prop.iter_enumerated() {
+                match mode {
+                    ConstPropMode::FullConstProp => {}
+                    ConstPropMode::NoPropagation | ConstPropMode::OnlyInsideOwnBlock => {
+                        self.ensure_not_propagated(local);
+                    }
                 }
             }
         }
-        self.ecx.machine.can_const_prop = can_const_prop;
     }
 }