compiler/rustc_passes/src/upvars.rs

   1 //! Upvar (closure capture) collection from cross-body HIR uses of `Res::Local`s.
   2
   3 use rustc_data_structures::fx::{FxHashSet, FxIndexMap};
   4 use rustc_hir as hir;
   5 use rustc_hir::def::Res;
   6 use rustc_hir::intravisit::{self, Visitor};
   7 use rustc_hir::{self, HirId};
   8 use rustc_middle::ty::query::Providers;
   9 use rustc_middle::ty::TyCtxt;
  10 use rustc_span::Span;
  11
  12 pub fn provide(providers: &mut Providers) {
  13     providers.upvars_mentioned = |tcx, def_id| {
  14         if !tcx.is_closure(def_id) {
  15             return None;
  16         }
  17
  18         let local_def_id = def_id.expect_local();
  19         let body = tcx.hir().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)]
  40 struct LocalCollector {
  41     // FIXME(eddyb) perhaps use `ItemLocalId` instead?
  42     locals: FxHashSet<HirId>,
  43 }
  44
  45 impl<'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
  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
  68 impl<'tcx> Visitor<'tcx> for CaptureCollector<'_, 'tcx> {
  69     fn visit_path(&mut self, path: &hir::Path<'tcx>, _: hir::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 }