Make query helpers on `TyCtxt` take `impl Into<DefId>`

1 files changed, 66 insertions, 23 deletions

diff --git a/src/librustc_middle/ty/query/plumbing.rs b/src/librustc_middle/ty/query/plumbing.rs
index 1bb392f436f..612507711fd 100644
--- a/src/librustc_middle/ty/query/plumbing.rs
+++ b/src/librustc_middle/ty/query/plumbing.rs
@@ -234,18 +234,73 @@ macro_rules! hash_result {
 
 macro_rules! define_queries {
     (<$tcx:tt> $($category:tt {
-        $($(#[$attr:meta])* [$($modifiers:tt)*] fn $name:ident: $node:ident($K:ty) -> $V:ty,)*
+        $($(#[$attr:meta])* [$($modifiers:tt)*] fn $name:ident: $node:ident($($K:tt)*) -> $V:ty,)*
     },)*) => {
         define_queries_inner! { <$tcx>
-            $($( $(#[$attr])* category<$category> [$($modifiers)*] fn $name: $node($K) -> $V,)*)*
+            $($( $(#[$attr])* category<$category> [$($modifiers)*] fn $name: $node($($K)*) -> $V,)*)*
         }
     }
 }
 
+macro_rules! define_query_helper {
+    (TyCtxtAt<$tcx:tt>, $(#[$attr:meta])* $name:ident(DefId) -> $V:ty) => {
+        $(#[$attr])*
+        #[inline(always)]
+        pub fn $name(self, key: impl Into<DefId>) -> $V {
+            fn mono(this: TyCtxtAt<$tcx>, key: DefId) -> $V {
+                get_query::<queries::$name<'_>, _>(this.tcx, this.span, key)
+            }
+
+            mono(self, key.into())
+        }
+    };
+    (TyCtxtAt<$tcx:tt>, $(#[$attr:meta])* $name:ident($K:ty) -> $V:ty) => {
+        $(#[$attr])*
+        #[inline(always)]
+        pub fn $name(self, key: $K) -> $V {
+            get_query::<queries::$name<'_>, _>(self.tcx, self.span, key)
+        }
+    };
+
+    (TyCtxt<$tcx:tt>, $(#[$attr:meta])* $name:ident(DefId) -> $V:ty) => {
+        $(#[$attr])*
+        #[inline(always)]
+        pub fn $name(self, key: impl Into<DefId>) -> $V {
+            self.at(DUMMY_SP).$name(key)
+        }
+    };
+    (TyCtxt<$tcx:tt>, $(#[$attr:meta])* $name:ident($K:ty) -> $V:ty) => {
+        $(#[$attr])*
+        #[inline(always)]
+        pub fn $name(self, key: $K) -> $V {
+            self.at(DUMMY_SP).$name(key)
+        }
+    };
+
+    (TyCtxtEnsure<$tcx:tt>, $(#[$attr:meta])* $name:ident(DefId) -> $V:ty) => {
+        $(#[$attr])*
+        #[inline(always)]
+        pub fn $name(self, key: impl Into<DefId>) {
+            fn mono(this: TyCtxtEnsure<$tcx>, key: DefId) {
+                ensure_query::<queries::$name<'_>, _>(this.tcx, key)
+            }
+
+            mono(self, key.into())
+        }
+    };
+    (TyCtxtEnsure<$tcx:tt>, $(#[$attr:meta])* $name:ident($K:ty) -> $V:ty) => {
+        $(#[$attr])*
+        #[inline(always)]
+        pub fn $name(self, key: $K) {
+            ensure_query::<queries::$name<'_>, _>(self.tcx, key)
+        }
+    };
+}
+
 macro_rules! define_queries_inner {
     (<$tcx:tt>
      $($(#[$attr:meta])* category<$category:tt>
-        [$($modifiers:tt)*] fn $name:ident: $node:ident($K:ty) -> $V:ty,)*) => {
+        [$($modifiers:tt)*] fn $name:ident: $node:ident($($K:tt)*) -> $V:ty,)*) => {
 
         use std::mem;
         use crate::{
@@ -263,7 +318,7 @@ macro_rules! define_queries_inner {
         #[allow(nonstandard_style)]
         #[derive(Clone, Debug)]
         pub enum Query<$tcx> {
-            $($(#[$attr])* $name($K)),*
+            $($(#[$attr])* $name($($K)*)),*
         }
 
         impl<$tcx> Query<$tcx> {
@@ -321,7 +376,7 @@ macro_rules! define_queries_inner {
         }
 
         $(impl<$tcx> QueryConfig<TyCtxt<$tcx>> for queries::$name<$tcx> {
-            type Key = $K;
+            type Key = $($K)*;
             type Value = $V;
             const NAME: &'static str = stringify!($name);
             const CATEGORY: ProfileCategory = $category;
@@ -332,7 +387,7 @@ macro_rules! define_queries_inner {
             const EVAL_ALWAYS: bool = is_eval_always!([$($modifiers)*]);
             const DEP_KIND: dep_graph::DepKind = dep_graph::DepKind::$node;
 
-            type Cache = query_storage!([$($modifiers)*][$K, $V]);
+            type Cache = query_storage!([$($modifiers)*][$($K)*, $V]);
 
             #[inline(always)]
             fn query_state<'a>(tcx: TyCtxt<$tcx>) -> &'a QueryState<TyCtxt<$tcx>, Self::Cache> {
@@ -377,12 +432,8 @@ macro_rules! define_queries_inner {
             pub tcx: TyCtxt<'tcx>,
         }
 
-        impl TyCtxtEnsure<$tcx> {
-            $($(#[$attr])*
-            #[inline(always)]
-            pub fn $name(self, key: $K) {
-                ensure_query::<queries::$name<'_>, _>(self.tcx, key)
-            })*
+        impl TyCtxtEnsure<'tcx> {
+            $( define_query_helper!(TyCtxtEnsure<'tcx>, $(#[$attr])* $name($($K)*) -> $V); )*
         }
 
         #[derive(Copy, Clone)]
@@ -419,11 +470,7 @@ macro_rules! define_queries_inner {
                 }
             }
 
-            $($(#[$attr])*
-            #[inline(always)]
-            pub fn $name(self, key: $K) -> $V {
-                self.at(DUMMY_SP).$name(key)
-            })*
+            $( define_query_helper!(TyCtxt<$tcx>, $(#[$attr])* $name($($K)*) -> $V); )*
 
             /// All self-profiling events generated by the query engine use
             /// virtual `StringId`s for their `event_id`. This method makes all
@@ -456,16 +503,12 @@ macro_rules! define_queries_inner {
         }
 
         impl TyCtxtAt<$tcx> {
-            $($(#[$attr])*
-            #[inline(always)]
-            pub fn $name(self, key: $K) -> $V {
-                get_query::<queries::$name<'_>, _>(self.tcx, self.span, key)
-            })*
+            $( define_query_helper!(TyCtxtAt<$tcx>, $(#[$attr])* $name($($K)*) -> $V); )*
         }
 
         define_provider_struct! {
             tcx: $tcx,
-            input: ($(([$($modifiers)*] [$name] [$K] [$V]))*)
+            input: ($(([$($modifiers)*] [$name] [$($K)*] [$V]))*)
         }
 
         impl<$tcx> Copy for Providers<$tcx> {}