src/librustc/hir/upvars.rs

   1 //! Upvar (closure capture) collection from cross-body HIR uses of `Res::Local`s.
   2
   3 use crate::hir::def::Res;
   4 use crate::hir::intravisit::{self, NestedVisitorMap, Visitor};
   5 use crate::hir::{self, HirId};
   6 use crate::ty::query::Providers;
   7 use crate::ty::TyCtxt;
   8 use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
   9 use syntax_pos::Span;
  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 hir_id = tcx.hir().as_local_hir_id(def_id).unwrap();
  18         let body = tcx.hir().body(tcx.hir().maybe_body_owned_by(hir_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<'tcx>) {
  50         if let hir::PatKind::Binding(_, hir_id, ..) = pat.kind {
  51             self.locals.insert(hir_id);
  52         }
  53         intravisit::walk_pat(self, pat);
  54     }
  55 }
  56
  57 struct CaptureCollector<'a, 'tcx> {
  58     tcx: TyCtxt<'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<'tcx>, _: 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<'tcx>) {
  85         if let hir::ExprKind::Closure(..) = expr.kind {
  86             let closure_def_id = self.tcx.hir().local_def_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 }