Wrap more into into closure_typeinfo query.

cjgillot · cjgillot · commit 0915d55d871e · 2023-02-26T10:30:27.000Z
diff --git a/compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs b/compiler/rustc_borrowck/src/diagnostics/conflict_errors.rs
@@ -6,7 +6,7 @@ use rustc_errors::{
     struct_span_err, Applicability, Diagnostic, DiagnosticBuilder, ErrorGuaranteed, MultiSpan,
 };
 use rustc_hir as hir;
-use rustc_hir::def::Res;
+use rustc_hir::def::{DefKind, Res};
 use rustc_hir::intravisit::{walk_block, walk_expr, Visitor};
 use rustc_hir::{AsyncGeneratorKind, GeneratorKind, LangItem};
 use rustc_infer::infer::TyCtxtInferExt;
@@ -236,10 +236,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
             let ty = used_place.ty(self.body, self.infcx.tcx).ty;
             let needs_note = match ty.kind() {
                 ty::Closure(id, _) => {
-                    let tables = self.infcx.tcx.typeck(id.expect_local());
-                    let hir_id = self.infcx.tcx.hir().local_def_id_to_hir_id(id.expect_local());
-
-                    tables.closure_kind_origins().get(hir_id).is_none()
+                    self.infcx.tcx.closure_kind_origin(id.expect_local()).is_none()
                 }
                 _ => true,
             };
@@ -1670,27 +1667,19 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                 format!("`{}` would have to be valid for `{}`...", name, region_name),
             );
 
-            let fn_hir_id = self.mir_hir_id();
             err.span_label(
                 drop_span,
                 format!(
                     "...but `{}` will be dropped here, when the {} returns",
                     name,
                     self.infcx
                         .tcx
-                        .hir()
-                        .opt_name(fn_hir_id)
+                        .opt_item_name(self.mir_def_id().to_def_id())
                         .map(|name| format!("function `{}`", name))
                         .unwrap_or_else(|| {
-                            match &self
-                                .infcx
-                                .tcx
-                                .typeck(self.mir_def_id())
-                                .node_type(fn_hir_id)
-                                .kind()
-                            {
-                                ty::Closure(..) => "enclosing closure",
-                                ty::Generator(..) => "enclosing generator",
+                            match &self.infcx.tcx.def_kind(self.mir_def_id()) {
+                                DefKind::Closure => "enclosing closure",
+                                DefKind::Generator => "enclosing generator",
                                 kind => bug!("expected closure or generator, found {:?}", kind),
                             }
                             .to_string()
diff --git a/compiler/rustc_borrowck/src/diagnostics/mod.rs b/compiler/rustc_borrowck/src/diagnostics/mod.rs
@@ -115,11 +115,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                     debug!("add_moved_or_invoked_closure_note: closure={:?}", closure);
                     if let ty::Closure(did, _) = self.body.local_decls[closure].ty.kind() {
                         let did = did.expect_local();
-                        let hir_id = self.infcx.tcx.hir().local_def_id_to_hir_id(did);
-
-                        if let Some((span, hir_place)) =
-                            self.infcx.tcx.typeck(did).closure_kind_origins().get(hir_id)
-                        {
+                        if let Some((span, hir_place)) = self.infcx.tcx.closure_kind_origin(did) {
                             diag.span_note(
                                 *span,
                                 &format!(
@@ -139,11 +135,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         if let Some(target) = target {
             if let ty::Closure(did, _) = self.body.local_decls[target].ty.kind() {
                 let did = did.expect_local();
-                let hir_id = self.infcx.tcx.hir().local_def_id_to_hir_id(did);
-
-                if let Some((span, hir_place)) =
-                    self.infcx.tcx.typeck(did).closure_kind_origins().get(hir_id)
-                {
+                if let Some((span, hir_place)) = self.infcx.tcx.closure_kind_origin(did) {
                     diag.span_note(
                         *span,
                         &format!(
@@ -373,14 +365,8 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
                     //
                     // We know the field exists so it's safe to call operator[] and `unwrap` here.
                     let def_id = def_id.expect_local();
-                    let var_id = self
-                        .infcx
-                        .tcx
-                        .typeck(def_id)
-                        .closure_min_captures_flattened(def_id)
-                        .nth(field.index())
-                        .unwrap()
-                        .get_root_variable();
+                    let var_id =
+                        self.infcx.tcx.closure_captures(def_id)[field.index()].get_root_variable();
 
                     Some(self.infcx.tcx.hir().name(var_id).to_string())
                 }
@@ -987,7 +973,7 @@ impl<'cx, 'tcx> MirBorrowckCtxt<'cx, 'tcx> {
         debug!("closure_span: hir_id={:?} expr={:?}", hir_id, expr);
         if let hir::ExprKind::Closure(&hir::Closure { body, fn_decl_span, .. }) = expr {
             for (captured_place, place) in
-                self.infcx.tcx.typeck(def_id).closure_min_captures_flattened(def_id).zip(places)
+                self.infcx.tcx.closure_captures(def_id).iter().zip(places)
             {
                 match place {
                     Operand::Copy(place) | Operand::Move(place)
diff --git a/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs b/compiler/rustc_borrowck/src/diagnostics/mutability_errors.rs
@@ -901,10 +901,7 @@ impl<'a, 'tcx> MirBorrowckCtxt<'a, 'tcx> {
         err: &mut Diagnostic,
     ) {
         let tables = tcx.typeck(closure_local_def_id);
-        let closure_hir_id = tcx.hir().local_def_id_to_hir_id(closure_local_def_id);
-        if let Some((span, closure_kind_origin)) =
-            &tables.closure_kind_origins().get(closure_hir_id)
-        {
+        if let Some((span, closure_kind_origin)) = tcx.closure_kind_origin(closure_local_def_id) {
             let reason = if let PlaceBase::Upvar(upvar_id) = closure_kind_origin.base {
                 let upvar = ty::place_to_string_for_capture(tcx, closure_kind_origin);
                 let root_hir_id = upvar_id.var_path.hir_id;
diff --git a/compiler/rustc_borrowck/src/type_check/input_output.rs b/compiler/rustc_borrowck/src/type_check/input_output.rs
@@ -26,11 +26,7 @@ impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
         if !self.tcx().is_closure(mir_def_id.to_def_id()) {
             return;
         }
-        let Some(user_provided_poly_sig) =
-            self.tcx().typeck(mir_def_id).user_provided_sigs.get(&mir_def_id)
-        else {
-            return;
-        };
+        let user_provided_poly_sig = self.tcx().closure_user_provided_sig(mir_def_id);
 
         // Instantiate the canonicalized variables from user-provided signature
         // (e.g., the `_` in the code above) with fresh variables.
diff --git a/compiler/rustc_middle/src/arena.rs b/compiler/rustc_middle/src/arena.rs
@@ -116,6 +116,7 @@ macro_rules! arena_types {
             [] bit_set_u32: rustc_index::bit_set::BitSet<u32>,
             [] external_constraints: rustc_middle::traits::solve::ExternalConstraintsData<'tcx>,
             [decode] doc_link_resolutions: rustc_hir::def::DocLinkResMap,
+            [] closure_kind_origin: (rustc_span::Span, rustc_middle::hir::place::Place<'tcx>),
         ]);
     )
 }
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -475,7 +475,7 @@ rustc_queries! {
         }
     }
 
-    query closure_captures(key: LocalDefId) -> &'tcx [&'tcx ty::CapturedPlace<'tcx>] {
+    query closure_typeinfo(key: LocalDefId) -> ty::ClosureTypeInfo<'tcx> {
         desc {
             |tcx| "finding symbols for captures of closure `{}`",
             tcx.def_path_str(key.to_def_id())
diff --git a/compiler/rustc_middle/src/ty/closure.rs b/compiler/rustc_middle/src/ty/closure.rs
@@ -6,7 +6,6 @@ use crate::{mir, ty};
 use std::fmt::Write;
 
 use rustc_data_structures::fx::{FxHashMap, FxIndexMap};
-use rustc_hir::def::DefKind;
 use rustc_hir::def_id::{DefId, LocalDefId};
 use rustc_hir::{self as hir, LangItem};
 use rustc_span::symbol::Ident;
@@ -234,14 +233,39 @@ impl<'tcx> CapturedPlace<'tcx> {
     }
 }
 
-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 &[] };
+#[derive(Copy, Clone, Debug, HashStable)]
+pub struct ClosureTypeInfo<'tcx> {
+    user_provided_sig: ty::CanonicalPolyFnSig<'tcx>,
+    captures: &'tcx [&'tcx ty::CapturedPlace<'tcx>],
+    kind_origin: Option<&'tcx (Span, HirPlace<'tcx>)>,
+}
+
+fn closure_typeinfo<'tcx>(tcx: TyCtxt<'tcx>, def: LocalDefId) -> ClosureTypeInfo<'tcx> {
+    debug_assert!(tcx.is_closure(def.to_def_id()));
     let typeck_results = tcx.typeck(def);
+    let user_provided_sig = typeck_results.user_provided_sigs[&def];
     let captures = typeck_results.closure_min_captures_flattened(def);
-    tcx.arena.alloc_from_iter(captures)
+    let captures = tcx.arena.alloc_from_iter(captures);
+    let hir_id = tcx.hir().local_def_id_to_hir_id(def);
+    let kind_origin = typeck_results.closure_kind_origins().get(hir_id);
+    ClosureTypeInfo { user_provided_sig, captures, kind_origin }
+}
+
+impl<'tcx> TyCtxt<'tcx> {
+    pub fn closure_kind_origin(self, def_id: LocalDefId) -> Option<&'tcx (Span, HirPlace<'tcx>)> {
+        self.closure_typeinfo(def_id).kind_origin
+    }
+
+    pub fn closure_user_provided_sig(self, def_id: LocalDefId) -> ty::CanonicalPolyFnSig<'tcx> {
+        self.closure_typeinfo(def_id).user_provided_sig
+    }
+
+    pub fn closure_captures(self, def_id: LocalDefId) -> &'tcx [&'tcx ty::CapturedPlace<'tcx>] {
+        if !self.is_closure(def_id.to_def_id()) {
+            return &[];
+        };
+        self.closure_typeinfo(def_id).captures
+    }
 }
 
 /// Return true if the `proj_possible_ancestor` represents an ancestor path
@@ -434,5 +458,5 @@ impl BorrowKind {
 }
 
 pub fn provide(providers: &mut ty::query::Providers) {
-    *providers = ty::query::Providers { closure_captures, ..*providers }
+    *providers = ty::query::Providers { closure_typeinfo, ..*providers }
 }
diff --git a/compiler/rustc_middle/src/ty/mod.rs b/compiler/rustc_middle/src/ty/mod.rs
@@ -73,7 +73,7 @@ pub use self::binding::BindingMode;
 pub use self::binding::BindingMode::*;
 pub use self::closure::{
     is_ancestor_or_same_capture, place_to_string_for_capture, BorrowKind, CaptureInfo,
-    CapturedPlace, ClosureKind, MinCaptureInformationMap, MinCaptureList,
+    CapturedPlace, ClosureKind, ClosureTypeInfo, MinCaptureInformationMap, MinCaptureList,
     RootVariableMinCaptureList, UpvarCapture, UpvarCaptureMap, UpvarId, UpvarListMap, UpvarPath,
     CAPTURE_STRUCT_LOCAL,
 };
diff --git a/compiler/rustc_middle/src/ty/typeck_results.rs b/compiler/rustc_middle/src/ty/typeck_results.rs
@@ -569,7 +569,7 @@ impl<'a, V> LocalTableInContext<'a, V> {
         self.data.contains_key(&id.local_id)
     }
 
-    pub fn get(&self, id: hir::HirId) -> Option<&V> {
+    pub fn get(&self, id: hir::HirId) -> Option<&'a V> {
         validate_hir_id_for_typeck_results(self.hir_owner, id);
         self.data.get(&id.local_id)
     }

Original file line number	Diff line number	Diff line change
`@@ -116,6 +116,7 @@ macro_rules! arena_types {`
`116`	`116`	`[] bit_set_u32: rustc_index::bit_set::BitSet<u32>,`
`117`	`117`	`[] external_constraints: rustc_middle::traits::solve::ExternalConstraintsData<'tcx>,`
`118`	`118`	`[decode] doc_link_resolutions: rustc_hir::def::DocLinkResMap,`
	`119`	`+ [] closure_kind_origin: (rustc_span::Span, rustc_middle::hir::place::Place<'tcx>),`
`119`	`120`	`]);`
`120`	`121`	`)`
`121`	`122`	`}`
Original file line number	Diff line number	Diff line change
`@@ -475,7 +475,7 @@ rustc_queries! {`
`475`	`475`	`}`
`476`	`476`	`}`
`477`	`477`
`478`		`- query closure_captures(key: LocalDefId) -> &'tcx [&'tcx ty::CapturedPlace<'tcx>] {`
	`478`	`+ query closure_typeinfo(key: LocalDefId) -> ty::ClosureTypeInfo<'tcx> {`
`479`	`479`	`desc {`
`480`	`480`	\|tcx\| "finding symbols for captures of closure `{}`",
`481`	`481`	`tcx.def_path_str(key.to_def_id())`
Original file line number	Diff line number	Diff line change
`@@ -569,7 +569,7 @@ impl<'a, V> LocalTableInContext<'a, V> {`
`569`	`569`	`self.data.contains_key(&id.local_id)`
`570`	`570`	`}`
`571`	`571`
`572`		`- pub fn get(&self, id: hir::HirId) -> Option<&V> {`
	`572`	`+ pub fn get(&self, id: hir::HirId) -> Option<&'a V> {`
`573`	`573`	`validate_hir_id_for_typeck_results(self.hir_owner, id);`
`574`	`574`	`self.data.get(&id.local_id)`
`575`	`575`	`}`