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`

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>;

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 {

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);
         }
     }

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);
     }
 

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 {

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);
         }
 

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);
     }
 }
 

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 {