Only consider yields coming after the expressions when computing generator interiors

7 files changed, 85 insertions, 33 deletions

diff --git a/src/librustc/middle/region.rs b/src/librustc/middle/region.rs
index 5b286c6593b..c2d178e2207 100644
--- a/src/librustc/middle/region.rs
+++ b/src/librustc/middle/region.rs
@@ -18,7 +18,6 @@ use ich::{StableHashingContext, NodeIdHashingMode};
 use util::nodemap::{FxHashMap, FxHashSet};
 use ty;
 
-use std::collections::hash_map::Entry;
 use std::mem;
 use std::rc::Rc;
 use syntax::codemap;
@@ -250,8 +249,9 @@ pub struct ScopeTree {
     closure_tree: FxHashMap<hir::ItemLocalId, hir::ItemLocalId>,
 
     /// If there are any `yield` nested within a scope, this map
-    /// stores the `Span` of the first one.
-    yield_in_scope: FxHashMap<Scope, Span>,
+    /// stores the `Span` of the last one and the number of expressions
+    /// which came before it in a generator body.
+    yield_in_scope: FxHashMap<Scope, (Span, usize)>,
 }
 
 #[derive(Debug, Copy, Clone)]
@@ -274,6 +274,9 @@ pub struct Context {
 struct RegionResolutionVisitor<'a, 'tcx: 'a> {
     tcx: TyCtxt<'a, 'tcx, 'tcx>,
 
+    // The number of expressions visited in the current body
+    expr_count: usize,
+
     // Generated scope tree:
     scope_tree: ScopeTree,
 
@@ -611,8 +614,9 @@ impl<'tcx> ScopeTree {
     }
 
     /// Checks whether the given scope contains a `yield`. If so,
-    /// returns `Some(span)` with the span of a yield we found.
-    pub fn yield_in_scope(&self, scope: Scope) -> Option<Span> {
+    /// returns `Some((span, expr_count))` with the span of a yield we found and
+    /// the number of expressions appearing before the `yield` in the body.
+    pub fn yield_in_scope(&self, scope: Scope) -> Option<(Span, usize)> {
         self.yield_in_scope.get(&scope).cloned()
     }
 }
@@ -738,6 +742,8 @@ fn resolve_stmt<'a, 'tcx>(visitor: &mut RegionResolutionVisitor<'a, 'tcx>, stmt:
 fn resolve_expr<'a, 'tcx>(visitor: &mut RegionResolutionVisitor<'a, 'tcx>, expr: &'tcx hir::Expr) {
     debug!("resolve_expr(expr.id={:?})", expr.id);
 
+    visitor.expr_count += 1;
+
     let prev_cx = visitor.cx;
     visitor.enter_node_scope_with_dtor(expr.hir_id.local_id);
 
@@ -808,14 +814,8 @@ fn resolve_expr<'a, 'tcx>(visitor: &mut RegionResolutionVisitor<'a, 'tcx>, expr:
                 // Mark this expr's scope and all parent scopes as containing `yield`.
                 let mut scope = Scope::Node(expr.hir_id.local_id);
                 loop {
-                    match visitor.scope_tree.yield_in_scope.entry(scope) {
-                        // Another `yield` has already been found.
-                        Entry::Occupied(_) => break,
-
-                        Entry::Vacant(entry) => {
-                            entry.insert(expr.span);
-                        }
-                    }
+                    visitor.scope_tree.yield_in_scope.insert(scope,
+                        (expr.span, visitor.expr_count));
 
                     // Keep traversing up while we can.
                     match visitor.scope_tree.parent_map.get(&scope) {
@@ -1120,6 +1120,7 @@ impl<'a, 'tcx> Visitor<'tcx> for RegionResolutionVisitor<'a, 'tcx> {
                body_id,
                self.cx.parent);
 
+        let outer_ec = mem::replace(&mut self.expr_count, 0);
         let outer_cx = self.cx;
         let outer_ts = mem::replace(&mut self.terminating_scopes, FxHashSet());
         self.terminating_scopes.insert(body.value.hir_id.local_id);
@@ -1166,6 +1167,7 @@ impl<'a, 'tcx> Visitor<'tcx> for RegionResolutionVisitor<'a, 'tcx> {
         }
 
         // Restore context we had at the start.
+        self.expr_count = outer_ec;
         self.cx = outer_cx;
         self.terminating_scopes = outer_ts;
     }
@@ -1200,6 +1202,7 @@ fn region_scope_tree<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, def_id: DefId)
         let mut visitor = RegionResolutionVisitor {
             tcx,
             scope_tree: ScopeTree::default(),
+            expr_count: 0,
             cx: Context {
                 root_id: None,
                 parent: None,
diff --git a/src/librustc_borrowck/borrowck/mod.rs b/src/librustc_borrowck/borrowck/mod.rs
index 6fb49a0908f..ef93e0365e6 100644
--- a/src/librustc_borrowck/borrowck/mod.rs
+++ b/src/librustc_borrowck/borrowck/mod.rs
@@ -857,7 +857,7 @@ impl<'a, 'tcx> BorrowckCtxt<'a, 'tcx> {
                     None
                 };
 
-                if let Some(yield_span) = maybe_borrow_across_yield {
+                if let Some((yield_span, _)) = maybe_borrow_across_yield {
                     debug!("err_out_of_scope: opt_yield_span = {:?}", yield_span);
                     struct_span_err!(self.tcx.sess,
                                      error_span,
diff --git a/src/librustc_typeck/check/generator_interior.rs b/src/librustc_typeck/check/generator_interior.rs
index 88219566792..c7971666f44 100644
--- a/src/librustc_typeck/check/generator_interior.rs
+++ b/src/librustc_typeck/check/generator_interior.rs
@@ -15,7 +15,7 @@
 
 use rustc::hir::def_id::DefId;
 use rustc::hir::intravisit::{self, Visitor, NestedVisitorMap};
-use rustc::hir::{self, Body, Pat, PatKind, Expr};
+use rustc::hir::{self, Pat, PatKind, Expr};
 use rustc::middle::region;
 use rustc::ty::Ty;
 use std::rc::Rc;
@@ -26,14 +26,27 @@ struct InteriorVisitor<'a, 'gcx: 'a+'tcx, 'tcx: 'a> {
     fcx: &'a FnCtxt<'a, 'gcx, 'tcx>,
     types: FxHashMap<Ty<'tcx>, usize>,
     region_scope_tree: Rc<region::ScopeTree>,
+    expr_count: usize,
 }
 
 impl<'a, 'gcx, 'tcx> InteriorVisitor<'a, 'gcx, 'tcx> {
     fn record(&mut self, ty: Ty<'tcx>, scope: Option<region::Scope>, expr: Option<&'tcx Expr>) {
         use syntax_pos::DUMMY_SP;
 
+        if self.fcx.tcx.sess.verbose() {
+        let span = scope.map_or(DUMMY_SP, |s| s.span(self.fcx.tcx, &self.region_scope_tree));
+        self.fcx.tcx.sess.span_warn(span, &format!("temporary scope for node id {:?}", expr));
+        }
+
         let live_across_yield = scope.map_or(Some(DUMMY_SP), |s| {
-            self.region_scope_tree.yield_in_scope(s)
+            self.region_scope_tree.yield_in_scope(s).and_then(|(span, expr_count)| {
+                // Check if the span in the region comes after the expression
+                if expr_count > self.expr_count {
+                    Some(span)
+                } else {
+                    None
+                }
+            })
         });
 
         if let Some(span) = live_across_yield {
@@ -60,6 +73,7 @@ pub fn resolve_interior<'a, 'gcx, 'tcx>(fcx: &'a FnCtxt<'a, 'gcx, 'tcx>,
         fcx,
         types: FxHashMap(),
         region_scope_tree: fcx.tcx.region_scope_tree(def_id),
+        expr_count: 0,
     };
     intravisit::walk_body(&mut visitor, body);
 
@@ -82,15 +96,14 @@ pub fn resolve_interior<'a, 'gcx, 'tcx>(fcx: &'a FnCtxt<'a, 'gcx, 'tcx>,
    }
 }
 
+// This visitor has to have the same visit_expr calls as RegionResolutionVisitor in
+// librustc/middle/region.rs since `expr_count` is compared against the results
+// there.
 impl<'a, 'gcx, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'gcx, 'tcx> {
     fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
         NestedVisitorMap::None
     }
 
-    fn visit_body(&mut self, _body: &'tcx Body) {
-        // Closures inside are not considered part of the generator interior
-    }
-
     fn visit_pat(&mut self, pat: &'tcx Pat) {
         if let PatKind::Binding(..) = pat.node {
             let scope = self.region_scope_tree.var_scope(pat.hir_id.local_id);
@@ -102,7 +115,15 @@ impl<'a, 'gcx, 'tcx> Visitor<'tcx> for InteriorVisitor<'a, 'gcx, 'tcx> {
     }
 
     fn visit_expr(&mut self, expr: &'tcx Expr) {
+        self.expr_count += 1;
+
+        if self.fcx.tcx.sess.verbose() {
+        self.fcx.tcx.sess.span_warn(expr.span, &format!("node id {:?}", expr.id));
+        }
+
         let scope = self.region_scope_tree.temporary_scope(expr.hir_id.local_id);
+
+
         let ty = self.fcx.tables.borrow().expr_ty_adjusted(expr);
         self.record(ty, scope, Some(expr));
 
diff --git a/src/test/run-pass/generator/borrow-in-tail-expr.rs b/src/test/run-pass/generator/borrow-in-tail-expr.rs
new file mode 100644
index 00000000000..df1a1dcebe6
--- /dev/null
+++ b/src/test/run-pass/generator/borrow-in-tail-expr.rs
@@ -0,0 +1,19 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![feature(generators)]
+
+fn main() {
+    let _a = || {
+        yield;
+        let a = String::new();
+        a.len()
+    };
+}
\ No newline at end of file
diff --git a/src/test/run-pass/generator/borrow-in-yield-expr.rs b/src/test/run-pass/generator/borrow-in-yield-expr.rs
new file mode 100644
index 00000000000..50981df7566
--- /dev/null
+++ b/src/test/run-pass/generator/borrow-in-yield-expr.rs
@@ -0,0 +1,21 @@
+// Copyright 2017 The Rust Project Developers. See the COPYRIGHT
+// file at the top-level directory of this distribution and at
+// http://rust-lang.org/COPYRIGHT.
+//
+// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
+// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
+// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
+// option. This file may not be copied, modified, or distributed
+// except according to those terms.
+
+#![feature(generators, generator_trait, conservative_impl_trait)]
+
+use std::ops::Generator;
+
+fn bar(baz: String) -> impl Generator<Yield=(), Return=()> {
+    move || {
+        yield drop(&baz);
+    }
+}
+
+fn main() {}
\ No newline at end of file
diff --git a/src/test/ui/generator/yield-in-args-rev.rs b/src/test/run-pass/generator/yield-in-args-rev.rs
index fb0e68136f5..e22759291d1 100644
--- a/src/test/ui/generator/yield-in-args-rev.rs
+++ b/src/test/run-pass/generator/yield-in-args-rev.rs
@@ -12,12 +12,10 @@
 
 fn foo(_a: (), _b: &bool) {}
 
-// Some examples that probably *could* be accepted, but which we reject for now.
-
 fn bar() {
     || {
         let b = true;
-        foo(yield, &b); //~ ERROR
+        foo(yield, &b);
     };
 }
 
diff --git a/src/test/ui/generator/yield-in-args-rev.stderr b/src/test/ui/generator/yield-in-args-rev.stderr
deleted file mode 100644
index 157f8968209..00000000000
--- a/src/test/ui/generator/yield-in-args-rev.stderr
+++ /dev/null
@@ -1,10 +0,0 @@
-error[E0626]: borrow may still be in use when generator yields
-  --> $DIR/yield-in-args-rev.rs:20:21
-   |
-20 |         foo(yield, &b); //~ ERROR
-   |             -----   ^
-   |             |
-   |             possible yield occurs here
-
-error: aborting due to previous error
-