Auto merge of #100089 - JakobDegen:no-invalidate-visitor, r=tmiasko

Add option to `mir::MutVisitor` to not invalidate CFG.

This also applies that option to some uses of the visitor. I had considered a design more similar to #100087 in which we detect if the CFG needs to be invalidated, but that is more difficult with the visitor API and so I decided against it. Another alternative to this design is to offer an API for "saving" and "restoring" CFG caches across arbitrary code. Such an API is more general, and so we may eventually want it anyway, but it seems overkill for this use case.

r? `@tmiasko`

8 files changed, 101 insertions, 72 deletions

diff --git a/compiler/rustc_middle/src/mir/visit.rs b/compiler/rustc_middle/src/mir/visit.rs
index 89160876401..4a85defb1ed 100644
--- a/compiler/rustc_middle/src/mir/visit.rs
+++ b/compiler/rustc_middle/src/mir/visit.rs
@@ -80,6 +80,8 @@ macro_rules! make_mir_visitor {
                 self.super_body(body);
             }
 
+            extra_body_methods!($($mutability)?);
+
             fn visit_basic_block_data(
                 &mut self,
                 block: BasicBlock,
@@ -287,63 +289,7 @@ macro_rules! make_mir_visitor {
                 &mut self,
                 body: &$($mutability)? Body<'tcx>,
             ) {
-                let span = body.span;
-                if let Some(gen) = &$($mutability)? body.generator {
-                    if let Some(yield_ty) = $(& $mutability)? gen.yield_ty {
-                        self.visit_ty(
-                            yield_ty,
-                            TyContext::YieldTy(SourceInfo::outermost(span))
-                        );
-                    }
-                }
-
-                // for best performance, we want to use an iterator rather
-                // than a for-loop, to avoid calling `body::Body::invalidate` for
-                // each basic block.
-                #[allow(unused_macro_rules)]
-                macro_rules! basic_blocks {
-                    (mut) => (body.basic_blocks_mut().iter_enumerated_mut());
-                    () => (body.basic_blocks().iter_enumerated());
-                }
-                for (bb, data) in basic_blocks!($($mutability)?) {
-                    self.visit_basic_block_data(bb, data);
-                }
-
-                for scope in &$($mutability)? body.source_scopes {
-                    self.visit_source_scope_data(scope);
-                }
-
-                self.visit_ty(
-                    $(& $mutability)? body.return_ty(),
-                    TyContext::ReturnTy(SourceInfo::outermost(body.span))
-                );
-
-                for local in body.local_decls.indices() {
-                    self.visit_local_decl(local, & $($mutability)? body.local_decls[local]);
-                }
-
-                #[allow(unused_macro_rules)]
-                macro_rules! type_annotations {
-                    (mut) => (body.user_type_annotations.iter_enumerated_mut());
-                    () => (body.user_type_annotations.iter_enumerated());
-                }
-
-                for (index, annotation) in type_annotations!($($mutability)?) {
-                    self.visit_user_type_annotation(
-                        index, annotation
-                    );
-                }
-
-                for var_debug_info in &$($mutability)? body.var_debug_info {
-                    self.visit_var_debug_info(var_debug_info);
-                }
-
-                self.visit_span($(& $mutability)? body.span);
-
-                for const_ in &$($mutability)? body.required_consts {
-                    let location = START_BLOCK.start_location();
-                    self.visit_constant(const_, location);
-                }
+                super_body!(self, body, $($mutability, true)?);
             }
 
             fn super_basic_block_data(&mut self,
@@ -982,12 +928,7 @@ macro_rules! make_mir_visitor {
                 body: &$($mutability)? Body<'tcx>,
                 location: Location
             ) {
-                #[allow(unused_macro_rules)]
-                macro_rules! basic_blocks {
-                    (mut) => (body.basic_blocks_mut());
-                    () => (body.basic_blocks());
-                }
-                let basic_block = & $($mutability)? basic_blocks!($($mutability)?)[location.block];
+                let basic_block = & $($mutability)? basic_blocks!(body, $($mutability, true)?)[location.block];
                 if basic_block.statements.len() == location.statement_index {
                     if let Some(ref $($mutability)? terminator) = basic_block.terminator {
                         self.visit_terminator(terminator, location)
@@ -1002,6 +943,94 @@ macro_rules! make_mir_visitor {
     }
 }
 
+macro_rules! basic_blocks {
+    ($body:ident, mut, true) => {
+        $body.basic_blocks.as_mut()
+    };
+    ($body:ident, mut, false) => {
+        $body.basic_blocks.as_mut_preserves_cfg()
+    };
+    ($body:ident,) => {
+        $body.basic_blocks()
+    };
+}
+
+macro_rules! basic_blocks_iter {
+    ($body:ident, mut, $invalidate:tt) => {
+        basic_blocks!($body, mut, $invalidate).iter_enumerated_mut()
+    };
+    ($body:ident,) => {
+        basic_blocks!($body,).iter_enumerated()
+    };
+}
+
+macro_rules! extra_body_methods {
+    (mut) => {
+        fn visit_body_preserves_cfg(&mut self, body: &mut Body<'tcx>) {
+            self.super_body_preserves_cfg(body);
+        }
+
+        fn super_body_preserves_cfg(&mut self, body: &mut Body<'tcx>) {
+            super_body!(self, body, mut, false);
+        }
+    };
+    () => {};
+}
+
+macro_rules! super_body {
+    ($self:ident, $body:ident, $($mutability:ident, $invalidate:tt)?) => {
+        let span = $body.span;
+        if let Some(gen) = &$($mutability)? $body.generator {
+            if let Some(yield_ty) = $(& $mutability)? gen.yield_ty {
+                $self.visit_ty(
+                    yield_ty,
+                    TyContext::YieldTy(SourceInfo::outermost(span))
+                );
+            }
+        }
+
+        for (bb, data) in basic_blocks_iter!($body, $($mutability, $invalidate)?) {
+            $self.visit_basic_block_data(bb, data);
+        }
+
+        for scope in &$($mutability)? $body.source_scopes {
+            $self.visit_source_scope_data(scope);
+        }
+
+        $self.visit_ty(
+            $(& $mutability)? $body.return_ty(),
+            TyContext::ReturnTy(SourceInfo::outermost($body.span))
+        );
+
+        for local in $body.local_decls.indices() {
+            $self.visit_local_decl(local, & $($mutability)? $body.local_decls[local]);
+        }
+
+        #[allow(unused_macro_rules)]
+        macro_rules! type_annotations {
+            (mut) => ($body.user_type_annotations.iter_enumerated_mut());
+            () => ($body.user_type_annotations.iter_enumerated());
+        }
+
+        for (index, annotation) in type_annotations!($($mutability)?) {
+            $self.visit_user_type_annotation(
+                index, annotation
+            );
+        }
+
+        for var_debug_info in &$($mutability)? $body.var_debug_info {
+            $self.visit_var_debug_info(var_debug_info);
+        }
+
+        $self.visit_span($(& $mutability)? $body.span);
+
+        for const_ in &$($mutability)? $body.required_consts {
+            let location = START_BLOCK.start_location();
+            $self.visit_constant(const_, location);
+        }
+    }
+}
+
 macro_rules! visit_place_fns {
     (mut) => {
         fn tcx<'a>(&'a self) -> TyCtxt<'tcx>;
diff --git a/compiler/rustc_mir_transform/src/cleanup_post_borrowck.rs b/compiler/rustc_mir_transform/src/cleanup_post_borrowck.rs
index 611d29a4ee2..3378923c22c 100644
--- a/compiler/rustc_mir_transform/src/cleanup_post_borrowck.rs
+++ b/compiler/rustc_mir_transform/src/cleanup_post_borrowck.rs
@@ -33,7 +33,7 @@ pub struct DeleteNonCodegenStatements<'tcx> {
 impl<'tcx> MirPass<'tcx> for CleanupNonCodegenStatements {
     fn run_pass(&self, tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
         let mut delete = DeleteNonCodegenStatements { tcx };
-        delete.visit_body(body);
+        delete.visit_body_preserves_cfg(body);
         body.user_type_annotations.raw.clear();
 
         for decl in &mut body.local_decls {
diff --git a/compiler/rustc_mir_transform/src/const_prop.rs b/compiler/rustc_mir_transform/src/const_prop.rs
index fbc0a767f07..98016659a05 100644
--- a/compiler/rustc_mir_transform/src/const_prop.rs
+++ b/compiler/rustc_mir_transform/src/const_prop.rs
@@ -951,7 +951,7 @@ impl<'tcx> MutVisitor<'tcx> for ConstPropagator<'_, 'tcx> {
     }
 
     fn visit_body(&mut self, body: &mut Body<'tcx>) {
-        for (bb, data) in body.basic_blocks_mut().iter_enumerated_mut() {
+        for (bb, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() {
             self.visit_basic_block_data(bb, data);
         }
     }
diff --git a/compiler/rustc_mir_transform/src/deref_separator.rs b/compiler/rustc_mir_transform/src/deref_separator.rs
index a00bb16f7ac..87d7b664015 100644
--- a/compiler/rustc_mir_transform/src/deref_separator.rs
+++ b/compiler/rustc_mir_transform/src/deref_separator.rs
@@ -90,7 +90,7 @@ pub fn deref_finder<'tcx>(tcx: TyCtxt<'tcx>, body: &mut Body<'tcx>) {
     let patch = MirPatch::new(body);
     let mut checker = DerefChecker { tcx, patcher: patch, local_decls: body.local_decls.clone() };
 
-    for (bb, data) in body.basic_blocks_mut().iter_enumerated_mut() {
+    for (bb, data) in body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut() {
         checker.visit_basic_block_data(bb, data);
     }
 
diff --git a/compiler/rustc_mir_transform/src/elaborate_box_derefs.rs b/compiler/rustc_mir_transform/src/elaborate_box_derefs.rs
index 44e3945d6fc..76b4cdd2ecd 100644
--- a/compiler/rustc_mir_transform/src/elaborate_box_derefs.rs
+++ b/compiler/rustc_mir_transform/src/elaborate_box_derefs.rs
@@ -116,7 +116,7 @@ impl<'tcx> MirPass<'tcx> for ElaborateBoxDerefs {
                 ElaborateBoxDerefVisitor { tcx, unique_did, nonnull_did, local_decls, patch };
 
             for (block, BasicBlockData { statements, terminator, .. }) in
-                body.basic_blocks.as_mut().iter_enumerated_mut()
+                body.basic_blocks.as_mut_preserves_cfg().iter_enumerated_mut()
             {
                 let mut index = 0;
                 for statement in statements {
diff --git a/compiler/rustc_mir_transform/src/nrvo.rs b/compiler/rustc_mir_transform/src/nrvo.rs
index bb063915f55..42d732730ec 100644
--- a/compiler/rustc_mir_transform/src/nrvo.rs
+++ b/compiler/rustc_mir_transform/src/nrvo.rs
@@ -53,10 +53,10 @@ impl<'tcx> MirPass<'tcx> for RenameReturnPlace {
             def_id, returned_local
         );
 
-        RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body(body);
+        RenameToReturnPlace { tcx, to_rename: returned_local }.visit_body_preserves_cfg(body);
 
         // Clean up the `NOP`s we inserted for statements made useless by our renaming.
-        for block_data in body.basic_blocks_mut() {
+        for block_data in body.basic_blocks.as_mut_preserves_cfg() {
             block_data.statements.retain(|stmt| stmt.kind != mir::StatementKind::Nop);
         }
 
diff --git a/compiler/rustc_mir_transform/src/reveal_all.rs b/compiler/rustc_mir_transform/src/reveal_all.rs
index 4919ad40098..abe6cb285f5 100644
--- a/compiler/rustc_mir_transform/src/reveal_all.rs
+++ b/compiler/rustc_mir_transform/src/reveal_all.rs
@@ -19,7 +19,7 @@ impl<'tcx> MirPass<'tcx> for RevealAll {
         }
 
         let param_env = tcx.param_env_reveal_all_normalized(body.source.def_id());
-        RevealAllVisitor { tcx, param_env }.visit_body(body);
+        RevealAllVisitor { tcx, param_env }.visit_body_preserves_cfg(body);
     }
 }
 
diff --git a/compiler/rustc_mir_transform/src/simplify.rs b/compiler/rustc_mir_transform/src/simplify.rs
index 180f4c7dcd6..7a6ca917d0f 100644
--- a/compiler/rustc_mir_transform/src/simplify.rs
+++ b/compiler/rustc_mir_transform/src/simplify.rs
@@ -412,7 +412,7 @@ pub fn simplify_locals<'tcx>(body: &mut Body<'tcx>, tcx: TyCtxt<'tcx>) {
     if map.iter().any(Option::is_none) {
         // Update references to all vars and tmps now
         let mut updater = LocalUpdater { map, tcx };
-        updater.visit_body(body);
+        updater.visit_body_preserves_cfg(body);
 
         body.local_decls.shrink_to_fit();
     }
@@ -548,7 +548,7 @@ fn remove_unused_definitions(used_locals: &mut UsedLocals, body: &mut Body<'_>)
     while modified {
         modified = false;
 
-        for data in body.basic_blocks_mut() {
+        for data in body.basic_blocks.as_mut_preserves_cfg() {
             // Remove unnecessary StorageLive and StorageDead annotations.
             data.statements.retain(|statement| {
                 let keep = match &statement.kind {