Wrap more into into closure_typeinfo query.

5 files changed, 36 insertions, 11 deletions

diff --git a/compiler/rustc_middle/src/arena.rs b/compiler/rustc_middle/src/arena.rs
index d4019b5bf17..62e44b6298b 100644
--- 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
index 03d2eea7d4a..2bb8a8b5b73 100644
--- 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
index 65df7204302..dc2bd54b7fe 100644
--- 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
index 17262a0be24..d8411042af0 100644
--- 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
index 9beaac87183..586958247fc 100644
--- 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)
     }