Auto merge of #98872 - JakobDegen:no-invalidate, r=davidtwco

Add method to mutate MIR body without invalidating CFG caches.

In addition to adding this method, a handful of passes are updated to use it. There's still quite a few passes that could in principle make use of this as well, but do not at the moment because they use `VisitorMut` or `MirPatch`, which needs additional support for this.

The method name is slightly unwieldy, but I don't expect anyone to be writing it a lot, and at least it says what it does. If anyone has a suggestion for a better name though, would be happy to rename.

r? rust-lang/mir-opt

7 files changed, 47 insertions, 15 deletions

diff --git a/compiler/rustc_middle/src/mir/mod.rs b/compiler/rustc_middle/src/mir/mod.rs
index ad33a54e0a2..11d20b886c9 100644
--- a/compiler/rustc_middle/src/mir/mod.rs
+++ b/compiler/rustc_middle/src/mir/mod.rs
@@ -357,10 +357,7 @@ impl<'tcx> Body<'tcx> {
         //
         // FIXME: Use a finer-grained API for this, so only transformations that alter terminators
         // invalidate the caches.
-        self.predecessor_cache.invalidate();
-        self.switch_source_cache.invalidate();
-        self.is_cyclic.invalidate();
-        self.postorder_cache.invalidate();
+        self.invalidate_cfg_cache();
         &mut self.basic_blocks
     }
 
@@ -368,10 +365,7 @@ impl<'tcx> Body<'tcx> {
     pub fn basic_blocks_and_local_decls_mut(
         &mut self,
     ) -> (&mut IndexVec<BasicBlock, BasicBlockData<'tcx>>, &mut LocalDecls<'tcx>) {
-        self.predecessor_cache.invalidate();
-        self.switch_source_cache.invalidate();
-        self.is_cyclic.invalidate();
-        self.postorder_cache.invalidate();
+        self.invalidate_cfg_cache();
         (&mut self.basic_blocks, &mut self.local_decls)
     }
 
@@ -383,11 +377,43 @@ impl<'tcx> Body<'tcx> {
         &mut LocalDecls<'tcx>,
         &mut Vec<VarDebugInfo<'tcx>>,
     ) {
+        self.invalidate_cfg_cache();
+        (&mut self.basic_blocks, &mut self.local_decls, &mut self.var_debug_info)
+    }
+
+    /// Get mutable access to parts of the Body without invalidating the CFG cache.
+    ///
+    /// By calling this method instead of eg [`Body::basic_blocks_mut`], you promise not to change
+    /// the CFG. This means that
+    ///
+    ///  1) The number of basic blocks remains unchanged
+    ///  2) The set of successors of each terminator remains unchanged.
+    ///  3) For each `TerminatorKind::SwitchInt`, the `targets` remains the same and the terminator
+    ///     kind is not changed.
+    ///
+    /// If any of these conditions cannot be upheld, you should call [`Body::invalidate_cfg_cache`].
+    #[inline]
+    pub fn basic_blocks_local_decls_mut_and_var_debug_info_no_invalidate(
+        &mut self,
+    ) -> (
+        &mut IndexVec<BasicBlock, BasicBlockData<'tcx>>,
+        &mut LocalDecls<'tcx>,
+        &mut Vec<VarDebugInfo<'tcx>>,
+    ) {
+        (&mut self.basic_blocks, &mut self.local_decls, &mut self.var_debug_info)
+    }
+
+    /// Invalidates cached information about the CFG.
+    ///
+    /// You will only ever need this if you have also called
+    /// [`Body::basic_blocks_local_decls_mut_and_var_debug_info_no_invalidate`]. All other methods
+    /// that allow you to mutate the body also call this method themselves, thereby avoiding any
+    /// risk of accidentaly cache invalidation.
+    pub fn invalidate_cfg_cache(&mut self) {
         self.predecessor_cache.invalidate();
         self.switch_source_cache.invalidate();
         self.is_cyclic.invalidate();
         self.postorder_cache.invalidate();
-        (&mut self.basic_blocks, &mut self.local_decls, &mut self.var_debug_info)
     }
 
     /// Returns `true` if a cycle exists in the control-flow graph that is reachable from the
diff --git a/compiler/rustc_mir_transform/src/dead_store_elimination.rs b/compiler/rustc_mir_transform/src/dead_store_elimination.rs
index 28f3790914b..c96497abf8f 100644
--- a/compiler/rustc_mir_transform/src/dead_store_elimination.rs
+++ b/compiler/rustc_mir_transform/src/dead_store_elimination.rs
@@ -66,7 +66,7 @@ pub fn eliminate<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>, borrowed: &BitS
         return;
     }
 
-    let bbs = body.basic_blocks_mut();
+    let bbs = body.basic_blocks_local_decls_mut_and_var_debug_info_no_invalidate().0;
     for Location { block, statement_index } in patch {
         bbs[block].statements[statement_index].make_nop();
     }
diff --git a/compiler/rustc_mir_transform/src/deaggregator.rs b/compiler/rustc_mir_transform/src/deaggregator.rs
index 616ba819982..01f490e23bf 100644
--- a/compiler/rustc_mir_transform/src/deaggregator.rs
+++ b/compiler/rustc_mir_transform/src/deaggregator.rs
@@ -11,7 +11,8 @@ impl<'tcx> MirPass<'tcx> for Deaggregator {
     }
 
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-        let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
+        let (basic_blocks, local_decls, _) =
+            body.basic_blocks_local_decls_mut_and_var_debug_info_no_invalidate();
         let local_decls = &*local_decls;
         for bb in basic_blocks {
             bb.expand_statements(|stmt| {
diff --git a/compiler/rustc_mir_transform/src/lower_slice_len.rs b/compiler/rustc_mir_transform/src/lower_slice_len.rs
index 07163cfe575..813ab4001a7 100644
--- a/compiler/rustc_mir_transform/src/lower_slice_len.rs
+++ b/compiler/rustc_mir_transform/src/lower_slice_len.rs
@@ -26,7 +26,9 @@ pub fn lower_slice_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         return;
     };
 
-    let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
+    // The one successor remains unchanged, so no need to invalidate
+    let (basic_blocks, local_decls, _) =
+        body.basic_blocks_local_decls_mut_and_var_debug_info_no_invalidate();
 
     for block in basic_blocks {
         // lower `<[_]>::len` calls
diff --git a/compiler/rustc_mir_transform/src/normalize_array_len.rs b/compiler/rustc_mir_transform/src/normalize_array_len.rs
index 0f45711baa3..3396a446df2 100644
--- a/compiler/rustc_mir_transform/src/normalize_array_len.rs
+++ b/compiler/rustc_mir_transform/src/normalize_array_len.rs
@@ -32,7 +32,9 @@ impl<'tcx> MirPass<'tcx> for NormalizeArrayLen {
 }
 
 pub fn normalize_array_len_calls<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
-    let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
+    // We don't ever touch terminators, so no need to invalidate the CFG cache
+    let (basic_blocks, local_decls, _) =
+        body.basic_blocks_local_decls_mut_and_var_debug_info_no_invalidate();
 
     // do a preliminary analysis to see if we ever have locals of type `[T;N]` or `&[T;N]`
     let mut interesting_locals = BitSet::new_empty(local_decls.len());
diff --git a/compiler/rustc_mir_transform/src/remove_storage_markers.rs b/compiler/rustc_mir_transform/src/remove_storage_markers.rs
index c9b6e1459d3..5bb4f8bb9b3 100644
--- a/compiler/rustc_mir_transform/src/remove_storage_markers.rs
+++ b/compiler/rustc_mir_transform/src/remove_storage_markers.rs
@@ -17,7 +17,7 @@ impl<'tcx> MirPass<'tcx> for RemoveStorageMarkers {
         }
 
         trace!("Running RemoveStorageMarkers on {:?}", body.source);
-        for data in body.basic_blocks_mut() {
+        for data in body.basic_blocks_local_decls_mut_and_var_debug_info_no_invalidate().0 {
             data.statements.retain(|statement| match statement.kind {
                 StatementKind::StorageLive(..)
                 | StatementKind::StorageDead(..)
diff --git a/compiler/rustc_mir_transform/src/remove_zsts.rs b/compiler/rustc_mir_transform/src/remove_zsts.rs
index aaee6f491cd..34941c1907d 100644
--- a/compiler/rustc_mir_transform/src/remove_zsts.rs
+++ b/compiler/rustc_mir_transform/src/remove_zsts.rs
@@ -18,7 +18,8 @@ impl<'tcx> MirPass<'tcx> for RemoveZsts {
             return;
         }
         let param_env = tcx.param_env(body.source.def_id());
-        let (basic_blocks, local_decls) = body.basic_blocks_and_local_decls_mut();
+        let (basic_blocks, local_decls, _) =
+            body.basic_blocks_local_decls_mut_and_var_debug_info_no_invalidate();
         for block in basic_blocks.iter_mut() {
             for statement in block.statements.iter_mut() {
                 if let StatementKind::Assign(box (place, _)) | StatementKind::Deinit(box place) =