Implement a cache for looking up yields

3 files changed, 53 insertions, 20 deletions

diff --git a/src/librustc/middle/region.rs b/src/librustc/middle/region.rs
index 82e2a2114da..1f75d09c317 100644
--- a/src/librustc/middle/region.rs
+++ b/src/librustc/middle/region.rs
@@ -24,6 +24,7 @@ use std::mem;
 use std::rc::Rc;
 use syntax::codemap;
 use syntax::ast;
+use syntax::ast::NodeId;
 use syntax_pos::Span;
 use ty::TyCtxt;
 use ty::maps::Providers;
@@ -1167,29 +1168,56 @@ fn region_maps<'a, 'tcx>(tcx: TyCtxt<'a, 'tcx, 'tcx>, def_id: DefId)
     Rc::new(maps)
 }
 
-struct YieldFinder(Option<Span>);
+struct YieldFinder<'a> {
+    cache: &'a mut FxHashMap<NodeId, Option<Span>>,
+    result: Option<Span>,
+}
+
+impl<'a> YieldFinder<'a> {
+    fn lookup<F: FnOnce(&mut Self)>(&mut self, id: NodeId, f: F) {
+        if let Some(result) = self.cache.get(&id) {
+            self.result = *result;
+            return;
+        }
+        if self.result.is_some() {
+            return;
+        }
+        f(self);
+        self.cache.insert(id, self.result);
+    }
+}
 
-impl<'tcx> Visitor<'tcx> for YieldFinder {
+impl<'a, 'tcx> Visitor<'tcx> for YieldFinder<'a> {
     fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
         NestedVisitorMap::None
     }
 
     fn visit_expr(&mut self, expr: &'tcx hir::Expr) {
         if let hir::ExprYield(..) = expr.node {
-            if self.0.is_none() {
-                self.0 = Some(expr.span);
-            }
+            self.result = Some(expr.span);
+            return;
         }
 
-        intravisit::walk_expr(self, expr);
+        self.lookup(expr.id, |this| {
+            intravisit::walk_expr(this, expr);
+        });
+    }
+
+    fn visit_block(&mut self, block: &'tcx hir::Block) {
+        self.lookup(block.id, |this| {
+            intravisit::walk_block(this, block);
+        });
     }
 }
 
 impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
     /// Checks whether the given code extent contains a `yield`. If so,
-    /// returns `Some(span)` with the span of the "first" yield we find.
-    pub fn yield_in_extent(self, extent: CodeExtent) -> Option<Span> {
-        let mut finder = YieldFinder(None);
+    /// returns `Some(span)` with the span of a yield we found.
+    pub fn yield_in_extent(self, extent: CodeExtent, cache: &mut FxHashMap<NodeId, Option<Span>>) -> Option<Span> {
+        let mut finder = YieldFinder {
+            cache,
+            result: None,
+        };
 
         match extent {
             CodeExtent::DestructionScope(node_id) |
@@ -1199,33 +1227,33 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
                     Node::NodeTraitItem(_) |
                     Node::NodeImplItem(_) => {
                         let body = self.hir.body(self.hir.body_owned_by(node_id));
-                        intravisit::walk_body(&mut finder, body);
+                        finder.visit_body(body);
                     }
-                    Node::NodeExpr(expr) => intravisit::walk_expr(&mut finder, expr),
-                    Node::NodeStmt(stmt) => intravisit::walk_stmt(&mut finder, stmt),
-                    Node::NodeBlock(block) => intravisit::walk_block(&mut finder, block),
+                    Node::NodeExpr(expr) => finder.visit_expr(expr),
+                    Node::NodeStmt(stmt) => finder.visit_stmt(stmt),
+                    Node::NodeBlock(block) => finder.visit_block(block),
                     _ => bug!(),
                 }
             }
 
             CodeExtent::CallSiteScope(body_id) |
             CodeExtent::ParameterScope(body_id) => {
-                intravisit::walk_body(&mut finder, self.hir.body(body_id))
+                finder.visit_body(self.hir.body(body_id))
             }
 
             CodeExtent::Remainder(r) => {
                 if let Node::NodeBlock(block) = self.hir.get(r.block) {
                     for stmt in &block.stmts[(r.first_statement_index as usize + 1)..] {
-                        intravisit::walk_stmt(&mut finder, stmt);
+                        finder.visit_stmt(stmt);
                     }
-                    block.expr.as_ref().map(|e| intravisit::walk_expr(&mut finder, e));
+                    block.expr.as_ref().map(|e| finder.visit_expr(e));
                 } else {
                     bug!()
                 }
             }
         }
 
-        finder.0
+        finder.result
     }
 }
 
diff --git a/src/librustc_borrowck/borrowck/mod.rs b/src/librustc_borrowck/borrowck/mod.rs
index 41771738ac6..ca047b0d488 100644
--- a/src/librustc_borrowck/borrowck/mod.rs
+++ b/src/librustc_borrowck/borrowck/mod.rs
@@ -36,7 +36,7 @@ use rustc::middle::region::{self, RegionMaps};
 use rustc::middle::free_region::RegionRelations;
 use rustc::ty::{self, TyCtxt};
 use rustc::ty::maps::Providers;
-
+use rustc::util::nodemap::FxHashMap;
 use std::fmt;
 use std::rc::Rc;
 use std::hash::{Hash, Hasher};
@@ -827,7 +827,7 @@ impl<'a, 'tcx> BorrowckCtxt<'a, 'tcx> {
                                 // block remainder that starts with
                                 // `let a`) for a yield. We can cite
                                 // that for the user.
-                                self.tcx.yield_in_extent(value_extent)
+                                self.tcx.yield_in_extent(value_extent, &mut FxHashMap())
                             } else {
                                 None
                             }
diff --git a/src/librustc_typeck/check/generator_interior.rs b/src/librustc_typeck/check/generator_interior.rs
index 6753d42b0c5..a5440e6d992 100644
--- a/src/librustc_typeck/check/generator_interior.rs
+++ b/src/librustc_typeck/check/generator_interior.rs
@@ -14,12 +14,16 @@ use rustc::hir::intravisit::{self, Visitor, NestedVisitorMap};
 use rustc::hir::{self, Body, Pat, PatKind, Expr};
 use rustc::middle::region::{RegionMaps, CodeExtent};
 use rustc::ty::Ty;
+use syntax::ast::NodeId;
+use syntax::codemap::Span;
 use std::rc::Rc;
 use super::FnCtxt;
 use util::nodemap::FxHashSet;
+use util::nodemap::FxHashMap;
 
 struct InteriorVisitor<'a, 'gcx: 'a+'tcx, 'tcx: 'a> {
     fcx: &'a FnCtxt<'a, 'gcx, 'tcx>,
+    cache: FxHashMap<NodeId, Option<Span>>,
     types: FxHashSet<Ty<'tcx>>,
     region_maps: Rc<RegionMaps>,
 }
@@ -28,7 +32,7 @@ impl<'a, 'gcx, 'tcx> InteriorVisitor<'a, 'gcx, 'tcx> {
     fn record(&mut self, ty: Ty<'tcx>, scope: Option<CodeExtent>, expr: Option<&'tcx Expr>) {
         use syntax_pos::DUMMY_SP;
 
-        if scope.map(|s| self.fcx.tcx.yield_in_extent(s).is_some()).unwrap_or(true) {
+        if scope.map(|s| self.fcx.tcx.yield_in_extent(s, &mut self.cache).is_some()).unwrap_or(true) {
             if log_enabled!(log::LogLevel::Debug) {
                 if let Some(s) = scope {
                     let span = s.span(&self.fcx.tcx.hir).unwrap_or(DUMMY_SP);
@@ -61,6 +65,7 @@ pub fn resolve_interior<'a, 'gcx, 'tcx>(fcx: &'a FnCtxt<'a, 'gcx, 'tcx>,
     let mut visitor = InteriorVisitor {
         fcx,
         types: FxHashSet(),
+        cache: FxHashMap(),
         region_maps: fcx.tcx.region_maps(def_id),
     };
     intravisit::walk_body(&mut visitor, body);