1 //! Upvar (closure capture) collection from cross-body HIR uses of `Res::Local`s.
2
3 use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
4 use rustc_hir as hir;
5 use rustc_hir::def::Res;
6 use rustc_hir::intravisit::{self, NestedVisitorMap, Visitor};
7 use rustc_hir::{self, HirId};
8 use rustc_middle::ty::query::Providers;
9 use rustc_middle::ty::TyCtxt;
10 use rustc_span::Span;
11
provide(providers: &mut Providers)12 pub fn provide(providers: &mut Providers) {
13 providers.upvars_mentioned = |tcx, def_id| {
14 if !tcx.is_closure(def_id) {
15 return None;
16 }
17
18 let hir_id = tcx.hir().local_def_id_to_hir_id(def_id.expect_local());
19 let body = tcx.hir().body(tcx.hir().maybe_body_owned_by(hir_id)?);
20
21 let mut local_collector = LocalCollector::default();
22 local_collector.visit_body(body);
23
24 let mut capture_collector = CaptureCollector {
25 tcx,
26 locals: &local_collector.locals,
27 upvars: FxIndexMap::default(),
28 };
29 capture_collector.visit_body(body);
30
31 if !capture_collector.upvars.is_empty() {
32 Some(tcx.arena.alloc(capture_collector.upvars))
33 } else {
34 None
35 }
36 };
37 }
38
39 #[derive(Default)]
40 struct LocalCollector {
41 // FIXME(eddyb) perhaps use `ItemLocalId` instead?
42 locals: FxHashSet<HirId>,
43 }
44
45 impl Visitor<'tcx> for LocalCollector {
46 type Map = intravisit::ErasedMap<'tcx>;
47
nested_visit_map(&mut self) -> NestedVisitorMap<Self::Map>48 fn nested_visit_map(&mut self) -> NestedVisitorMap<Self::Map> {
49 NestedVisitorMap::None
50 }
51
visit_pat(&mut self, pat: &'tcx hir::Pat<'tcx>)52 fn visit_pat(&mut self, pat: &'tcx hir::Pat<'tcx>) {
53 if let hir::PatKind::Binding(_, hir_id, ..) = pat.kind {
54 self.locals.insert(hir_id);
55 }
56 intravisit::walk_pat(self, pat);
57 }
58 }
59
60 struct CaptureCollector<'a, 'tcx> {
61 tcx: TyCtxt<'tcx>,
62 locals: &'a FxHashSet<HirId>,
63 upvars: FxIndexMap<HirId, hir::Upvar>,
64 }
65
66 impl CaptureCollector<'_, '_> {
visit_local_use(&mut self, var_id: HirId, span: Span)67 fn visit_local_use(&mut self, var_id: HirId, span: Span) {
68 if !self.locals.contains(&var_id) {
69 self.upvars.entry(var_id).or_insert(hir::Upvar { span });
70 }
71 }
72 }
73
74 impl Visitor<'tcx> for CaptureCollector<'a, 'tcx> {
75 type Map = intravisit::ErasedMap<'tcx>;
76
nested_visit_map(&mut self) -> NestedVisitorMap<Self::Map>77 fn nested_visit_map(&mut self) -> NestedVisitorMap<Self::Map> {
78 NestedVisitorMap::None
79 }
80
visit_path(&mut self, path: &'tcx hir::Path<'tcx>, _: hir::HirId)81 fn visit_path(&mut self, path: &'tcx hir::Path<'tcx>, _: hir::HirId) {
82 if let Res::Local(var_id) = path.res {
83 self.visit_local_use(var_id, path.span);
84 }
85
86 intravisit::walk_path(self, path);
87 }
88
visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>)89 fn visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>) {
90 if let hir::ExprKind::Closure(..) = expr.kind {
91 let closure_def_id = self.tcx.hir().local_def_id(expr.hir_id);
92 if let Some(upvars) = self.tcx.upvars_mentioned(closure_def_id) {
93 // Every capture of a closure expression is a local in scope,
94 // that is moved/copied/borrowed into the closure value, and
95 // for this analysis they are like any other access to a local.
96 //
97 // E.g. in `|b| |c| (a, b, c)`, the upvars of the inner closure
98 // are `a` and `b`, and while `a` is not directly used in the
99 // outer closure, it needs to be an upvar there too, so that
100 // the inner closure can take it (from the outer closure's env).
101 for (&var_id, upvar) in upvars {
102 self.visit_local_use(var_id, upvar.span);
103 }
104 }
105 }
106
107 intravisit::walk_expr(self, expr);
108 }
109 }
110