1 //! Upvar (closure capture) collection from cross-body HIR uses of `Res::Local`s.
3 use crate::hir::map::Map;
4 use crate::ty::query::Providers;
6 use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
8 use rustc_hir::def::Res;
9 use rustc_hir::intravisit::{self, NestedVisitorMap, Visitor};
10 use rustc_hir::{self, HirId};
13 pub fn provide(providers: &mut Providers<'_>) {
14 providers.upvars = |tcx, def_id| {
15 if !tcx.is_closure(def_id) {
19 let hir_id = tcx.hir().as_local_hir_id(def_id).unwrap();
20 let body = tcx.hir().body(tcx.hir().maybe_body_owned_by(hir_id)?);
22 let mut local_collector = LocalCollector::default();
23 local_collector.visit_body(body);
25 let mut capture_collector = CaptureCollector {
27 locals: &local_collector.locals,
28 upvars: FxIndexMap::default(),
30 capture_collector.visit_body(body);
32 if !capture_collector.upvars.is_empty() {
33 Some(tcx.arena.alloc(capture_collector.upvars))
41 struct LocalCollector {
42 // FIXME(eddyb) perhaps use `ItemLocalId` instead?
43 locals: FxHashSet<HirId>,
46 impl Visitor<'tcx> for LocalCollector {
49 fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, Self::Map> {
50 NestedVisitorMap::None
53 fn visit_pat(&mut self, pat: &'tcx hir::Pat<'tcx>) {
54 if let hir::PatKind::Binding(_, hir_id, ..) = pat.kind {
55 self.locals.insert(hir_id);
57 intravisit::walk_pat(self, pat);
61 struct CaptureCollector<'a, 'tcx> {
63 locals: &'a FxHashSet<HirId>,
64 upvars: FxIndexMap<HirId, hir::Upvar>,
67 impl CaptureCollector<'_, '_> {
68 fn visit_local_use(&mut self, var_id: HirId, span: Span) {
69 if !self.locals.contains(&var_id) {
70 self.upvars.entry(var_id).or_insert(hir::Upvar { span });
75 impl Visitor<'tcx> for CaptureCollector<'a, 'tcx> {
78 fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, Self::Map> {
79 NestedVisitorMap::None
82 fn visit_path(&mut self, path: &'tcx hir::Path<'tcx>, _: hir::HirId) {
83 if let Res::Local(var_id) = path.res {
84 self.visit_local_use(var_id, path.span);
87 intravisit::walk_path(self, path);
90 fn visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>) {
91 if let hir::ExprKind::Closure(..) = expr.kind {
92 let closure_def_id = self.tcx.hir().local_def_id(expr.hir_id);
93 if let Some(upvars) = self.tcx.upvars(closure_def_id) {
94 // Every capture of a closure expression is a local in scope,
95 // that is moved/copied/borrowed into the closure value, and
96 // for this analysis they are like any other access to a local.
98 // E.g. in `|b| |c| (a, b, c)`, the upvars of the inner closure
99 // are `a` and `b`, and while `a` is not directly used in the
100 // outer closure, it needs to be an upvar there too, so that
101 // the inner closure can take it (from the outer closure's env).
102 for (&var_id, upvar) in upvars {
103 self.visit_local_use(var_id, upvar.span);
108 intravisit::walk_expr(self, expr);