From 8a7801908c6b8f142a0b31f0582e526fb7369833 Mon Sep 17 00:00:00 2001
From: Matthew Jasper <mjjasper1@gmail.com>
Date: Sun, 3 Mar 2019 19:27:41 +0000
Subject: [PATCH] Don't incorrectly mark blocks in generator drop shims as
 cleanup

This also ensure that dropping a generator won't leak upvars if dropping
one of them panics
---
 src/librustc_mir/transform/generator.rs    | 46 +++++++++++-----------
 src/test/mir-opt/generator-drop-cleanup.rs | 43 ++++++++++++++++++++
 2 files changed, 65 insertions(+), 24 deletions(-)
 create mode 100644 src/test/mir-opt/generator-drop-cleanup.rs

diff --git a/src/librustc_mir/transform/generator.rs b/src/librustc_mir/transform/generator.rs
index c455d38cebce8..9c212e36761c0 100644
--- a/src/librustc_mir/transform/generator.rs
+++ b/src/librustc_mir/transform/generator.rs
@@ -592,8 +592,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 {
@@ -604,31 +611,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>(
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 0000000000000..48398691271ba
--- /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