]>
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; | |
5099ac24 | 6 | use rustc_hir::intravisit::{self, Visitor}; |
dfeec247 | 7 | use rustc_hir::{self, HirId}; |
49aad941 | 8 | use rustc_middle::query::Providers; |
ba9703b0 | 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 | ||
064997fb FG |
18 | let local_def_id = def_id.expect_local(); |
19 | let body = tcx.hir().body(tcx.hir().maybe_body_owned_by(local_def_id)?); | |
dc9dc135 XL |
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 | ||
a2a8927a | 45 | impl<'tcx> Visitor<'tcx> for LocalCollector { |
dfeec247 | 46 | fn visit_pat(&mut self, pat: &'tcx hir::Pat<'tcx>) { |
e74abb32 | 47 | if let hir::PatKind::Binding(_, hir_id, ..) = pat.kind { |
dc9dc135 XL |
48 | self.locals.insert(hir_id); |
49 | } | |
50 | intravisit::walk_pat(self, pat); | |
51 | } | |
52 | } | |
53 | ||
54 | struct CaptureCollector<'a, 'tcx> { | |
55 | tcx: TyCtxt<'tcx>, | |
56 | locals: &'a FxHashSet<HirId>, | |
57 | upvars: FxIndexMap<HirId, hir::Upvar>, | |
58 | } | |
59 | ||
60 | impl CaptureCollector<'_, '_> { | |
61 | fn visit_local_use(&mut self, var_id: HirId, span: Span) { | |
62 | if !self.locals.contains(&var_id) { | |
63 | self.upvars.entry(var_id).or_insert(hir::Upvar { span }); | |
64 | } | |
65 | } | |
66 | } | |
67 | ||
a2a8927a | 68 | impl<'tcx> Visitor<'tcx> for CaptureCollector<'_, 'tcx> { |
487cf647 | 69 | fn visit_path(&mut self, path: &hir::Path<'tcx>, _: hir::HirId) { |
dc9dc135 XL |
70 | if let Res::Local(var_id) = path.res { |
71 | self.visit_local_use(var_id, path.span); | |
72 | } | |
73 | ||
74 | intravisit::walk_path(self, path); | |
75 | } | |
76 | ||
dfeec247 | 77 | fn visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>) { |
487cf647 FG |
78 | if let hir::ExprKind::Closure(closure) = expr.kind { |
79 | if let Some(upvars) = self.tcx.upvars_mentioned(closure.def_id) { | |
dc9dc135 XL |
80 | // Every capture of a closure expression is a local in scope, |
81 | // that is moved/copied/borrowed into the closure value, and | |
82 | // for this analysis they are like any other access to a local. | |
83 | // | |
84 | // E.g. in `|b| |c| (a, b, c)`, the upvars of the inner closure | |
85 | // are `a` and `b`, and while `a` is not directly used in the | |
86 | // outer closure, it needs to be an upvar there too, so that | |
87 | // the inner closure can take it (from the outer closure's env). | |
88 | for (&var_id, upvar) in upvars { | |
89 | self.visit_local_use(var_id, upvar.span); | |
90 | } | |
91 | } | |
92 | } | |
93 | ||
94 | intravisit::walk_expr(self, expr); | |
95 | } | |
96 | } |