Match query key types as $K:tt instead of $($K:tt)*

compiler/rustc_macros/src/query.rs

@@ -30,14 +30,28 @@ fn check_attributes(attrs: Vec<Attribute>) -> Result<Vec<Attribute>> {
     attrs.into_iter().map(inner).collect()
 }
 
-/// A compiler query. `query ... { ... }`
+/// Declaration of a compiler query.
+///
+/// ```ignore (illustrative)
+/// /// Doc comment for `my_query`.
+/// //  ^^^^^^^^^^^^^^^^^^^^^^^^^^^              doc_comments
+/// query my_query(key: DefId) -> Value { anon }
+/// //    ^^^^^^^^                               name
+/// //             ^^^                           key_pat
+/// //                  ^^^^^                    key_ty
+/// //                         ^^^^^^^^          return_ty
+/// //                                    ^^^^   modifiers
+/// ```
 struct Query {
     doc_comments: Vec<Attribute>,
-    modifiers: QueryModifiers,
     name: Ident,
-    key: Pat,
-    arg: Type,
-    result: ReturnType,
+
+    /// Parameter name for the key, or an arbitrary irrefutable pattern (e.g. `_`).
+    key_pat: Pat,
+    key_ty: Type,
+    return_ty: ReturnType,
+
+    modifiers: QueryModifiers,
 }
 
 impl Parse for Query {
@@ -47,26 +61,30 @@ impl Parse for Query {
         // Parse the query declaration. Like `query type_of(key: DefId) -> Ty<'tcx>`
         input.parse::<kw::query>()?;
         let name: Ident = input.parse()?;
-        let arg_content;
-        parenthesized!(arg_content in input);
-        let key = Pat::parse_single(&arg_content)?;
-        arg_content.parse::<Token![:]>()?;
-        let arg = arg_content.parse()?;
-        let _ = arg_content.parse::<Option<Token![,]>>()?;
-        let result = input.parse()?;
+
+        // `(key: DefId)`
+        let parens_content;
+        parenthesized!(parens_content in input);
+        let key_pat = Pat::parse_single(&parens_content)?;
+        parens_content.parse::<Token![:]>()?;
+        let key_ty = parens_content.parse::<Type>()?;
+        let _trailing_comma = parens_content.parse::<Option<Token![,]>>()?;
+
+        // `-> Value`
+        let return_ty = input.parse::<ReturnType>()?;
 
         // Parse the query modifiers
-        let content;
-        braced!(content in input);
-        let modifiers = parse_query_modifiers(&content)?;
+        let braces_content;
+        braced!(braces_content in input);
+        let modifiers = parse_query_modifiers(&braces_content)?;
 
         // If there are no doc-comments, give at least some idea of what
         // it does by showing the query description.
         if doc_comments.is_empty() {
             doc_comments.push(doc_comment_from_desc(&modifiers.desc.expr_list)?);
         }
 
-        Ok(Query { doc_comments, modifiers, name, key, arg, result })
+        Ok(Query { doc_comments, modifiers, name, key_pat, key_ty, return_ty })
     }
 }
 
@@ -288,7 +306,7 @@ struct HelperTokenStreams {
 }
 
 fn make_helpers_for_query(query: &Query, streams: &mut HelperTokenStreams) {
-    let Query { name, key, modifiers, arg, .. } = &query;
+    let Query { name, key_pat, key_ty, modifiers, .. } = &query;
 
     // Replace span for `name` to make rust-analyzer ignore it.
     let mut erased_name = name.clone();
@@ -301,7 +319,7 @@ fn make_helpers_for_query(query: &Query, streams: &mut HelperTokenStreams) {
         streams.cache_on_disk_if_fns_stream.extend(quote! {
             #[allow(unused_variables, rustc::pass_by_value)]
             #[inline]
-            pub fn #erased_name<'tcx>(#tcx: TyCtxt<'tcx>, #key: &crate::queries::#name::Key<'tcx>) -> bool
+            pub fn #erased_name<'tcx>(#tcx: TyCtxt<'tcx>, #key_pat: &crate::queries::#name::Key<'tcx>) -> bool
             #block
         });
     }
@@ -311,8 +329,8 @@ fn make_helpers_for_query(query: &Query, streams: &mut HelperTokenStreams) {
 
     let desc = quote! {
         #[allow(unused_variables)]
-        pub fn #erased_name<'tcx>(tcx: TyCtxt<'tcx>, key: #arg) -> String {
-            let (#tcx, #key) = (tcx, key);
+        pub fn #erased_name<'tcx>(tcx: TyCtxt<'tcx>, key: #key_ty) -> String {
+            let (#tcx, #key_pat) = (tcx, key);
             format!(#expr_list)
         }
     };
@@ -373,7 +391,7 @@ fn add_to_analyzer_stream(query: &Query, analyzer_stream: &mut proc_macro2::Toke
     let mut erased_name = name.clone();
     erased_name.set_span(Span::call_site());
 
-    let result = &query.result;
+    let result = &query.return_ty;
 
     // This dead code exists to instruct rust-analyzer about the link between the `rustc_queries`
     // query names and the corresponding produced provider. The issue is that by nature of this
@@ -417,19 +435,20 @@ pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {
     }
 
     for query in queries.0 {
-        let Query { name, arg, modifiers, .. } = &query;
-        let result_full = &query.result;
-        let result = match query.result {
+        let Query { doc_comments, name, key_ty, return_ty, modifiers, .. } = &query;
+
+        // Normalize an absent return type into `-> ()` to make macro-rules parsing easier.
+        let return_ty = match return_ty {
             ReturnType::Default => quote! { -> () },
-            _ => quote! { #result_full },
+            ReturnType::Type(..) => quote! { #return_ty },
         };
 
-        let mut attributes = Vec::new();
+        let mut modifiers_out = vec![];
 
         macro_rules! passthrough {
             ( $( $modifier:ident ),+ $(,)? ) => {
                 $( if let Some($modifier) = &modifiers.$modifier {
-                    attributes.push(quote! { (#$modifier) });
+                    modifiers_out.push(quote! { (#$modifier) });
                 }; )+
             }
         }
@@ -452,7 +471,7 @@ pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {
         // on a synthetic `(cache_on_disk)` modifier that can be inspected by
         // macro-rules macros.
         if modifiers.cache_on_disk_if.is_some() {
-            attributes.push(quote! { (cache_on_disk) });
+            modifiers_out.push(quote! { (cache_on_disk) });
         }
 
         // This uses the span of the query definition for the commas,
@@ -462,12 +481,15 @@ pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {
         // at the entire `rustc_queries!` invocation, which wouldn't
         // be very useful.
         let span = name.span();
-        let attribute_stream = quote_spanned! {span=> #(#attributes),*};
-        let doc_comments = &query.doc_comments;
-        // Add the query to the group
+        let modifiers_stream = quote_spanned! { span => #(#modifiers_out),* };
+
+        // Add the query to the group.
+        // Surround `key_ty` with parentheses so it can be matched as a single
+        // token tree `$K:tt`, instead of the clunkier `$($K:tt)*`.
         query_stream.extend(quote! {
             #(#doc_comments)*
-            [#attribute_stream] fn #name(#arg) #result,
+            [#modifiers_stream]
+            fn #name((#key_ty)) #return_ty,
         });
 
         if let Some(feedable) = &modifiers.feedable {
@@ -482,7 +504,8 @@ pub(super) fn rustc_queries(input: TokenStream) -> TokenStream {
                 "Query {name} cannot be both `feedable` and `eval_always`."
             );
             feedable_queries.extend(quote! {
-                [#attribute_stream] fn #name(#arg) #result,
+                [#modifiers_stream]
+                fn #name((#key_ty)) #return_ty,
             });
         }
 

compiler/rustc_middle/src/dep_graph/dep_node.rs

@@ -360,7 +360,8 @@ macro_rules! define_dep_nodes {
     (
         $(
             $(#[$attr:meta])*
-            [$($modifiers:tt)*] fn $variant:ident($($K:tt)*) -> $V:ty,
+            [$($modifiers:tt)*]
+            fn $variant:ident($K:tt) -> $V:ty,
         )*
     ) => {
 
@@ -430,15 +431,15 @@ macro_rules! define_dep_nodes {
 // that aren't queries.
 rustc_with_all_queries!(define_dep_nodes![
     /// We use this for most things when incr. comp. is turned off.
-    [] fn Null() -> (),
+    [] fn Null(()) -> (),
     /// We use this to create a forever-red node.
-    [] fn Red() -> (),
-    [] fn SideEffect() -> (),
-    [] fn AnonZeroDeps() -> (),
-    [] fn TraitSelect() -> (),
-    [] fn CompileCodegenUnit() -> (),
-    [] fn CompileMonoItem() -> (),
-    [] fn Metadata() -> (),
+    [] fn Red(()) -> (),
+    [] fn SideEffect(()) -> (),
+    [] fn AnonZeroDeps(()) -> (),
+    [] fn TraitSelect(()) -> (),
+    [] fn CompileCodegenUnit(()) -> (),
+    [] fn CompileMonoItem(()) -> (),
+    [] fn Metadata(()) -> (),
 ]);
 
 // WARNING: `construct` is generic and does not know that `CompileCodegenUnit` takes `Symbol`s as keys.

compiler/rustc_middle/src/query/plumbing.rs

@@ -276,9 +276,16 @@ macro_rules! query_ensure_select {
     };
 }
 
-macro_rules! query_helper_param_ty {
-    (DefId) => { impl $crate::query::IntoQueryParam<DefId> };
-    (LocalDefId) => { impl $crate::query::IntoQueryParam<LocalDefId> };
+/// Expands to `impl IntoQueryParam<$K>` for a small allowlist of key types
+/// that benefit from the implicit conversion, or otherwise expands to `$K`
+/// itself.
+///
+/// This macro is why `define_callbacks!` needs to parse the key type as
+/// `$K:tt` and not `$K:ty`, because otherwise it would be unable to inspect
+/// the tokens inside `$K`.
+macro_rules! maybe_into_query_param {
+    ((DefId)) => { impl $crate::query::IntoQueryParam<DefId> };
+    ((LocalDefId)) => { impl $crate::query::IntoQueryParam<LocalDefId> };
     ($K:ty) => { $K };
 }
 
@@ -297,14 +304,12 @@ macro_rules! query_if_arena {
 /// If `separate_provide_extern`, then the key can be projected to its
 /// local key via `<$K as AsLocalKey>::LocalKey`.
 macro_rules! local_key_if_separate_extern {
-    ([] $($K:tt)*) => {
-        $($K)*
-    };
-    ([(separate_provide_extern) $($rest:tt)*] $($K:tt)*) => {
-        <$($K)* as $crate::query::AsLocalKey>::LocalKey
+    ([] $K:ty) => { $K };
+    ([(separate_provide_extern) $($rest:tt)*] $K:ty) => {
+        <$K as $crate::query::AsLocalKey>::LocalKey
     };
-    ([$other:tt $($modifiers:tt)*] $($K:tt)*) => {
-        local_key_if_separate_extern!([$($modifiers)*] $($K)*)
+    ([$other:tt $($modifiers:tt)*] $K:ty) => {
+        local_key_if_separate_extern!([$($modifiers)*] $K)
     };
 }
 
@@ -349,19 +354,26 @@ macro_rules! separate_provide_extern_default {
 
 macro_rules! define_callbacks {
     (
+        // Match the key type as `tt` so that we can inspect its tokens.
+        // (See `maybe_into_query_param`.)
+        //
+        // We don't need to write `$($K:tt)*` because `rustc_macros::query`
+        // takes care to wrap the key type in an extra set of parentheses,
+        // so it's always a single token tree.
         $(
             $(#[$attr:meta])*
-            [$($modifiers:tt)*] fn $name:ident($($K:tt)*) -> $V:ty,
+            [$($modifiers:tt)*]
+            fn $name:ident($K:tt) -> $V:ty,
         )*
     ) => {
         $(#[allow(unused_lifetimes)] pub mod $name {
             use super::*;
             use $crate::query::erase::{self, Erased};
 
-            pub type Key<'tcx> = $($K)*;
+            pub type Key<'tcx> = $K;
             pub type Value<'tcx> = $V;
 
-            pub type LocalKey<'tcx> = local_key_if_separate_extern!([$($modifiers)*] $($K)*);
+            pub type LocalKey<'tcx> = local_key_if_separate_extern!([$($modifiers)*] $K);
 
             /// This type alias specifies the type returned from query providers and the type
             /// used for decoding. For regular queries this is the declared returned type `V`,
@@ -397,7 +409,7 @@ macro_rules! define_callbacks {
                 erase::erase_val(value)
             }
 
-            pub type Storage<'tcx> = <$($K)* as $crate::query::Key>::Cache<Erased<$V>>;
+            pub type Storage<'tcx> = <$K as $crate::query::Key>::Cache<Erased<$V>>;
 
             // Ensure that keys grow no larger than 88 bytes by accident.
             // Increase this limit if necessary, but do try to keep the size low if possible
@@ -408,7 +420,7 @@ macro_rules! define_callbacks {
                         "the query `",
                         stringify!($name),
                         "` has a key type `",
-                        stringify!($($K)*),
+                        stringify!($K),
                         "` that is too large"
                     ));
                 }
@@ -455,7 +467,7 @@ macro_rules! define_callbacks {
             #[inline(always)]
             pub fn $name(
                 self,
-                key: query_helper_param_ty!($($K)*),
+                key: maybe_into_query_param!($K),
             ) -> ensure_ok_result!([$($modifiers)*]) {
                 query_ensure_select!(
                     [$($modifiers)*]
@@ -471,7 +483,10 @@ macro_rules! define_callbacks {
         impl<'tcx> $crate::query::TyCtxtEnsureDone<'tcx> {
             $($(#[$attr])*
             #[inline(always)]
-            pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
+            pub fn $name(
+                self,
+                key: maybe_into_query_param!($K),
+            ) {
                 crate::query::inner::query_ensure(
                     self.tcx,
                     self.tcx.query_system.fns.engine.$name,
@@ -486,17 +501,22 @@ macro_rules! define_callbacks {
             $($(#[$attr])*
             #[inline(always)]
             #[must_use]
-            pub fn $name(self, key: query_helper_param_ty!($($K)*)) -> $V
-            {
+            pub fn $name(
+                self,
+                key: maybe_into_query_param!($K),
+            ) -> $V {
+                let key = $crate::query::IntoQueryParam::into_query_param(key);
                 self.at(DUMMY_SP).$name(key)
             })*
         }
 
         impl<'tcx> $crate::query::TyCtxtAt<'tcx> {
             $($(#[$attr])*
             #[inline(always)]
-            pub fn $name(self, key: query_helper_param_ty!($($K)*)) -> $V
-            {
+            pub fn $name(
+                self,
+                key: maybe_into_query_param!($K),
+            ) -> $V {
                 use $crate::query::{erase, inner};
 
                 erase::restore_val::<$V>(inner::query_get_at(
@@ -521,7 +541,7 @@ macro_rules! define_callbacks {
         #[derive(Default)]
         pub struct QueryStates<'tcx> {
             $(
-                pub $name: $crate::query::QueryState<'tcx, $($K)*>,
+                pub $name: $crate::query::QueryState<'tcx, $K>,
             )*
         }
 
@@ -578,8 +598,14 @@ macro_rules! define_callbacks {
 }
 
 macro_rules! define_feedable {
-    ($($(#[$attr:meta])* [$($modifiers:tt)*] fn $name:ident($($K:tt)*) -> $V:ty,)*) => {
-        $(impl<'tcx, K: $crate::query::IntoQueryParam<$($K)*> + Copy> TyCtxtFeed<'tcx, K> {
+    (
+        $(
+            $(#[$attr:meta])*
+            [$($modifiers:tt)*]
+            fn $name:ident($K:tt) -> $V:ty,
+        )*
+    ) => {
+        $(impl<'tcx, K: $crate::query::IntoQueryParam<$K> + Copy> TyCtxtFeed<'tcx, K> {
             $(#[$attr])*
             #[inline(always)]
             pub fn $name(self, value: $name::ProvidedValue<'tcx>) {

compiler/rustc_query_impl/src/plumbing.rs

@@ -479,7 +479,8 @@ macro_rules! define_queries {
     (
         $(
             $(#[$attr:meta])*
-            [$($modifiers:tt)*] fn $name:ident($($K:tt)*) -> $V:ty,
+            [$($modifiers:tt)*]
+            fn $name:ident($K:tt) -> $V:ty,
         )*
     ) => {