rustc: collect upvars from HIR, instead of during name resolution.

Author: eddyb

src/librustc/arena.rs

@@ -109,6 +109,7 @@ macro_rules! arena_types {
             >,
             [few] crate_variances: rustc::ty::CrateVariancesMap<'tcx>,
             [few] inferred_outlives_crate: rustc::ty::CratePredicatesMap<'tcx>,
+            [] upvars: rustc_data_structures::fx::FxIndexMap<rustc::hir::HirId, rustc::hir::Upvar>,
         ], $tcx);
     )
 }

src/librustc/hir/map/mod.rs

@@ -26,7 +26,7 @@ use crate::util::common::time;
 
 use std::io;
 use std::result::Result::Err;
-use crate::ty::TyCtxt;
+use crate::ty::query::Providers;
 
 pub mod blocks;
 mod collector;
@@ -1450,11 +1450,13 @@ fn hir_id_to_string(map: &Map<'_>, id: HirId, include_id: bool) -> String {
     }
 }
 
-pub fn def_kind(tcx: TyCtxt<'_, '_, '_>, def_id: DefId) -> Option<DefKind> {
-    if let Some(node_id) = tcx.hir().as_local_node_id(def_id) {
-        tcx.hir().def_kind(node_id)
-    } else {
-        bug!("Calling local def_kind query provider for upstream DefId: {:?}",
-             def_id)
-    }
+pub fn provide(providers: &mut Providers<'_>) {
+    providers.def_kind = |tcx, def_id| {
+        if let Some(node_id) = tcx.hir().as_local_node_id(def_id) {
+            tcx.hir().def_kind(node_id)
+        } else {
+            bug!("Calling local def_kind query provider for upstream DefId: {:?}",
+                def_id)
+        }
+    };
 }

src/librustc/hir/mod.rs

@@ -30,7 +30,6 @@ use syntax::util::parser::ExprPrecedence;
 use crate::ty::AdtKind;
 use crate::ty::query::Providers;
 
-use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::sync::{par_for_each_in, Send, Sync};
 use rustc_data_structures::thin_vec::ThinVec;
 use rustc_macros::HashStable;
@@ -64,6 +63,7 @@ pub mod lowering;
 pub mod map;
 pub mod pat_util;
 pub mod print;
+pub mod upvars;
 
 /// Uniquely identifies a node in the HIR of the current crate. It is
 /// composed of the `owner`, which is the `DefIndex` of the directly enclosing
@@ -2498,8 +2498,6 @@ pub struct Upvar {
     pub span: Span
 }
 
-pub type UpvarMap = NodeMap<FxIndexMap<ast::NodeId, Upvar>>;
-
 pub type CaptureModeMap = NodeMap<CaptureClause>;
 
  // The TraitCandidate's import_ids is empty if the trait is defined in the same module, and
@@ -2518,10 +2516,10 @@ pub type TraitMap = NodeMap<Vec<TraitCandidate>>;
 // imported.
 pub type GlobMap = NodeMap<FxHashSet<Name>>;
 
-
 pub fn provide(providers: &mut Providers<'_>) {
     check_attr::provide(providers);
-    providers.def_kind = map::def_kind;
+    map::provide(providers);
+    upvars::provide(providers);
 }
 
 #[derive(Clone, RustcEncodable, RustcDecodable, HashStable)]

src/librustc/hir/upvars.rs

@@ -0,0 +1,104 @@
+//! Upvar (closure capture) collection from cross-body HIR uses of `Res::Local`s.
+
+use crate::hir::{self, HirId};
+use crate::hir::def::Res;
+use crate::hir::intravisit::{self, Visitor, NestedVisitorMap};
+use crate::ty::TyCtxt;
+use crate::ty::query::Providers;
+use syntax_pos::Span;
+use rustc_data_structures::fx::{FxIndexMap, FxHashSet};
+
+pub fn provide(providers: &mut Providers<'_>) {
+    providers.upvars = |tcx, def_id| {
+        if !tcx.is_closure(def_id) {
+            return None;
+        }
+
+        let node_id = tcx.hir().as_local_node_id(def_id).unwrap();
+        let body = tcx.hir().body(tcx.hir().maybe_body_owned_by(node_id)?);
+
+        let mut local_collector = LocalCollector::default();
+        local_collector.visit_body(body);
+
+        let mut capture_collector = CaptureCollector {
+            tcx,
+            locals: &local_collector.locals,
+            upvars: FxIndexMap::default(),
+        };
+        capture_collector.visit_body(body);
+
+        if !capture_collector.upvars.is_empty() {
+            Some(tcx.arena.alloc(capture_collector.upvars))
+        } else {
+            None
+        }
+    };
+}
+
+#[derive(Default)]
+struct LocalCollector {
+    // FIXME(eddyb) perhaps use `ItemLocalId` instead?
+    locals: FxHashSet<HirId>,
+}
+
+impl Visitor<'tcx> for LocalCollector {
+    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
+        NestedVisitorMap::None
+    }
+
+    fn visit_pat(&mut self, pat: &'tcx hir::Pat) {
+        if let hir::PatKind::Binding(_, hir_id, ..) = pat.node {
+            self.locals.insert(hir_id);
+        }
+        intravisit::walk_pat(self, pat);
+    }
+}
+
+struct CaptureCollector<'a, 'tcx> {
+    tcx: TyCtxt<'a, 'tcx, 'tcx>,
+    locals: &'a FxHashSet<HirId>,
+    upvars: FxIndexMap<HirId, hir::Upvar>,
+}
+
+impl CaptureCollector<'_, '_> {
+    fn visit_local_use(&mut self, var_id: HirId, span: Span) {
+        if !self.locals.contains(&var_id) {
+            self.upvars.entry(var_id).or_insert(hir::Upvar { span });
+        }
+    }
+}
+
+impl Visitor<'tcx> for CaptureCollector<'a, 'tcx> {
+    fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
+        NestedVisitorMap::None
+    }
+
+    fn visit_path(&mut self, path: &'tcx hir::Path, _: hir::HirId) {
+        if let Res::Local(var_id) = path.res {
+            self.visit_local_use(var_id, path.span);
+        }
+
+        intravisit::walk_path(self, path);
+    }
+
+    fn visit_expr(&mut self, expr: &'tcx hir::Expr) {
+        if let hir::ExprKind::Closure(..) = expr.node {
+            let closure_def_id = self.tcx.hir().local_def_id_from_hir_id(expr.hir_id);
+            if let Some(upvars) = self.tcx.upvars(closure_def_id) {
+                // Every capture of a closure expression is a local in scope,
+                // that is moved/copied/borrowed into the closure value, and
+                // for this analysis they are like any other access to a local.
+                //
+                // E.g. in `|b| |c| (a, b, c)`, the upvars of the inner closure
+                // are `a` and `b`, and while `a` is not directly used in the
+                // outer closure, it needs to be an upvar there too, so that
+                // the inner closure can take it (from the outer closure's env).
+                for (&var_id, upvar) in upvars {
+                    self.visit_local_use(var_id, upvar.span);
+                }
+            }
+        }
+
+        intravisit::walk_expr(self, expr);
+    }
+}

src/librustc/ty/context.rs

@@ -54,7 +54,6 @@ use rustc_data_structures::stable_hasher::{HashStable, hash_stable_hashmap,
                                            StableHasher, StableHasherResult,
                                            StableVec};
 use arena::SyncDroplessArena;
-use rustc_data_structures::fx::FxIndexMap;
 use rustc_data_structures::indexed_vec::{Idx, IndexVec};
 use rustc_data_structures::sync::{Lrc, Lock, WorkerLocal};
 use std::any::Any;
@@ -1063,11 +1062,6 @@ pub struct GlobalCtxt<'tcx> {
 
     pub queries: query::Queries<'tcx>,
 
-    // Records the captured variables referenced by every closure
-    // expression. Do not track deps for this, just recompute it from
-    // scratch every time.
-    upvars: FxHashMap<DefId, FxIndexMap<hir::HirId, hir::Upvar>>,
-
     maybe_unused_trait_imports: FxHashSet<DefId>,
     maybe_unused_extern_crates: Vec<(DefId, Span)>,
     /// A map of glob use to a set of names it actually imports. Currently only
@@ -1298,12 +1292,6 @@ impl<'a, 'gcx, 'tcx> TyCtxt<'a, 'gcx, 'tcx> {
                 }).collect();
                 (k, exports)
             }).collect(),
-            upvars: resolutions.upvars.into_iter().map(|(k, upvars)| {
-                let upvars: FxIndexMap<_, _> = upvars.into_iter().map(|(var_id, upvar)| {
-                    (hir.node_to_hir_id(var_id), upvar)
-                }).collect();
-                (hir.local_def_id(k), upvars)
-            }).collect(),
             maybe_unused_trait_imports:
                 resolutions.maybe_unused_trait_imports
                     .into_iter()
@@ -3024,7 +3012,6 @@ pub fn provide(providers: &mut ty::query::Providers<'_>) {
         assert_eq!(id, LOCAL_CRATE);
         tcx.arena.alloc(middle::lang_items::collect(tcx))
     };
-    providers.upvars = |tcx, id| tcx.gcx.upvars.get(&id);
     providers.maybe_unused_trait_import = |tcx, id| {
         tcx.maybe_unused_trait_imports.contains(&id)
     };

src/librustc/ty/mod.rs

@@ -8,7 +8,7 @@ pub use self::BorrowKind::*;
 pub use self::IntVarValue::*;
 pub use self::fold::TypeFoldable;
 
-use crate::hir::{map as hir_map, UpvarMap, GlobMap, TraitMap};
+use crate::hir::{map as hir_map, GlobMap, TraitMap};
 use crate::hir::Node;
 use crate::hir::def::{Res, DefKind, CtorOf, CtorKind, ExportMap};
 use crate::hir::def_id::{CrateNum, DefId, LocalDefId, CRATE_DEF_INDEX, LOCAL_CRATE};
@@ -123,7 +123,6 @@ mod sty;
 
 #[derive(Clone)]
 pub struct Resolutions {
-    pub upvars: UpvarMap,
     pub trait_map: TraitMap,
     pub maybe_unused_trait_imports: NodeSet,
     pub maybe_unused_extern_crates: Vec<(NodeId, Span)>,

src/librustc_interface/passes.rs

@@ -178,7 +178,6 @@ impl ExpansionResult {
         ExpansionResult {
             defs: Steal::new(resolver.definitions),
             resolutions: Steal::new(Resolutions {
-                upvars: resolver.upvars,
                 export_map: resolver.export_map,
                 trait_map: resolver.trait_map,
                 glob_map: resolver.glob_map,
@@ -197,7 +196,6 @@ impl ExpansionResult {
         ExpansionResult {
             defs: Steal::new(resolver.definitions.clone()),
             resolutions: Steal::new(Resolutions {
-                upvars: resolver.upvars.clone(),
                 export_map: resolver.export_map.clone(),
                 trait_map: resolver.trait_map.clone(),
                 glob_map: resolver.glob_map.clone(),