rustc_passes/
upvars.rs

1//! Upvar (closure capture) collection from cross-body HIR uses of `Res::Local`s.
2
3use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
4use rustc_hir as hir;
5use rustc_hir::def::Res;
6use rustc_hir::intravisit::{self, Visitor};
7use rustc_hir::{self, HirId};
8use rustc_middle::query::Providers;
9use rustc_middle::ty::TyCtxt;
10use rustc_span::Span;
11
12pub(crate) fn provide(providers: &mut Providers) {
13    providers.upvars_mentioned = |tcx, def_id| {
14        if !tcx.is_closure_like(def_id) {
15            return None;
16        }
17
18        let local_def_id = def_id.expect_local();
19        let body = tcx.hir().maybe_body_owned_by(local_def_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)]
40struct LocalCollector {
41    // FIXME(eddyb) perhaps use `ItemLocalId` instead?
42    locals: FxHashSet<HirId>,
43}
44
45impl<'tcx> Visitor<'tcx> for LocalCollector {
46    fn visit_pat(&mut self, pat: &'tcx hir::Pat<'tcx>) {
47        if let hir::PatKind::Binding(_, hir_id, ..) = pat.kind {
48            self.locals.insert(hir_id);
49        }
50        intravisit::walk_pat(self, pat);
51    }
52}
53
54struct CaptureCollector<'a, 'tcx> {
55    tcx: TyCtxt<'tcx>,
56    locals: &'a FxHashSet<HirId>,
57    upvars: FxIndexMap<HirId, hir::Upvar>,
58}
59
60impl 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
68impl<'tcx> Visitor<'tcx> for CaptureCollector<'_, 'tcx> {
69    fn visit_path(&mut self, path: &hir::Path<'tcx>, _: HirId) {
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
77    fn visit_expr(&mut self, expr: &'tcx hir::Expr<'tcx>) {
78        if let hir::ExprKind::Closure(closure) = expr.kind {
79            if let Some(upvars) = self.tcx.upvars_mentioned(closure.def_id) {
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}