Access upvars through a query.

Author: cjgillot

compiler/rustc_borrowck/src/lib.rs

@@ -202,14 +202,14 @@ fn do_mir_borrowck<'tcx>(
     let mut errors = error::BorrowckErrors::new(infcx.tcx);
 
     // Gather the upvars of a closure, if any.
-    let tables = tcx.typeck_opt_const_arg(def);
-    if let Some(e) = tables.tainted_by_errors {
+    if let Some(e) = input_body.tainted_by_errors {
         infcx.set_tainted_by_errors(e);
         errors.set_tainted_by_errors(e);
     }
-    let upvars: Vec<_> = tables
-        .closure_min_captures_flattened(def.did)
-        .map(|captured_place| {
+    let upvars: Vec<_> = tcx
+        .closure_captures(def.did)
+        .iter()
+        .map(|&captured_place| {
             let capture = captured_place.info.capture_kind;
             let by_ref = match capture {
                 ty::UpvarCapture::ByValue => false,

compiler/rustc_const_eval/src/interpret/validity.rs

@@ -240,10 +240,8 @@ impl<'rt, 'mir, 'tcx: 'mir, M: Machine<'mir, 'tcx>> ValidityVisitor<'rt, 'mir, '
                 // FIXME this should be more descriptive i.e. CapturePlace instead of CapturedVar
                 // https://github.com/rust-lang/project-rfc-2229/issues/46
                 if let Some(local_def_id) = def_id.as_local() {
-                    let tables = self.ecx.tcx.typeck(local_def_id);
-                    if let Some(captured_place) =
-                        tables.closure_min_captures_flattened(local_def_id).nth(field)
-                    {
+                    let captures = self.ecx.tcx.closure_captures(local_def_id);
+                    if let Some(captured_place) = captures.get(field) {
                         // Sometimes the index is beyond the number of upvars (seen
                         // for a generator).
                         let var_hir_id = captured_place.get_root_variable();

compiler/rustc_hir_typeck/src/upvar.rs

@@ -526,10 +526,12 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
                 PlaceBase::Upvar(upvar_id) => upvar_id.var_path.hir_id,
                 base => bug!("Expected upvar, found={:?}", base),
             };
+            let var_ident = self.tcx.hir().ident(var_hir_id);
 
             let Some(min_cap_list) = root_var_min_capture_list.get_mut(&var_hir_id) else {
                 let mutability = self.determine_capture_mutability(&typeck_results, &place);
                 let min_cap_list = vec![ty::CapturedPlace {
+                    var_ident,
                     place,
                     info: capture_info,
                     mutability,
@@ -628,6 +630,7 @@ impl<'a, 'tcx> FnCtxt<'a, 'tcx> {
             if !ancestor_found {
                 let mutability = self.determine_capture_mutability(&typeck_results, &place);
                 let captured_place = ty::CapturedPlace {
+                    var_ident,
                     place,
                     info: updated_capture_info,
                     mutability,

compiler/rustc_middle/src/hir/map/mod.rs

@@ -902,6 +902,11 @@ impl<'hir> Map<'hir> {
         self.opt_ident(id).map(|ident| ident.span)
     }
 
+    #[inline]
+    pub fn ident(self, id: HirId) -> Ident {
+        self.opt_ident(id).unwrap()
+    }
+
     #[inline]
     pub fn opt_name(self, id: HirId) -> Option<Symbol> {
         self.opt_ident(id).map(|ident| ident.name)

compiler/rustc_middle/src/query/mod.rs

@@ -475,14 +475,10 @@ rustc_queries! {
         }
     }
 
-    query symbols_for_closure_captures(
-        key: (LocalDefId, LocalDefId)
-    ) -> &'tcx Vec<rustc_span::Symbol> {
-        arena_cache
+    query closure_captures(key: LocalDefId) -> &'tcx [&'tcx ty::CapturedPlace<'tcx>] {
         desc {
-            |tcx| "finding symbols for captures of closure `{}` in `{}`",
-            tcx.def_path_str(key.1.to_def_id()),
-            tcx.def_path_str(key.0.to_def_id())
+            |tcx| "finding symbols for captures of closure `{}`",
+            tcx.def_path_str(key.to_def_id())
         }
     }
 

compiler/rustc_middle/src/ty/closure.rs

@@ -5,10 +5,11 @@ use crate::{mir, ty};
 
 use std::fmt::Write;
 
-use hir::LangItem;
 use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
-use rustc_hir as hir;
+use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId};
+use rustc_hir::{self as hir, LangItem};
+use rustc_span::symbol::Ident;
 use rustc_span::{Span, Symbol};
 
 use super::{Ty, TyCtxt};
@@ -129,6 +130,9 @@ impl<'tcx> ClosureKind {
 #[derive(PartialEq, Clone, Debug, TyEncodable, TyDecodable, HashStable)]
 #[derive(TypeFoldable, TypeVisitable)]
 pub struct CapturedPlace<'tcx> {
+    /// Name and span where the binding happens.
+    pub var_ident: Ident,
+
     /// The `Place` that is captured.
     pub place: HirPlace<'tcx>,
 
@@ -148,12 +152,8 @@ impl<'tcx> CapturedPlace<'tcx> {
     }
 
     /// Returns a symbol of the captured upvar, which looks like `name__field1__field2`.
-    fn to_symbol(&self, tcx: TyCtxt<'tcx>) -> Symbol {
-        let hir_id = match self.place.base {
-            HirPlaceBase::Upvar(upvar_id) => upvar_id.var_path.hir_id,
-            base => bug!("Expected an upvar, found {:?}", base),
-        };
-        let mut symbol = tcx.hir().name(hir_id).as_str().to_string();
+    pub fn to_symbol(&self) -> Symbol {
+        let mut symbol = self.var_ident.to_string();
 
         let mut ty = self.place.base_ty;
         for proj in self.place.projections.iter() {
@@ -169,11 +169,7 @@ impl<'tcx> CapturedPlace<'tcx> {
                         .unwrap();
                     }
                     ty => {
-                        span_bug!(
-                            self.get_capture_kind_span(tcx),
-                            "Unexpected type {:?} for `Field` projection",
-                            ty
-                        )
+                        bug!("Unexpected type {:?} for `Field` projection", ty)
                     }
                 },
 
@@ -238,10 +234,14 @@ impl<'tcx> CapturedPlace<'tcx> {
     }
 }
 
-fn symbols_for_closure_captures(tcx: TyCtxt<'_>, def_id: (LocalDefId, LocalDefId)) -> Vec<Symbol> {
-    let typeck_results = tcx.typeck(def_id.0);
-    let captures = typeck_results.closure_min_captures_flattened(def_id.1);
-    captures.into_iter().map(|captured_place| captured_place.to_symbol(tcx)).collect()
+fn closure_captures<'tcx>(
+    tcx: TyCtxt<'tcx>,
+    def: LocalDefId,
+) -> &'tcx [&'tcx ty::CapturedPlace<'tcx>] {
+    let (DefKind::Closure | DefKind::Generator) = tcx.def_kind(def) else { return &[] };
+    let typeck_results = tcx.typeck(def);
+    let captures = typeck_results.closure_min_captures_flattened(def);
+    tcx.arena.alloc_from_iter(captures)
 }
 
 /// Return true if the `proj_possible_ancestor` represents an ancestor path
@@ -434,5 +434,5 @@ impl BorrowKind {
 }
 
 pub fn provide(providers: &mut ty::query::Providers) {
-    *providers = ty::query::Providers { symbols_for_closure_captures, ..*providers }
+    *providers = ty::query::Providers { closure_captures, ..*providers }
 }

compiler/rustc_middle/src/ty/structural_impls.rs

@@ -263,6 +263,7 @@ TrivialTypeTraversalAndLiftImpls! {
     crate::ty::UniverseIndex,
     crate::ty::Variance,
     ::rustc_span::Span,
+    ::rustc_span::symbol::Ident,
     ::rustc_errors::ErrorGuaranteed,
     Field,
     interpret::Scalar,

compiler/rustc_mir_build/src/build/mod.rs

@@ -20,7 +20,7 @@ use rustc_middle::mir::*;
 use rustc_middle::thir::{
     self, BindingMode, Expr, ExprId, LintLevel, LocalVarId, Param, ParamId, PatKind, Thir,
 };
-use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitableExt, TypeckResults};
+use rustc_middle::ty::{self, Ty, TyCtxt, TypeVisitableExt};
 use rustc_span::symbol::sym;
 use rustc_span::Span;
 use rustc_span::Symbol;
@@ -155,13 +155,13 @@ struct BlockContext(Vec<BlockFrame>);
 struct Builder<'a, 'tcx> {
     tcx: TyCtxt<'tcx>,
     infcx: InferCtxt<'tcx>,
-    typeck_results: &'tcx TypeckResults<'tcx>,
     region_scope_tree: &'tcx region::ScopeTree,
     param_env: ty::ParamEnv<'tcx>,
 
     thir: &'a Thir<'tcx>,
     cfg: CFG<'tcx>,
 
+    def: ty::WithOptConstParam<LocalDefId>,
     def_id: DefId,
     hir_id: hir::HirId,
     parent_module: DefId,
@@ -522,13 +522,7 @@ fn construct_fn<'tcx>(
         let return_block =
             unpack!(builder.in_breakable_scope(None, Place::return_place(), fn_end, |builder| {
                 Some(builder.in_scope(arg_scope_s, LintLevel::Inherited, |builder| {
-                    builder.args_and_body(
-                        START_BLOCK,
-                        fn_def.did,
-                        arguments,
-                        arg_scope,
-                        &thir[expr],
-                    )
+                    builder.args_and_body(START_BLOCK, arguments, arg_scope, &thir[expr])
                 }))
             }));
         let source_info = builder.source_info(fn_end);
@@ -704,9 +698,9 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             thir,
             tcx,
             infcx,
-            typeck_results: tcx.typeck_opt_const_arg(def),
             region_scope_tree: tcx.region_scope_tree(def.did),
             param_env,
+            def,
             def_id: def.did.to_def_id(),
             hir_id,
             parent_module: tcx.parent_module(hir_id).to_def_id(),
@@ -756,14 +750,78 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             self.var_debug_info,
             self.fn_span,
             self.generator_kind,
-            self.typeck_results.tainted_by_errors,
+            None,
         )
     }
 
+    fn insert_upvar_arg(&mut self) {
+        let Some(closure_arg) = self.local_decls.get(ty::CAPTURE_STRUCT_LOCAL) else { return };
+
+        let mut closure_ty = closure_arg.ty;
+        let mut closure_env_projs = vec![];
+        if let ty::Ref(_, ty, _) = closure_ty.kind() {
+            closure_env_projs.push(ProjectionElem::Deref);
+            closure_ty = *ty;
+        }
+
+        let upvar_substs = match closure_ty.kind() {
+            ty::Closure(_, substs) => ty::UpvarSubsts::Closure(substs),
+            ty::Generator(_, substs, _) => ty::UpvarSubsts::Generator(substs),
+            _ => return,
+        };
+
+        // In analyze_closure() in upvar.rs we gathered a list of upvars used by an
+        // indexed closure and we stored in a map called closure_min_captures in TypeckResults
+        // with the closure's DefId. Here, we run through that vec of UpvarIds for
+        // the given closure and use the necessary information to create upvar
+        // debuginfo and to fill `self.upvars`.
+        let capture_tys = upvar_substs.upvar_tys();
+
+        let tcx = self.tcx;
+        self.upvars = tcx
+            .closure_captures(self.def.did)
+            .iter()
+            .zip(capture_tys)
+            .enumerate()
+            .map(|(i, (captured_place, ty))| {
+                let name = captured_place.to_symbol();
+
+                let capture = captured_place.info.capture_kind;
+                let var_id = match captured_place.place.base {
+                    HirPlaceBase::Upvar(upvar_id) => upvar_id.var_path.hir_id,
+                    _ => bug!("Expected an upvar"),
+                };
+
+                let mutability = captured_place.mutability;
+
+                let mut projs = closure_env_projs.clone();
+                projs.push(ProjectionElem::Field(Field::new(i), ty));
+                match capture {
+                    ty::UpvarCapture::ByValue => {}
+                    ty::UpvarCapture::ByRef(..) => {
+                        projs.push(ProjectionElem::Deref);
+                    }
+                };
+
+                let use_place = Place {
+                    local: ty::CAPTURE_STRUCT_LOCAL,
+                    projection: tcx.mk_place_elems(&projs),
+                };
+                self.var_debug_info.push(VarDebugInfo {
+                    name,
+                    source_info: SourceInfo::outermost(captured_place.var_ident.span),
+                    value: VarDebugInfoContents::Place(use_place),
+                });
+
+                let capture = Capture { captured_place, use_place, mutability };
+                (var_id, capture)
+            })
+            .collect();
+    }
+
     fn args_and_body(
         &mut self,
         mut block: BasicBlock,
-        fn_def_id: LocalDefId,
         arguments: &IndexVec<ParamId, Param<'tcx>>,
         argument_scope: region::Scope,
         expr: &Expr<'tcx>,
@@ -785,69 +843,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
             }
         }
 
-        let tcx = self.tcx;
-        let tcx_hir = tcx.hir();
-        let hir_typeck_results = self.typeck_results;
-
-        // In analyze_closure() in upvar.rs we gathered a list of upvars used by an
-        // indexed closure and we stored in a map called closure_min_captures in TypeckResults
-        // with the closure's DefId. Here, we run through that vec of UpvarIds for
-        // the given closure and use the necessary information to create upvar
-        // debuginfo and to fill `self.upvars`.
-        if hir_typeck_results.closure_min_captures.get(&fn_def_id).is_some() {
-            let mut closure_env_projs = vec![];
-            let mut closure_ty = self.local_decls[ty::CAPTURE_STRUCT_LOCAL].ty;
-            if let ty::Ref(_, ty, _) = closure_ty.kind() {
-                closure_env_projs.push(ProjectionElem::Deref);
-                closure_ty = *ty;
-            }
-            let upvar_substs = match closure_ty.kind() {
-                ty::Closure(_, substs) => ty::UpvarSubsts::Closure(substs),
-                ty::Generator(_, substs, _) => ty::UpvarSubsts::Generator(substs),
-                _ => span_bug!(self.fn_span, "upvars with non-closure env ty {:?}", closure_ty),
-            };
-            let def_id = self.def_id.as_local().unwrap();
-            let capture_syms = tcx.symbols_for_closure_captures((def_id, fn_def_id));
-            let capture_tys = upvar_substs.upvar_tys();
-            let captures_with_tys = hir_typeck_results
-                .closure_min_captures_flattened(fn_def_id)
-                .zip(capture_tys.zip(capture_syms));
-
-            self.upvars = captures_with_tys
-                .enumerate()
-                .map(|(i, (captured_place, (ty, sym)))| {
-                    let capture = captured_place.info.capture_kind;
-                    let var_id = match captured_place.place.base {
-                        HirPlaceBase::Upvar(upvar_id) => upvar_id.var_path.hir_id,
-                        _ => bug!("Expected an upvar"),
-                    };
-
-                    let mutability = captured_place.mutability;
-
-                    let mut projs = closure_env_projs.clone();
-                    projs.push(ProjectionElem::Field(Field::new(i), ty));
-                    match capture {
-                        ty::UpvarCapture::ByValue => {}
-                        ty::UpvarCapture::ByRef(..) => {
-                            projs.push(ProjectionElem::Deref);
-                        }
-                    };
-
-                    let use_place = Place {
-                        local: ty::CAPTURE_STRUCT_LOCAL,
-                        projection: tcx.mk_place_elems(&projs),
-                    };
-                    self.var_debug_info.push(VarDebugInfo {
-                        name: *sym,
-                        source_info: SourceInfo::outermost(tcx_hir.span(var_id)),
-                        value: VarDebugInfoContents::Place(use_place),
-                    });
-
-                    let capture = Capture { captured_place, use_place, mutability };
-                    (var_id, capture)
-                })
-                .collect();
-        }
+        self.insert_upvar_arg();
 
         let mut scope = None;
         // Bind the argument patterns

compiler/rustc_mir_build/src/thir/cx/expr.rs

@@ -541,8 +541,9 @@ impl<'tcx> Cx<'tcx> {
                 let def_id = def_id.expect_local();
 
                 let upvars = self
-                    .typeck_results
-                    .closure_min_captures_flattened(def_id)
+                    .tcx
+                    .closure_captures(def_id)
+                    .iter()
                     .zip(substs.upvar_tys())
                     .map(|(captured_place, ty)| {
                         let upvars = self.capture_upvar(expr, captured_place, ty);