]> git.lizzy.rs Git - rust.git/blob - src/librustc/hir/upvars.rs
Rollup merge of #67671 - estebank:type-impl-trait, r=davidtwco
[rust.git] / src / librustc / hir / upvars.rs
1 //! Upvar (closure capture) collection from cross-body HIR uses of `Res::Local`s.
2
3 use crate::hir::intravisit::{self, NestedVisitorMap, Visitor};
4 use crate::ty::query::Providers;
5 use crate::ty::TyCtxt;
6 use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
7 use rustc_hir as hir;
8 use rustc_hir::def::Res;
9 use rustc_hir::{self, HirId};
10 use rustc_span::Span;
11
12 pub fn provide(providers: &mut Providers<'_>) {
13     providers.upvars = |tcx, def_id| {
14         if !tcx.is_closure(def_id) {
15             return None;
16         }
17
18         let hir_id = tcx.hir().as_local_hir_id(def_id).unwrap();
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     fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
47         NestedVisitorMap::None
48     }
49
50     fn visit_pat(&mut self, pat: &'tcx hir::Pat<'tcx>) {
51         if let hir::PatKind::Binding(_, hir_id, ..) = pat.kind {
52             self.locals.insert(hir_id);
53         }
54         intravisit::walk_pat(self, pat);
55     }
56 }
57
58 struct CaptureCollector<'a, 'tcx> {
59     tcx: TyCtxt<'tcx>,
60     locals: &'a FxHashSet<HirId>,
61     upvars: FxIndexMap<HirId, hir::Upvar>,
62 }
63
64 impl CaptureCollector<'_, '_> {
65     fn visit_local_use(&mut self, var_id: HirId, span: Span) {
66         if !self.locals.contains(&var_id) {
67             self.upvars.entry(var_id).or_insert(hir::Upvar { span });
68         }
69     }
70 }
71
72 impl Visitor<'tcx> for CaptureCollector<'a, 'tcx> {
73     fn nested_visit_map<'this>(&'this mut self) -> NestedVisitorMap<'this, 'tcx> {
74         NestedVisitorMap::None
75     }
76
77     fn visit_path(&mut self, path: &'tcx hir::Path<'tcx>, _: hir::HirId) {
78         if let Res::Local(var_id) = path.res {
79             self.visit_local_use(var_id, path.span);
80         }
81
82         intravisit::walk_path(self, path);
83     }
84
85     fn visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>) {
86         if let hir::ExprKind::Closure(..) = expr.kind {
87             let closure_def_id = self.tcx.hir().local_def_id(expr.hir_id);
88             if let Some(upvars) = self.tcx.upvars(closure_def_id) {
89                 // Every capture of a closure expression is a local in scope,
90                 // that is moved/copied/borrowed into the closure value, and
91                 // for this analysis they are like any other access to a local.
92                 //
93                 // E.g. in `|b| |c| (a, b, c)`, the upvars of the inner closure
94                 // are `a` and `b`, and while `a` is not directly used in the
95                 // outer closure, it needs to be an upvar there too, so that
96                 // the inner closure can take it (from the outer closure's env).
97                 for (&var_id, upvar) in upvars {
98                     self.visit_local_use(var_id, upvar.span);
99                 }
100             }
101         }
102
103         intravisit::walk_expr(self, expr);
104     }
105 }