Allow arbitrary keys in feeding API

2 files changed, 19 insertions, 12 deletions

diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
index c5683a9db94..40207c335a4 100644
--- a/compiler/rustc_middle/src/ty/context.rs
+++ b/compiler/rustc_middle/src/ty/context.rs
@@ -1034,16 +1034,23 @@ pub struct FreeRegionInfo {
 
 /// This struct should only be created by `create_def`.
 #[derive(Copy, Clone)]
-pub struct TyCtxtFeed<'tcx> {
+pub struct TyCtxtFeed<'tcx, KEY: Copy> {
     pub tcx: TyCtxt<'tcx>,
     // Do not allow direct access, as downstream code must not mutate this field.
-    def_id: LocalDefId,
+    key: KEY,
 }
 
-impl<'tcx> TyCtxtFeed<'tcx> {
+impl<'tcx, KEY: Copy> TyCtxtFeed<'tcx, KEY> {
+    #[inline(always)]
+    pub fn key(&self) -> KEY {
+        self.key
+    }
+}
+
+impl<'tcx> TyCtxtFeed<'tcx, LocalDefId> {
     #[inline(always)]
     pub fn def_id(&self) -> LocalDefId {
-        self.def_id
+        self.key
     }
 }
 
@@ -1515,7 +1522,7 @@ impl<'tcx> TyCtxtAt<'tcx> {
         self,
         parent: LocalDefId,
         data: hir::definitions::DefPathData,
-    ) -> TyCtxtFeed<'tcx> {
+    ) -> TyCtxtFeed<'tcx, LocalDefId> {
         // This function modifies `self.definitions` using a side-effect.
         // We need to ensure that these side effects are re-run by the incr. comp. engine.
         // Depending on the forever-red node will tell the graph that the calling query
@@ -1536,9 +1543,9 @@ impl<'tcx> TyCtxtAt<'tcx> {
         // This is fine because:
         // - those queries are `eval_always` so we won't miss their result changing;
         // - this write will have happened before these queries are called.
-        let def_id = self.definitions.write().create_def(parent, data);
+        let key = self.definitions.write().create_def(parent, data);
 
-        let feed = TyCtxtFeed { tcx: self.tcx, def_id };
+        let feed = TyCtxtFeed { tcx: self.tcx, key };
         feed.def_span(self.span);
         feed
     }
diff --git a/compiler/rustc_middle/src/ty/query.rs b/compiler/rustc_middle/src/ty/query.rs
index a7fd1754960..e073fca92e7 100644
--- a/compiler/rustc_middle/src/ty/query.rs
+++ b/compiler/rustc_middle/src/ty/query.rs
@@ -330,11 +330,11 @@ macro_rules! define_callbacks {
 
 macro_rules! define_feedable {
     ($($(#[$attr:meta])* [$($modifiers:tt)*] fn $name:ident($($K:tt)*) -> $V:ty,)*) => {
-        impl<'tcx> TyCtxtFeed<'tcx> {
-            $($(#[$attr])*
+        $(impl<'tcx, K: IntoQueryParam<$($K)*> + Copy> TyCtxtFeed<'tcx, K> {
+            $(#[$attr])*
             #[inline(always)]
             pub fn $name(self, value: $V) -> query_stored::$name<'tcx> {
-                let key = self.def_id().into_query_param();
+                let key = self.key().into_query_param();
                 opt_remap_env_constness!([$($modifiers)*][key]);
 
                 let tcx = self.tcx;
@@ -361,8 +361,8 @@ macro_rules! define_feedable {
                     dep_graph::hash_result,
                 );
                 cache.complete(key, value, dep_node_index)
-            })*
-        }
+            }
+        })*
     }
 }