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