src/librustc_passes/upvars.rs

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