Store map of region variable ids to lookups in nll visitor and remove reference to mir

2 files changed, 46 insertions, 26 deletions

diff --git a/src/librustc/mir/visit.rs b/src/librustc/mir/visit.rs
index fd3a9f8cd2d..2b3bb098a80 100644
--- a/src/librustc/mir/visit.rs
+++ b/src/librustc/mir/visit.rs
@@ -741,6 +741,7 @@ macro_rules! make_mir_visitor {
 make_mir_visitor!(Visitor,);
 make_mir_visitor!(MutVisitor,mut);
 
+#[derive(Copy, Clone, Debug)]
 pub enum Lookup {
     Loc(Location),
     Src(SourceInfo),
diff --git a/src/librustc_mir/transform/nll.rs b/src/librustc_mir/transform/nll.rs
index 0793bdabf0b..9699495eb0e 100644
--- a/src/librustc_mir/transform/nll.rs
+++ b/src/librustc_mir/transform/nll.rs
@@ -9,58 +9,78 @@
 // except according to those terms.
 
 use rustc::ty::TypeFoldable;
-use rustc::ty::subst::Substs;
-use rustc::ty::{Ty, TyCtxt, ClosureSubsts};
+use rustc::ty::subst::{Kind, Substs};
+use rustc::ty::{Ty, TyCtxt, ClosureSubsts, RegionVid, RegionKind};
 use rustc::mir::{Mir, Location, Rvalue, BasicBlock, Statement, StatementKind};
 use rustc::mir::visit::{MutVisitor, Lookup};
 use rustc::mir::transform::{MirPass, MirSource};
 use rustc::infer::{self, InferCtxt};
-use syntax_pos::Span;
+use syntax_pos::DUMMY_SP;
+use std::collections::HashMap;
 
 #[allow(dead_code)]
 struct NLLVisitor<'a, 'gcx: 'a + 'tcx, 'tcx: 'a> {
+    pub lookup_map: HashMap<RegionVid, Lookup>,
     infcx: InferCtxt<'a, 'gcx, 'tcx>,
-    source: &'a Mir<'tcx>
 }
 
 impl<'a, 'gcx, 'tcx> NLLVisitor<'a, 'gcx, 'tcx> {
-    pub fn new(infcx: InferCtxt<'a, 'gcx, 'tcx>, source: &'a Mir<'tcx>) -> Self {
+    pub fn new(infcx: InferCtxt<'a, 'gcx, 'tcx>) -> Self {
         NLLVisitor {
             infcx: infcx,
-            source: source,
+            lookup_map: HashMap::new(),
         }
     }
 
-    fn renumber_regions<T>(&self, value: &T, span: Span) -> T where T: TypeFoldable<'tcx> {
+    fn renumber_regions<T>(&self, value: &T) -> T where T: TypeFoldable<'tcx> {
         self.infcx.tcx.fold_regions(value, &mut false, |_region, _depth| {
-            self.infcx.next_region_var(infer::MiscVariable(span))
+            self.infcx.next_region_var(infer::MiscVariable(DUMMY_SP))
         })
     }
-}
 
-fn span_from_location<'tcx>(source: &Mir<'tcx>, location: Location) -> Span {
-    source[location.block].statements[location.statement_index].source_info.span
+    fn store_region(&mut self, region: &RegionKind, lookup: Lookup) {
+        if let RegionKind::ReVar(rid) = *region {
+            self.lookup_map.entry(rid).or_insert(lookup);
+        }
+    }
+
+    fn store_ty_regions(&mut self, ty: &Ty<'tcx>, lookup: Lookup) {
+        for region in ty.regions() {
+            self.store_region(region, lookup);
+        }
+    }
+
+    fn store_kind_regions(&mut self, kind: &'tcx Kind, lookup: Lookup) {
+        if let Some(ty) = kind.as_type() {
+            self.store_ty_regions(&ty, lookup);
+        } else if let Some(region) = kind.as_region() {
+            self.store_region(region, lookup);
+        }
+    }
 }
 
 impl<'a, 'gcx, 'tcx> MutVisitor<'tcx> for NLLVisitor<'a, 'gcx, 'tcx> {
     fn visit_ty(&mut self, ty: &mut Ty<'tcx>, lookup: Lookup) {
         let old_ty = *ty;
-        let span = match lookup {
-            Lookup::Loc(location) => span_from_location(self.source, location),
-            Lookup::Src(source_info) => source_info.span,
-        };
-        *ty = self.renumber_regions(&old_ty, span);
+        *ty = self.renumber_regions(&old_ty);
+        self.store_ty_regions(ty, lookup);
     }
 
     fn visit_substs(&mut self, substs: &mut &'tcx Substs<'tcx>, location: Location) {
-        *substs = self.renumber_regions(&{*substs}, span_from_location(self.source, location));
+        *substs = self.renumber_regions(&{*substs});
+        let lookup = Lookup::Loc(location);
+        for kind in *substs {
+            self.store_kind_regions(kind, lookup);
+        }
     }
 
     fn visit_rvalue(&mut self, rvalue: &mut Rvalue<'tcx>, location: Location) {
         match *rvalue {
             Rvalue::Ref(ref mut r, _, _) => {
                 let old_r = *r;
-                *r = self.renumber_regions(&old_r, span_from_location(self.source, location));
+                *r = self.renumber_regions(&old_r);
+                let lookup = Lookup::Loc(location);
+                self.store_region(r, lookup);
             }
             Rvalue::Use(..) |
             Rvalue::Repeat(..) |
@@ -81,7 +101,11 @@ impl<'a, 'gcx, 'tcx> MutVisitor<'tcx> for NLLVisitor<'a, 'gcx, 'tcx> {
     fn visit_closure_substs(&mut self,
                             substs: &mut ClosureSubsts<'tcx>,
                             location: Location) {
-        *substs = self.renumber_regions(substs, span_from_location(self.source, location));
+        *substs = self.renumber_regions(substs);
+        let lookup = Lookup::Loc(location);
+        for kind in substs.substs {
+            self.store_kind_regions(kind, lookup);
+        }
     }
 
     fn visit_statement(&mut self,
@@ -108,14 +132,9 @@ impl MirPass for NLL {
         }
 
         tcx.infer_ctxt().enter(|infcx| {
-            // Clone mir so we can mutate it without disturbing the rest
-            // of the compiler
+            // Clone mir so we can mutate it without disturbing the rest of the compiler
             let mut renumbered_mir = mir.clone();
-
-            // Note that we're using the passed-in mir for the visitor. This is
-            // so we can lookup locations during traversal without worrying about
-            // maintaing both a mutable and immutable reference to the same object
-            let mut visitor = NLLVisitor::new(infcx, &mir);
+            let mut visitor = NLLVisitor::new(infcx);
             visitor.visit_mir(&mut renumbered_mir);
         })
     }