Track drops across multiple yields

5 files changed, 75 insertions, 28 deletions

diff --git a/compiler/rustc_middle/src/middle/region.rs b/compiler/rustc_middle/src/middle/region.rs
index 39ca41c92ff..75dd223d014 100644
--- a/compiler/rustc_middle/src/middle/region.rs
+++ b/compiler/rustc_middle/src/middle/region.rs
@@ -308,7 +308,7 @@ pub struct ScopeTree {
     /// The reason is that semantically, until the `box` expression returns,
     /// the values are still owned by their containing expressions. So
     /// we'll see that `&x`.
-    pub yield_in_scope: FxHashMap<Scope, YieldData>,
+    pub yield_in_scope: FxHashMap<Scope, Vec<YieldData>>,
 
     /// The number of visit_expr and visit_pat calls done in the body.
     /// Used to sanity check visit_expr/visit_pat call count when
@@ -423,8 +423,8 @@ impl ScopeTree {
 
     /// Checks whether the given scope contains a `yield`. If so,
     /// returns `Some(YieldData)`. If not, returns `None`.
-    pub fn yield_in_scope(&self, scope: Scope) -> Option<YieldData> {
-        self.yield_in_scope.get(&scope).cloned()
+    pub fn yield_in_scope(&self, scope: Scope) -> Option<&Vec<YieldData>> {
+        self.yield_in_scope.get(&scope)
     }
 
     /// Gives the number of expressions visited in a body.
diff --git a/compiler/rustc_passes/src/region.rs b/compiler/rustc_passes/src/region.rs
index f8989200d7e..8b22c46f01b 100644
--- a/compiler/rustc_passes/src/region.rs
+++ b/compiler/rustc_passes/src/region.rs
@@ -365,7 +365,8 @@ fn resolve_expr<'tcx>(visitor: &mut RegionResolutionVisitor<'tcx>, expr: &'tcx h
             let target_scopes = visitor.fixup_scopes.drain(start_point..);
 
             for scope in target_scopes {
-                let mut yield_data = visitor.scope_tree.yield_in_scope.get_mut(&scope).unwrap();
+                let mut yield_data =
+                    visitor.scope_tree.yield_in_scope.get_mut(&scope).unwrap().last_mut().unwrap();
                 let count = yield_data.expr_and_pat_count;
                 let span = yield_data.span;
 
@@ -428,7 +429,13 @@ fn resolve_expr<'tcx>(visitor: &mut RegionResolutionVisitor<'tcx>, expr: &'tcx h
             };
             let data =
                 YieldData { span, expr_and_pat_count: visitor.expr_and_pat_count, source: *source };
-            visitor.scope_tree.yield_in_scope.insert(scope, data);
+            match visitor.scope_tree.yield_in_scope.get_mut(&scope) {
+                Some(yields) => yields.push(data),
+                None => {
+                    visitor.scope_tree.yield_in_scope.insert(scope, vec![data]);
+                }
+            }
+
             if visitor.pessimistic_yield {
                 debug!("resolve_expr in pessimistic_yield - marking scope {:?} for fixup", scope);
                 visitor.fixup_scopes.push(scope);
diff --git a/compiler/rustc_typeck/src/check/generator_interior.rs b/compiler/rustc_typeck/src/check/generator_interior.rs
index a406a1a8ecd..ba4080031a2 100644
--- a/compiler/rustc_typeck/src/check/generator_interior.rs
+++ b/compiler/rustc_typeck/src/check/generator_interior.rs
@@ -69,29 +69,29 @@ impl<'a, 'tcx> InteriorVisitor<'a, 'tcx> {
                     //
                     // See the mega-comment at `yield_in_scope` for a proof.
 
-                    debug!(
-                        "comparing counts yield: {} self: {}, source_span = {:?}",
-                        yield_data.expr_and_pat_count, self.expr_count, source_span
-                    );
-
-                    match self.drop_ranges.get(&hir_id) {
-                        Some(range) if range.contains(yield_data.expr_and_pat_count) => {
-                            debug!("value is dropped at yield point; not recording");
-                            return None
-                        }
-                        _ => (),
-                    }
-
-                    // If it is a borrowing happening in the guard,
-                    // it needs to be recorded regardless because they
-                    // do live across this yield point.
-                    if guard_borrowing_from_pattern
-                        || yield_data.expr_and_pat_count >= self.expr_count
-                    {
-                        Some(yield_data)
-                    } else {
-                        None
-                    }
+                    yield_data
+                        .iter()
+                        .find(|yield_data| {
+                            debug!(
+                                "comparing counts yield: {} self: {}, source_span = {:?}",
+                                yield_data.expr_and_pat_count, self.expr_count, source_span
+                            );
+
+                            match self.drop_ranges.get(&hir_id) {
+                                Some(range) if range.contains(yield_data.expr_and_pat_count) => {
+                                    debug!("value is dropped at yield point; not recording");
+                                    return false;
+                                }
+                                _ => (),
+                            }
+
+                            // If it is a borrowing happening in the guard,
+                            // it needs to be recorded regardless because they
+                            // do live across this yield point.
+                            guard_borrowing_from_pattern
+                                || yield_data.expr_and_pat_count >= self.expr_count
+                        })
+                        .cloned()
                 })
             })
             .unwrap_or_else(|| {
diff --git a/src/test/ui/generator/drop-yield-twice.rs b/src/test/ui/generator/drop-yield-twice.rs
new file mode 100644
index 00000000000..f484cbb8d67
--- /dev/null
+++ b/src/test/ui/generator/drop-yield-twice.rs
@@ -0,0 +1,15 @@
+#![feature(negative_impls, generators)]
+
+struct Foo(i32);
+impl !Send for Foo {}
+
+fn main() {
+    assert_send(|| { //~ ERROR generator cannot be sent between threads safely
+        let guard = Foo(42);
+        yield;
+        drop(guard);
+        yield;
+    })
+}
+
+fn assert_send<T: Send>(_: T) {}
diff --git a/src/test/ui/generator/drop-yield-twice.stderr b/src/test/ui/generator/drop-yield-twice.stderr
new file mode 100644
index 00000000000..f821f2f4005
--- /dev/null
+++ b/src/test/ui/generator/drop-yield-twice.stderr
@@ -0,0 +1,25 @@
+error: generator cannot be sent between threads safely
+  --> $DIR/drop-yield-twice.rs:7:5
+   |
+LL |     assert_send(|| {
+   |     ^^^^^^^^^^^ generator is not `Send`
+   |
+   = help: within `[generator@$DIR/drop-yield-twice.rs:7:17: 12:6]`, the trait `Send` is not implemented for `Foo`
+note: generator is not `Send` as this value is used across a yield
+  --> $DIR/drop-yield-twice.rs:9:9
+   |
+LL |         let guard = Foo(42);
+   |             ----- has type `Foo` which is not `Send`
+LL |         yield;
+   |         ^^^^^ yield occurs here, with `guard` maybe used later
+...
+LL |     })
+   |     - `guard` is later dropped here
+note: required by a bound in `assert_send`
+  --> $DIR/drop-yield-twice.rs:15:19
+   |
+LL | fn assert_send<T: Send>(_: T) {}
+   |                   ^^^^ required by this bound in `assert_send`
+
+error: aborting due to previous error
+