Auto merge of #58902 - matthewjasper:generator-cleanup-blocks, r=davidtwco

Fixes for the generator transform

* Moves cleanup annotations in pretty printed MIR so that they can be tested
* Correctly determines which drops are in cleanup blocks when elaborating generator drops
* Use the correct state for poisoning a generator

Closes #58892

13 files changed, 129 insertions, 60 deletions

diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index 33645b57589..b22258a49b2 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -26,7 +26,7 @@
 //!     }
 //!
 //! This pass computes the meaning of the state field and the MIR locals which are live
-//! across a suspension point. There are however two hardcoded generator states:
+//! across a suspension point. There are however three hardcoded generator states:
 //!     0 - Generator have not been resumed yet
 //!     1 - Generator has returned / is completed
 //!     2 - Generator has been poisoned
@@ -144,6 +144,13 @@ fn self_arg() -> Local {
     Local::new(1)
 }
 
+/// Generator have not been resumed yet
+const UNRESUMED: u32 = 0;
+/// Generator has returned / is completed
+const RETURNED: u32 = 1;
+/// Generator has been poisoned
+const POISONED: u32 = 2;
+
 struct SuspensionPoint {
     state: u32,
     resume: BasicBlock,
@@ -278,7 +285,7 @@ impl<'a, 'tcx> MutVisitor<'tcx> for TransformVisitor<'a, 'tcx> {
 
                 state
             } else { // Return
-                 1 // state for returned
+                RETURNED // state for returned
             };
             data.statements.push(self.set_state(state, source_info));
             data.terminator.as_mut().unwrap().kind = TerminatorKind::Return;
@@ -590,8 +597,15 @@ fn elaborate_generator_drops<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
     let param_env = tcx.param_env(def_id);
     let gen = self_arg();
 
-    for block in mir.basic_blocks().indices() {
-        let (target, unwind, source_info) = match mir.basic_blocks()[block].terminator() {
+    let mut elaborator = DropShimElaborator {
+        mir: mir,
+        patch: MirPatch::new(mir),
+        tcx,
+        param_env
+    };
+
+    for (block, block_data) in mir.basic_blocks().iter_enumerated() {
+        let (target, unwind, source_info) = match block_data.terminator() {
             &Terminator {
                 source_info,
                 kind: TerminatorKind::Drop {
@@ -602,31 +616,22 @@ fn elaborate_generator_drops<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>,
             } if local == gen => (target, unwind, source_info),
             _ => continue,
         };
-        let unwind = if let Some(unwind) = unwind {
-            Unwind::To(unwind)
-        } else {
+        let unwind = if block_data.is_cleanup {
             Unwind::InCleanup
+        } else {
+            Unwind::To(unwind.unwrap_or_else(|| elaborator.patch.resume_block()))
         };
-        let patch = {
-            let mut elaborator = DropShimElaborator {
-                mir: &mir,
-                patch: MirPatch::new(mir),
-                tcx,
-                param_env
-            };
-            elaborate_drop(
-                &mut elaborator,
-                source_info,
-                &Place::Base(PlaceBase::Local(gen)),
-                (),
-                target,
-                unwind,
-                block
-            );
-            elaborator.patch
-        };
-        patch.apply(mir);
+        elaborate_drop(
+            &mut elaborator,
+            source_info,
+            &Place::Base(PlaceBase::Local(gen)),
+            (),
+            target,
+            unwind,
+            block,
+        );
     }
+    elaborator.patch.apply(mir);
 }
 
 fn create_generator_drop_shim<'a, 'tcx>(
@@ -643,10 +648,10 @@ fn create_generator_drop_shim<'a, 'tcx>(
 
     let mut cases = create_cases(&mut mir, transform, |point| point.drop);
 
-    cases.insert(0, (0, drop_clean));
+    cases.insert(0, (UNRESUMED, drop_clean));
 
-    // The returned state (1) and the poisoned state (2) falls through to
-    // the default case which is just to return
+    // The returned state and the poisoned state fall through to the default
+    // case which is just to return
 
     insert_switch(tcx, &mut mir, cases, &transform, TerminatorKind::Return);
 
@@ -762,7 +767,7 @@ fn create_generator_resume_function<'a, 'tcx>(
     for block in mir.basic_blocks_mut() {
         let source_info = block.terminator().source_info;
         if let &TerminatorKind::Resume = &block.terminator().kind {
-            block.statements.push(transform.set_state(1, source_info));
+            block.statements.push(transform.set_state(POISONED, source_info));
         }
     }
 
@@ -773,12 +778,12 @@ fn create_generator_resume_function<'a, 'tcx>(
         GeneratorResumedAfterReturn,
     };
 
-    // Jump to the entry point on the 0 state
-    cases.insert(0, (0, BasicBlock::new(0)));
-    // Panic when resumed on the returned (1) state
-    cases.insert(1, (1, insert_panic_block(tcx, mir, GeneratorResumedAfterReturn)));
-    // Panic when resumed on the poisoned (2) state
-    cases.insert(2, (2, insert_panic_block(tcx, mir, GeneratorResumedAfterPanic)));
+    // Jump to the entry point on the unresumed
+    cases.insert(0, (UNRESUMED, BasicBlock::new(0)));
+    // Panic when resumed on the returned state
+    cases.insert(1, (RETURNED, insert_panic_block(tcx, mir, GeneratorResumedAfterReturn)));
+    // Panic when resumed on the poisoned state
+    cases.insert(2, (POISONED, insert_panic_block(tcx, mir, GeneratorResumedAfterPanic)));
 
     insert_switch(tcx, mir, cases, &transform, TerminatorKind::Unreachable);
 
@@ -942,7 +947,7 @@ impl MirPass for StateTransform {
         mir.generator_layout = Some(layout);
 
         // Insert `drop(generator_struct)` which is used to drop upvars for generators in
-        // the unresumed (0) state.
+        // the unresumed state.
         // This is expanded to a drop ladder in `elaborate_generator_drops`.
         let drop_clean = insert_clean_drop(mir);
 
diff --git a/src/librustc_mir/util/pretty.rs b/src/librustc_mir/util/pretty.rs
index 4663f90eb3e..13bcdc26a5e 100644
--- a/src/librustc_mir/util/pretty.rs
+++ b/src/librustc_mir/util/pretty.rs
@@ -317,9 +317,8 @@ where
     let data = &mir[block];
 
     // Basic block label at the top.
-    let cleanup_text = if data.is_cleanup { " // cleanup" } else { "" };
-    let lbl = format!("{}{:?}: {{", INDENT, block);
-    writeln!(w, "{0:1$}{2}", lbl, ALIGN, cleanup_text)?;
+    let cleanup_text = if data.is_cleanup { " (cleanup)" } else { "" };
+    writeln!(w, "{}{:?}{}: {{", INDENT, block, cleanup_text)?;
 
     // List of statements in the middle.
     let mut current_location = Location {
diff --git a/src/test/mir-opt/basic_assignment.rs b/src/test/mir-opt/basic_assignment.rs
index 3ce43cc4a22..c771013f728 100644
--- a/src/test/mir-opt/basic_assignment.rs
+++ b/src/test/mir-opt/basic_assignment.rs
@@ -48,7 +48,7 @@ fn main() {
 //        drop(_6) -> [return: bb6, unwind: bb4];
 //    }
 //    ...
-//    bb5: {
+//    bb5 (cleanup): {
 //        drop(_6) -> bb4;
 //    }
 // END rustc.main.SimplifyCfg-initial.after.mir
diff --git a/src/test/mir-opt/box_expr.rs b/src/test/mir-opt/box_expr.rs
index ad5cf42029e..14d302f0eea 100644
--- a/src/test/mir-opt/box_expr.rs
+++ b/src/test/mir-opt/box_expr.rs
@@ -38,7 +38,7 @@ impl Drop for S {
 //         (*_2) = const S::new() -> [return: bb2, unwind: bb3];
 //     }
 //
-//     bb1: {
+//     bb1 (cleanup): {
 //         resume;
 //     }
 //
@@ -47,7 +47,7 @@ impl Drop for S {
 //         drop(_2) -> bb4;
 //     }
 //
-//     bb3: {
+//     bb3 (cleanup): {
 //         drop(_2) -> bb1;
 //     }
 //
@@ -62,11 +62,11 @@ impl Drop for S {
 //         drop(_4) -> [return: bb8, unwind: bb6];
 //     }
 //
-//     bb6: {
+//     bb6 (cleanup): {
 //         drop(_1) -> bb1;
 //     }
 //
-//     bb7: {
+//     bb7 (cleanup): {
 //         drop(_4) -> bb6;
 //     }
 //
diff --git a/src/test/mir-opt/generator-drop-cleanup.rs b/src/test/mir-opt/generator-drop-cleanup.rs
new file mode 100644
index 00000000000..48398691271
--- /dev/null
+++ b/src/test/mir-opt/generator-drop-cleanup.rs
@@ -0,0 +1,43 @@
+#![feature(generators, generator_trait)]
+
+// Regression test for #58892, generator drop shims should not have blocks
+// spuriously marked as cleanup
+
+fn main() {
+    let gen = || {
+        yield;
+    };
+}
+
+// END RUST SOURCE
+
+// START rustc.main-{{closure}}.generator_drop.0.mir
+// bb0: {
+//     switchInt(((*_1).0: u32)) -> [0u32: bb4, 3u32: bb7, otherwise: bb8];
+// }
+// bb1: {
+//     goto -> bb5;
+// }
+// bb2: {
+//     return;
+// }
+// bb3: {
+//     return;
+// }
+// bb4: {
+//     goto -> bb6;
+// }
+// bb5: {
+//     goto -> bb2;
+// }
+// bb6: {
+//     goto -> bb3;
+// }
+// bb7: {
+//     StorageLive(_3);
+//     goto -> bb1;
+// }
+// bb8: {
+//     return;
+// }
+// END rustc.main-{{closure}}.generator_drop.0.mir
diff --git a/src/test/mir-opt/issue-38669.rs b/src/test/mir-opt/issue-38669.rs
index 618ee2f7407..047e623941b 100644
--- a/src/test/mir-opt/issue-38669.rs
+++ b/src/test/mir-opt/issue-38669.rs
@@ -18,7 +18,7 @@ fn main() {
 //         FakeRead(ForLet, _1);
 //         goto -> bb2;
 //     }
-//     bb1: {
+//     bb1 (cleanup): {
 //         resume;
 //     }
 //     bb2: {
diff --git a/src/test/mir-opt/issue-49232.rs b/src/test/mir-opt/issue-49232.rs
index 0f0401a55ea..5f4f4ab82af 100644
--- a/src/test/mir-opt/issue-49232.rs
+++ b/src/test/mir-opt/issue-49232.rs
@@ -43,7 +43,7 @@ fn main() {
 //         FakeRead(ForMatchedPlace, _3);
 //         switchInt(_3) -> [false: bb9, otherwise: bb8];
 //     }
-//     bb4: {
+//     bb4 (cleanup): {
 //         resume;
 //     }
 //     bb5: {
diff --git a/src/test/mir-opt/loop_test.rs b/src/test/mir-opt/loop_test.rs
index e44743aa4b7..34891ee70b5 100644
--- a/src/test/mir-opt/loop_test.rs
+++ b/src/test/mir-opt/loop_test.rs
@@ -18,7 +18,7 @@ fn main() {
 // END RUST SOURCE
 // START rustc.main.SimplifyCfg-qualify-consts.after.mir
 //    ...
-//    bb1: { // The cleanup block
+//    bb1 (cleanup): {
 //        resume;
 //    }
 //    ...
diff --git a/src/test/mir-opt/match_false_edges.rs b/src/test/mir-opt/match_false_edges.rs
index 9eeef8570a3..0cbf048697a 100644
--- a/src/test/mir-opt/match_false_edges.rs
+++ b/src/test/mir-opt/match_false_edges.rs
@@ -47,7 +47,7 @@ fn main() {
 //      _3 = discriminant(_2);
 //      switchInt(move _3) -> [0isize: bb4, 1isize: bb2, otherwise: bb7];
 //  }
-//  bb1: {
+//  bb1 (cleanup): {
 //      resume;
 //  }
 //  bb2: {
@@ -116,7 +116,7 @@ fn main() {
 //      _3 = discriminant(_2);
 //      switchInt(move _3) -> [0isize: bb3, 1isize: bb2, otherwise: bb7];
 //  }
-//  bb1: {
+//  bb1 (cleanup): {
 //      resume;
 //  }
 //  bb2: {
@@ -185,7 +185,7 @@ fn main() {
 //      _3 = discriminant(_2);
 //      switchInt(move _3) -> [1isize: bb2, otherwise: bb3];
 //  }
-//  bb1: {
+//  bb1 (cleanup): {
 //      resume;
 //  }
 //  bb2: {
diff --git a/src/test/mir-opt/packed-struct-drop-aligned.rs b/src/test/mir-opt/packed-struct-drop-aligned.rs
index 01402f26156..167a6eb349e 100644
--- a/src/test/mir-opt/packed-struct-drop-aligned.rs
+++ b/src/test/mir-opt/packed-struct-drop-aligned.rs
@@ -38,14 +38,14 @@ impl Drop for Droppy {
 //         _6 = move (_1.0: Aligned);
 //         drop(_6) -> [return: bb4, unwind: bb3];
 //     }
-//     bb1: {
+//     bb1 (cleanup): {
 //         resume;
 //     }
 //     bb2: {
 //         StorageDead(_1);
 //         return;
 //     }
-//     bb3: {
+//     bb3 (cleanup): {
 //         (_1.0: Aligned) = move _4;
 //         drop(_1) -> bb1;
 //     }
diff --git a/src/test/mir-opt/remove_fake_borrows.rs b/src/test/mir-opt/remove_fake_borrows.rs
index 48d1c991b62..144348450a9 100644
--- a/src/test/mir-opt/remove_fake_borrows.rs
+++ b/src/test/mir-opt/remove_fake_borrows.rs
@@ -63,7 +63,7 @@ fn main() {
 //     StorageDead(_8);
 //     return;
 // }
-// bb10: {
+// bb10 (cleanup): {
 //     resume;
 // }
 // END rustc.match_guard.CleanupNonCodegenStatements.before.mir
@@ -114,7 +114,7 @@ fn main() {
 //     StorageDead(_8);
 //     return;
 // }
-// bb10: {
+// bb10 (cleanup): {
 //     resume;
 // }
 // END rustc.match_guard.CleanupNonCodegenStatements.after.mir
diff --git a/src/test/mir-opt/unusual-item-types.rs b/src/test/mir-opt/unusual-item-types.rs
index ced30381fda..ef41373d774 100644
--- a/src/test/mir-opt/unusual-item-types.rs
+++ b/src/test/mir-opt/unusual-item-types.rs
@@ -29,7 +29,7 @@ fn main() {
 //     _0 = const 2i32;
 //     return;
 // }
-// bb1: {
+// bb1 (cleanup): {
 //     resume;
 // }
 // END rustc.{{impl}}-ASSOCIATED_CONSTANT.mir_map.0.mir
@@ -39,7 +39,7 @@ fn main() {
 //     _0 = const 5isize;
 //     return;
 // }
-// bb1: {
+// bb1 (cleanup): {
 //     resume;
 // }
 // END rustc.E-V-{{constant}}.mir_map.0.mir
@@ -51,16 +51,16 @@ fn main() {
 // bb1: {
 //     return;
 // }
-// bb2: {
+// bb2 (cleanup): {
 //     resume;
 // }
 // bb3: {
 //     goto -> bb1;
 // }
-// bb4: {
+// bb4 (cleanup): {
 //     goto -> bb2;
 // }
-// bb5: {
+// bb5 (cleanup): {
 //     drop(((*_1).0: alloc::raw_vec::RawVec<i32>)) -> bb4;
 // }
 // bb6: {
diff --git a/src/test/run-fail/generator-resume-after-panic.rs b/src/test/run-fail/generator-resume-after-panic.rs
new file mode 100644
index 00000000000..910b4903bf6
--- /dev/null
+++ b/src/test/run-fail/generator-resume-after-panic.rs
@@ -0,0 +1,22 @@
+// error-pattern:generator resumed after panicking
+
+// Test that we get the correct message for resuming a panicked generator.
+
+#![feature(generators, generator_trait)]
+
+use std::{
+    ops::Generator,
+    pin::Pin,
+    panic,
+};
+
+fn main() {
+    let mut g = || {
+        panic!();
+        yield;
+    };
+    panic::catch_unwind(panic::AssertUnwindSafe(|| {
+        let x = Pin::new(&mut g).resume();
+    }));
+    Pin::new(&mut g).resume();
+}