Auto merge of #100087 - JakobDegen:mir-patch, r=tmiasko

Avoid invalidating the CFG in `MirPatch`

As a part of this change, we adjust `MirPatch` to not needlessly create unnecessary resume blocks.

r? `@tmiasko`

Author: bors

compiler/rustc_middle/src/mir/patch.rs

@@ -11,7 +11,8 @@ pub struct MirPatch<'tcx> {
     new_blocks: Vec<BasicBlockData<'tcx>>,
     new_statements: Vec<(Location, StatementKind<'tcx>)>,
     new_locals: Vec<LocalDecl<'tcx>>,
-    resume_block: BasicBlock,
+    resume_block: Option<BasicBlock>,
+    body_span: Span,
     next_local: usize,
 }
 
@@ -23,47 +24,36 @@ impl<'tcx> MirPatch<'tcx> {
             new_statements: vec![],
             new_locals: vec![],
             next_local: body.local_decls.len(),
-            resume_block: START_BLOCK,
+            resume_block: None,
+            body_span: body.span,
         };
 
-        // make sure the MIR we create has a resume block. It is
-        // completely legal to convert jumps to the resume block
-        // to jumps to None, but we occasionally have to add
-        // instructions just before that.
-
-        let mut resume_block = None;
-        let mut resume_stmt_block = None;
+        // Check if we already have a resume block
         for (bb, block) in body.basic_blocks().iter_enumerated() {
-            if let TerminatorKind::Resume = block.terminator().kind {
-                if !block.statements.is_empty() {
-                    assert!(resume_stmt_block.is_none());
-                    resume_stmt_block = Some(bb);
-                } else {
-                    resume_block = Some(bb);
-                }
+            if let TerminatorKind::Resume = block.terminator().kind && block.statements.is_empty() {
+                result.resume_block = Some(bb);
                 break;
             }
         }
-        let resume_block = resume_block.unwrap_or_else(|| {
-            result.new_block(BasicBlockData {
-                statements: vec![],
-                terminator: Some(Terminator {
-                    source_info: SourceInfo::outermost(body.span),
-                    kind: TerminatorKind::Resume,
-                }),
-                is_cleanup: true,
-            })
-        });
-        result.resume_block = resume_block;
-        if let Some(resume_stmt_block) = resume_stmt_block {
-            result
-                .patch_terminator(resume_stmt_block, TerminatorKind::Goto { target: resume_block });
-        }
+
         result
     }
 
-    pub fn resume_block(&self) -> BasicBlock {
-        self.resume_block
+    pub fn resume_block(&mut self) -> BasicBlock {
+        if let Some(bb) = self.resume_block {
+            return bb;
+        }
+
+        let bb = self.new_block(BasicBlockData {
+            statements: vec![],
+            terminator: Some(Terminator {
+                source_info: SourceInfo::outermost(self.body_span),
+                kind: TerminatorKind::Resume,
+            }),
+            is_cleanup: true,
+        });
+        self.resume_block = Some(bb);
+        bb
     }
 
     pub fn is_patched(&self, bb: BasicBlock) -> bool {
@@ -138,12 +128,17 @@ impl<'tcx> MirPatch<'tcx> {
             self.new_blocks.len(),
             body.basic_blocks().len()
         );
-        body.basic_blocks_mut().extend(self.new_blocks);
+        let bbs = if self.patch_map.is_empty() && self.new_blocks.is_empty() {
+            body.basic_blocks.as_mut_preserves_cfg()
+        } else {
+            body.basic_blocks.as_mut()
+        };
+        bbs.extend(self.new_blocks);
         body.local_decls.extend(self.new_locals);
         for (src, patch) in self.patch_map.into_iter_enumerated() {
             if let Some(patch) = patch {
                 debug!("MirPatch: patching block {:?}", src);
-                body[src].terminator_mut().kind = patch;
+                bbs[src].terminator_mut().kind = patch;
             }
         }
 

compiler/rustc_mir_transform/src/remove_noop_landing_pads.rs

@@ -83,9 +83,9 @@ impl RemoveNoopLandingPads {
     fn remove_nop_landing_pads(&self, body: &mut Body<'_>) {
         debug!("body: {:#?}", body);
 
-        // make sure there's a single resume block
+        // make sure there's a resume block
         let resume_block = {
-            let patch = MirPatch::new(body);
+            let mut patch = MirPatch::new(body);
             let resume_block = patch.resume_block();
             patch.apply(body);
             resume_block

src/test/mir-opt/derefer_complex_case.main.Derefer.diff

@@ -102,10 +102,6 @@
           StorageDead(_6);                 // scope 1 at $DIR/derefer_complex_case.rs:+1:39: +1:40
           _5 = const ();                   // scope 1 at $DIR/derefer_complex_case.rs:+1:5: +1:40
           goto -> bb2;                     // scope 1 at $DIR/derefer_complex_case.rs:+1:5: +1:40
-+     }
-+ 
-+     bb8 (cleanup): {
-+         resume;                          // scope 0 at $DIR/derefer_complex_case.rs:+0:1: +2:2
       }
   }
 

src/test/mir-opt/derefer_terminator_test.main.Derefer.diff

@@ -94,10 +94,6 @@
           StorageDead(_2);                 // scope 1 at $DIR/derefer_terminator_test.rs:+8:1: +8:2
           StorageDead(_1);                 // scope 0 at $DIR/derefer_terminator_test.rs:+8:1: +8:2
           return;                          // scope 0 at $DIR/derefer_terminator_test.rs:+8:2: +8:2
-+     }
-+ 
-+     bb6 (cleanup): {
-+         resume;                          // scope 0 at $DIR/derefer_terminator_test.rs:+0:1: +8:2
       }
   }
 

src/test/mir-opt/derefer_test.main.Derefer.diff

@@ -49,10 +49,6 @@
           StorageDead(_2);                 // scope 1 at $DIR/derefer_test.rs:+5:1: +5:2
           StorageDead(_1);                 // scope 0 at $DIR/derefer_test.rs:+5:1: +5:2
           return;                          // scope 0 at $DIR/derefer_test.rs:+5:2: +5:2
-+     }
-+ 
-+     bb1 (cleanup): {
-+         resume;                          // scope 0 at $DIR/derefer_test.rs:+0:1: +5:2
       }
   }
 

src/test/mir-opt/derefer_test_multiple.main.Derefer.diff

@@ -87,10 +87,6 @@
           StorageDead(_2);                 // scope 1 at $DIR/derefer_test_multiple.rs:+7:1: +7:2
           StorageDead(_1);                 // scope 0 at $DIR/derefer_test_multiple.rs:+7:1: +7:2
           return;                          // scope 0 at $DIR/derefer_test_multiple.rs:+7:2: +7:2
-+     }
-+ 
-+     bb1 (cleanup): {
-+         resume;                          // scope 0 at $DIR/derefer_test_multiple.rs:+0:1: +7:2
       }
   }
 

src/test/mir-opt/inline/dyn_trait.get_query.Inline.diff

@@ -57,10 +57,6 @@
           StorageDead(_4);                 // scope 1 at $DIR/dyn-trait.rs:+2:24: +2:25
           StorageDead(_2);                 // scope 0 at $DIR/dyn-trait.rs:+3:1: +3:2
           return;                          // scope 0 at $DIR/dyn-trait.rs:+3:2: +3:2
-+     }
-+ 
-+     bb3 (cleanup): {
-+         resume;                          // scope 0 at $DIR/dyn-trait.rs:+0:1: +3:2
       }
   }
 

src/test/mir-opt/inline/dyn_trait.try_execute_query.Inline.diff

@@ -32,10 +32,6 @@
 +         StorageDead(_4);                 // scope 1 at $DIR/dyn-trait.rs:+0:21: +0:22
           StorageDead(_2);                 // scope 0 at $DIR/dyn-trait.rs:+1:15: +1:16
           return;                          // scope 0 at $DIR/dyn-trait.rs:+2:2: +2:2
-+     }
-+ 
-+     bb2 (cleanup): {
-+         resume;                          // scope 0 at $DIR/dyn-trait.rs:+0:1: +2:2
       }
   }
 

src/test/mir-opt/inline/inline_any_operand.bar.Inline.after.mir

@@ -41,8 +41,4 @@ fn bar() -> bool {
         StorageDead(_1);                 // scope 0 at $DIR/inline-any-operand.rs:+3:1: +3:2
         return;                          // scope 0 at $DIR/inline-any-operand.rs:+3:2: +3:2
     }
-
-    bb1 (cleanup): {
-        resume;                          // scope 0 at $DIR/inline-any-operand.rs:+0:1: +3:2
-    }
 }

src/test/mir-opt/inline/inline_closure.foo.Inline.after.mir

@@ -46,8 +46,4 @@ fn foo(_1: T, _2: i32) -> i32 {
         StorageDead(_3);                 // scope 0 at $DIR/inline-closure.rs:+3:1: +3:2
         return;                          // scope 0 at $DIR/inline-closure.rs:+3:2: +3:2
     }
-
-    bb1 (cleanup): {
-        resume;                          // scope 0 at $DIR/inline-closure.rs:+0:1: +3:2
-    }
 }