]>
Commit | Line | Data |
---|---|---|
dc9dc135 XL |
1 | //! Upvar (closure capture) collection from cross-body HIR uses of `Res::Local`s. |
2 | ||
dfeec247 XL |
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}; | |
ba9703b0 XL |
8 | use rustc_middle::ty::query::Providers; |
9 | use rustc_middle::ty::TyCtxt; | |
dfeec247 | 10 | use rustc_span::Span; |
dc9dc135 | 11 | |
f035d41b | 12 | pub fn provide(providers: &mut Providers) { |
f9f354fc | 13 | providers.upvars_mentioned = |tcx, def_id| { |
dc9dc135 XL |
14 | if !tcx.is_closure(def_id) { |
15 | return None; | |
16 | } | |
17 | ||
3dfed10e | 18 | let hir_id = tcx.hir().local_def_id_to_hir_id(def_id.expect_local()); |
dc9dc135 XL |
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 { | |
ba9703b0 | 46 | type Map = intravisit::ErasedMap<'tcx>; |
dfeec247 | 47 | |
ba9703b0 | 48 | fn nested_visit_map(&mut self) -> NestedVisitorMap<Self::Map> { |
dc9dc135 XL |
49 | NestedVisitorMap::None |
50 | } | |
51 | ||
dfeec247 | 52 | fn visit_pat(&mut self, pat: &'tcx hir::Pat<'tcx>) { |
e74abb32 | 53 | if let hir::PatKind::Binding(_, hir_id, ..) = pat.kind { |
dc9dc135 XL |
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<'_, '_> { | |
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> { | |
ba9703b0 | 75 | type Map = intravisit::ErasedMap<'tcx>; |
dfeec247 | 76 | |
ba9703b0 | 77 | fn nested_visit_map(&mut self) -> NestedVisitorMap<Self::Map> { |
dc9dc135 XL |
78 | NestedVisitorMap::None |
79 | } | |
80 | ||
dfeec247 | 81 | fn visit_path(&mut self, path: &'tcx hir::Path<'tcx>, _: hir::HirId) { |
dc9dc135 XL |
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 | ||
dfeec247 | 89 | fn visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>) { |
e74abb32 | 90 | if let hir::ExprKind::Closure(..) = expr.kind { |
416331ca | 91 | let closure_def_id = self.tcx.hir().local_def_id(expr.hir_id); |
f9f354fc | 92 | if let Some(upvars) = self.tcx.upvars_mentioned(closure_def_id) { |
dc9dc135 XL |
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 | } |