Rollup merge of #56906 - blitzerr:master, r=nikomatsakis

Issue #56905

Adding a map to TypeckTables to get the list of all the Upvars
given a closureID. This is help us get rid of the recurring
pattern in the codebase of iterating over the free vars
using with_freevars.

5 files changed, 81 insertions, 26 deletions

diff --git a/src/librustc/ty/context.rs b/src/librustc/ty/context.rs
index c0ba4329ae0..8d4b8aae8b1 100644
--- a/src/librustc/ty/context.rs
+++ b/src/librustc/ty/context.rs
@@ -417,6 +417,12 @@ pub struct TypeckTables<'tcx> {
     /// All the existential types that are restricted to concrete types
     /// by this function
     pub concrete_existential_types: FxHashMap<DefId, Ty<'tcx>>,
+
+    /// Given the closure ID this map provides the list of UpvarIDs used by it.
+    /// The upvarID contains the HIR node ID and it also contains the full path
+    /// leading to the member of the struct or tuple that is used instead of the
+    /// entire variable.
+    pub upvar_list: ty::UpvarListMap,
 }
 
 impl<'tcx> TypeckTables<'tcx> {
@@ -441,6 +447,7 @@ impl<'tcx> TypeckTables<'tcx> {
             tainted_by_errors: false,
             free_region_map: Default::default(),
             concrete_existential_types: Default::default(),
+            upvar_list: Default::default(),
         }
     }
 
@@ -741,6 +748,8 @@ impl<'a, 'gcx> HashStable<StableHashingContext<'a>> for TypeckTables<'gcx> {
             tainted_by_errors,
             ref free_region_map,
             ref concrete_existential_types,
+            ref upvar_list,
+
         } = *self;
 
         hcx.with_node_id_hashing_mode(NodeIdHashingMode::HashDefPath, |hcx| {
@@ -783,6 +792,7 @@ impl<'a, 'gcx> HashStable<StableHashingContext<'a>> for TypeckTables<'gcx> {
             tainted_by_errors.hash_stable(hcx, hasher);
             free_region_map.hash_stable(hcx, hasher);
             concrete_existential_types.hash_stable(hcx, hasher);
+            upvar_list.hash_stable(hcx, hasher);
         })
     }
 }
diff --git a/src/librustc/ty/mod.rs b/src/librustc/ty/mod.rs
index a2c96e7cf9f..cfd99948e43 100644
--- a/src/librustc/ty/mod.rs
+++ b/src/librustc/ty/mod.rs
@@ -808,6 +808,7 @@ pub struct UpvarBorrow<'tcx> {
     pub region: ty::Region<'tcx>,
 }
 
+pub type UpvarListMap = FxHashMap<DefId, Vec<UpvarId>>;
 pub type UpvarCaptureMap<'tcx> = FxHashMap<UpvarId, UpvarCapture<'tcx>>;
 
 #[derive(Copy, Clone)]
diff --git a/src/librustc_mir/build/mod.rs b/src/librustc_mir/build/mod.rs
index 727b769cf4d..65ae111fbc0 100644
--- a/src/librustc_mir/build/mod.rs
+++ b/src/librustc_mir/build/mod.rs
@@ -4,7 +4,7 @@ use hair::cx::Cx;
 use hair::{LintLevel, BindingMode, PatternKind};
 use rustc::hir;
 use rustc::hir::Node;
-use rustc::hir::def_id::{DefId, LocalDefId};
+use rustc::hir::def_id::DefId;
 use rustc::middle::region;
 use rustc::mir::*;
 use rustc::mir::visit::{MutVisitor, TyContext};
@@ -640,21 +640,29 @@ fn construct_fn<'a, 'gcx, 'tcx, A>(hir: Cx<'a, 'gcx, 'tcx>,
     let arguments: Vec<_> = arguments.collect();
 
     let tcx = hir.tcx();
-    let span = tcx.hir().span(fn_id);
+    let tcx_hir = tcx.hir();
+    let span = tcx_hir.span(fn_id);
+
+    let hir_tables = hir.tables();
+    let fn_def_id = tcx_hir.local_def_id(fn_id);
 
     // Gather the upvars of a closure, if any.
-    let upvar_decls: Vec<_> = tcx.with_freevars(fn_id, |freevars| {
-        freevars.iter().map(|fv| {
-            let var_id = fv.var_id();
-            let var_hir_id = tcx.hir().node_to_hir_id(var_id);
-            let closure_expr_id = tcx.hir().local_def_id(fn_id);
-            let capture = hir.tables().upvar_capture(ty::UpvarId {
-                var_path: ty::UpvarPath {hir_id: var_hir_id},
-                closure_expr_id: LocalDefId::from_def_id(closure_expr_id),
-            });
+    // In analyze_closure() in upvar.rs we gathered a list of upvars used by a
+    // closure and we stored in a map called upvar_list in TypeckTables indexed
+    // with the closure's DefId. Here, we run through that vec of UpvarIds for
+    // the given closure and use the necessary information to create UpvarDecl.
+    let upvar_decls: Vec<_> = hir_tables
+        .upvar_list
+        .get(&fn_def_id)
+        .into_iter()
+        .flatten()
+        .map(|upvar_id| {
+            let var_hir_id = upvar_id.var_path.hir_id;
+            let var_node_id = tcx_hir.hir_to_node_id(var_hir_id);
+            let capture = hir_tables.upvar_capture(*upvar_id);
             let by_ref = match capture {
                 ty::UpvarCapture::ByValue => false,
-                ty::UpvarCapture::ByRef(..) => true
+                ty::UpvarCapture::ByRef(..) => true,
             };
             let mut decl = UpvarDecl {
                 debug_name: keywords::Invalid.name(),
@@ -662,10 +670,9 @@ fn construct_fn<'a, 'gcx, 'tcx, A>(hir: Cx<'a, 'gcx, 'tcx>,
                 by_ref,
                 mutability: Mutability::Not,
             };
-            if let Some(Node::Binding(pat)) = tcx.hir().find(var_id) {
+            if let Some(Node::Binding(pat)) = tcx_hir.find(var_node_id) {
                 if let hir::PatKind::Binding(_, _, ident, _) = pat.node {
                     decl.debug_name = ident.name;
-
                     if let Some(&bm) = hir.tables.pat_binding_modes().get(pat.hir_id) {
                         if bm == ty::BindByValue(hir::MutMutable) {
                             decl.mutability = Mutability::Mut;
@@ -678,8 +685,8 @@ fn construct_fn<'a, 'gcx, 'tcx, A>(hir: Cx<'a, 'gcx, 'tcx>,
                 }
             }
             decl
-        }).collect()
-    });
+        })
+        .collect();
 
     let mut builder = Builder::new(hir,
         span,
@@ -689,7 +696,6 @@ fn construct_fn<'a, 'gcx, 'tcx, A>(hir: Cx<'a, 'gcx, 'tcx>,
         return_ty_span,
         upvar_decls);
 
-    let fn_def_id = tcx.hir().local_def_id(fn_id);
     let call_site_scope = region::Scope {
         id: body.value.hir_id.local_id,
         data: region::ScopeData::CallSite
@@ -732,7 +738,7 @@ fn construct_fn<'a, 'gcx, 'tcx, A>(hir: Cx<'a, 'gcx, 'tcx>,
         // RustCall pseudo-ABI untuples the last argument.
         spread_arg = Some(Local::new(arguments.len()));
     }
-    let closure_expr_id = tcx.hir().local_def_id(fn_id);
+    let closure_expr_id = tcx_hir.local_def_id(fn_id);
     info!("fn_id {:?} has attrs {:?}", closure_expr_id,
           tcx.get_attrs(closure_expr_id));
 
diff --git a/src/librustc_typeck/check/upvar.rs b/src/librustc_typeck/check/upvar.rs
index 86165d50b27..ffd7c2114e5 100644
--- a/src/librustc_typeck/check/upvar.rs
+++ b/src/librustc_typeck/check/upvar.rs
@@ -122,14 +122,18 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
         };
 
         self.tcx.with_freevars(closure_node_id, |freevars| {
+            let mut freevar_list: Vec<ty::UpvarId> = Vec::with_capacity(freevars.len());
             for freevar in freevars {
                 let upvar_id = ty::UpvarId {
                     var_path: ty::UpvarPath {
-                        hir_id : self.tcx.hir().node_to_hir_id(freevar.var_id()),
+                        hir_id: self.tcx.hir().node_to_hir_id(freevar.var_id()),
                     },
                     closure_expr_id: LocalDefId::from_def_id(closure_def_id),
                 };
                 debug!("seed upvar_id {:?}", upvar_id);
+                // Adding the upvar Id to the list of Upvars, which will be added
+                // to the map for the closure at the end of the for loop.
+                freevar_list.push(upvar_id);
 
                 let capture_kind = match capture_clause {
                     hir::CaptureByValue => ty::UpvarCapture::ByValue,
@@ -149,6 +153,15 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
                     .upvar_capture_map
                     .insert(upvar_id, capture_kind);
             }
+            // Add the vector of freevars to the map keyed with the closure id.
+            // This gives us an easier access to them without having to call
+            // with_freevars again..
+            if !freevar_list.is_empty() {
+                self.tables
+                    .borrow_mut()
+                    .upvar_list
+                    .insert(closure_def_id, freevar_list);
+            }
         });
 
         let body_owner_def_id = self.tcx.hir().body_owner_def_id(body.id());
@@ -166,7 +179,8 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
             self.param_env,
             region_scope_tree,
             &self.tables.borrow(),
-        ).consume_body(body);
+        )
+        .consume_body(body);
 
         if let Some(closure_substs) = infer_kind {
             // Unify the (as yet unbound) type variable in the closure
@@ -240,9 +254,7 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
                     let var_hir_id = tcx.hir().node_to_hir_id(var_node_id);
                     let freevar_ty = self.node_ty(var_hir_id);
                     let upvar_id = ty::UpvarId {
-                        var_path: ty::UpvarPath {
-                            hir_id: var_hir_id,
-                        },
+                        var_path: ty::UpvarPath { hir_id: var_hir_id },
                         closure_expr_id: LocalDefId::from_def_id(closure_def_index),
                     };
                     let capture = self.tables.borrow().upvar_capture(upvar_id);
@@ -262,7 +274,8 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
                             },
                         ),
                     }
-                }).collect()
+                })
+                .collect()
         })
     }
 }
diff --git a/src/librustc_typeck/check/writeback.rs b/src/librustc_typeck/check/writeback.rs
index 38de936a027..c61159eb494 100644
--- a/src/librustc_typeck/check/writeback.rs
+++ b/src/librustc_typeck/check/writeback.rs
@@ -21,6 +21,15 @@ use syntax_pos::Span;
 ///////////////////////////////////////////////////////////////////////////
 // Entry point
 
+// During type inference, partially inferred types are
+// represented using Type variables (ty::Infer). These don't appear in
+// the final TypeckTables since all of the types should have been
+// inferred once typeck_tables_of is done.
+// When type inference is running however, having to update the typeck
+// tables every time a new type is inferred would be unreasonably slow,
+// so instead all of the replacement happens at the end in
+// resolve_type_vars_in_body, which creates a new TypeTables which
+// doesn't contain any inference types.
 impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
     pub fn resolve_type_vars_in_body(&self, body: &'gcx hir::Body) -> &'gcx ty::TypeckTables<'gcx> {
         let item_id = self.tcx.hir().body_owner(body.id());
@@ -35,7 +44,8 @@ impl<'a, 'gcx, 'tcx> FnCtxt<'a, 'gcx, 'tcx> {
             wbcx.visit_node_id(arg.pat.span, arg.hir_id);
         }
         wbcx.visit_body(body);
-        wbcx.visit_upvar_borrow_map();
+        wbcx.visit_upvar_capture_map();
+        wbcx.visit_upvar_list_map();
         wbcx.visit_closures();
         wbcx.visit_liberated_fn_sigs();
         wbcx.visit_fru_field_types();
@@ -291,7 +301,7 @@ impl<'cx, 'gcx, 'tcx> Visitor<'gcx> for WritebackCx<'cx, 'gcx, 'tcx> {
 }
 
 impl<'cx, 'gcx, 'tcx> WritebackCx<'cx, 'gcx, 'tcx> {
-    fn visit_upvar_borrow_map(&mut self) {
+    fn visit_upvar_capture_map(&mut self) {
         for (upvar_id, upvar_capture) in self.fcx.tables.borrow().upvar_capture_map.iter() {
             let new_upvar_capture = match *upvar_capture {
                 ty::UpvarCapture::ByValue => ty::UpvarCapture::ByValue,
@@ -314,6 +324,21 @@ impl<'cx, 'gcx, 'tcx> WritebackCx<'cx, 'gcx, 'tcx> {
         }
     }
 
+    /// Runs through the function context's upvar list map and adds the same to
+    /// the TypeckTables. upvarlist is a hashmap of the list of upvars referred
+    /// to in a closure..
+    fn visit_upvar_list_map(&mut self) {
+        for (closure_def_id, upvar_list) in self.fcx.tables.borrow().upvar_list.iter() {
+            debug!(
+                "UpvarIDs captured by closure {:?} are: {:?}",
+                closure_def_id, upvar_list
+            );
+            self.tables
+                .upvar_list
+                .insert(*closure_def_id, upvar_list.to_vec());
+        }
+    }
+
     fn visit_closures(&mut self) {
         let fcx_tables = self.fcx.tables.borrow();
         debug_assert_eq!(fcx_tables.local_id_root, self.tables.local_id_root);