Repository: rust-lang/rust
Pull request: `define_callbacks` tweaks

Details in the individual commits.

r? @Zalathar

Changed files:
filename: compiler/rustc_middle/src/queries.rs
patch: @@ -101,7 +101,7 @@ use crate::mir::mono::{
     CodegenUnit, CollectionMode, MonoItem, MonoItemPartitions, NormalizationErrorInMono,
 };
 use crate::query::describe_as_module;
-use crate::query::plumbing::{define_callbacks, query_helper_param_ty};
+use crate::query::plumbing::{define_callbacks, maybe_into_query_key};
 use crate::traits::query::{
     CanonicalAliasGoal, CanonicalDropckOutlivesGoal, CanonicalImpliedOutlivesBoundsGoal,
     CanonicalMethodAutoderefStepsGoal, CanonicalPredicateGoal, CanonicalTypeOpAscribeUserTypeGoal,

filename: compiler/rustc_middle/src/query/plumbing.rs
patch: @@ -201,6 +201,13 @@ pub struct TyCtxtEnsureDone<'tcx> {
 }
 
 impl<'tcx> TyCtxt<'tcx> {
+    /// Returns a transparent wrapper for `TyCtxt` which uses
+    /// `span` as the location of queries performed through it.
+    #[inline(always)]
+    pub fn at(self, span: Span) -> TyCtxtAt<'tcx> {
+        TyCtxtAt { tcx: self, span }
+    }
+
     /// FIXME: `ensure_ok`'s effects are subtle. Is this comment fully accurate?
     ///
     /// Wrapper that calls queries in a special "ensure OK" mode, for callers
@@ -258,16 +265,9 @@ impl<'tcx> TyCtxt<'tcx> {
     pub fn ensure_done(self) -> TyCtxtEnsureDone<'tcx> {
         TyCtxtEnsureDone { tcx: self }
     }
-
-    /// Returns a transparent wrapper for `TyCtxt` which uses
-    /// `span` as the location of queries performed through it.
-    #[inline(always)]
-    pub fn at(self, span: Span) -> TyCtxtAt<'tcx> {
-        TyCtxtAt { tcx: self, span }
-    }
 }
 
-macro_rules! query_helper_param_ty {
+macro_rules! maybe_into_query_key {
     (DefId) => { impl $crate::query::IntoQueryKey<DefId> };
     (LocalDefId) => { impl $crate::query::IntoQueryKey<LocalDefId> };
     ($K:ty) => { $K };
@@ -276,7 +276,7 @@ macro_rules! query_helper_param_ty {
 macro_rules! define_callbacks {
     (
         // You might expect the key to be `$K:ty`, but it needs to be `$($K:tt)*` so that
-        // `query_helper_param_ty!` can match on specific type names.
+        // `maybe_into_query_key!` can match on specific type names.
         queries {
             $(
                 $(#[$attr:meta])*
@@ -323,6 +323,9 @@ macro_rules! define_callbacks {
                 #[cfg(not($arena_cache))]
                 pub type ProvidedValue<'tcx> = Value<'tcx>;
 
+                pub type Cache<'tcx> =
+                    <Key<'tcx> as $crate::query::QueryKey>::Cache<Erased<Value<'tcx>>>;
+
                 /// This helper function takes a value returned by the query provider
                 /// (or loaded from disk, or supplied by query feeding), allocates
                 /// it in an arena if requested by the `arena_cache` modifier, and
@@ -354,9 +357,6 @@ macro_rules! define_callbacks {
                     erase::erase_val(value)
                 }
 
-                pub type Cache<'tcx> =
-                    <Key<'tcx> as $crate::query::QueryKey>::Cache<Erased<Value<'tcx>>>;
-
                 // 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
                 #[cfg(target_pointer_width = "64")]
@@ -390,116 +390,8 @@ macro_rules! define_callbacks {
             }
         )*
 
-        /// Holds per-query arenas for queries with the `arena_cache` modifier.
-        #[derive(Default)]
-        pub struct QueryArenas<'tcx> {
-            $(
-                // Use the `ArenaCached` helper trait to determine the arena's value type.
-                #[cfg($arena_cache)]
-                pub $name: TypedArena<
-                    <$V as $crate::query::arena_cached::ArenaCached<'tcx>>::Allocated,
-                >,
-            )*
-        }
-
-        impl<'tcx> $crate::query::TyCtxtEnsureOk<'tcx> {
-            $(
-                $(#[$attr])*
-                #[inline(always)]
-                pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
-                    crate::query::inner::query_ensure_ok_or_done(
-                        self.tcx,
-                        &self.tcx.query_system.query_vtables.$name,
-                        $crate::query::IntoQueryKey::into_query_key(key),
-                        $crate::query::EnsureMode::Ok,
-                    )
-                }
-            )*
-        }
-
-        // Only defined when the `ensure_result` modifier is present.
-        impl<'tcx> $crate::query::TyCtxtEnsureResult<'tcx> {
-            $(
-                #[cfg($returns_error_guaranteed)]
-                $(#[$attr])*
-                #[inline(always)]
-                pub fn $name(
-                    self,
-                    key: query_helper_param_ty!($($K)*),
-                ) -> Result<(), rustc_errors::ErrorGuaranteed> {
-                    crate::query::inner::query_ensure_result(
-                        self.tcx,
-                        &self.tcx.query_system.query_vtables.$name,
-                        $crate::query::IntoQueryKey::into_query_key(key),
-                    )
-                }
-            )*
-        }
-
-        impl<'tcx> $crate::query::TyCtxtEnsureDone<'tcx> {
-            $(
-                $(#[$attr])*
-                #[inline(always)]
-                pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
-                    crate::query::inner::query_ensure_ok_or_done(
-                        self.tcx,
-                        &self.tcx.query_system.query_vtables.$name,
-                        $crate::query::IntoQueryKey::into_query_key(key),
-                        $crate::query::EnsureMode::Done,
-                    );
-                }
-            )*
-        }
-
-        impl<'tcx> TyCtxt<'tcx> {
-            $(
-                $(#[$attr])*
-                #[inline(always)]
-                #[must_use]
-                pub fn $name(self, key: query_helper_param_ty!($($K)*)) -> $V {
-                    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 {
-                    use $crate::query::{erase, inner};
-
-                    erase::restore_val::<$V>(inner::query_get_at(
-                        self.tcx,
-                        self.span,
-                        &self.tcx.query_system.query_vtables.$name,
-                        $crate::query::IntoQueryKey::into_query_key(key),
-                    ))
-                }
-            )*
-        }
-
-        $(
-            #[cfg($feedable)]
-            impl<'tcx, K: $crate::query::IntoQueryKey<$name::Key<'tcx>> + Copy>
-                TyCtxtFeed<'tcx, K>
-            {
-                $(#[$attr])*
-                #[inline(always)]
-                pub fn $name(self, value: $name::ProvidedValue<'tcx>) {
-                    let key = self.key().into_query_key();
-                    let erased_value = $name::provided_to_erased(self.tcx, value);
-                    $crate::query::inner::query_feed(
-                        self.tcx,
-                        &self.tcx.query_system.query_vtables.$name,
-                        key,
-                        erased_value,
-                    );
-                }
-            }
-        )*
-
-        /// Identifies a query by kind and key. This is in contrast to `QueryJobId` which is just a number.
+        /// Identifies a query by kind and key. This is in contrast to `QueryJobId` which is just a
+        /// number.
         #[allow(non_camel_case_types)]
         #[derive(Clone, Copy, Debug)]
         pub enum TaggedQueryKey<'tcx> {
@@ -528,7 +420,8 @@ macro_rules! define_callbacks {
             pub fn description(&self, tcx: TyCtxt<'tcx>) -> String {
                 let (name, description) = ty::print::with_no_queries!(match self {
                     $(
-                        TaggedQueryKey::$name(key) => (stringify!($name), _description_fns::$name(tcx, *key)),
+                        TaggedQueryKey::$name(key) =>
+                            (stringify!($name), _description_fns::$name(tcx, *key)),
                     )*
                 });
                 if tcx.sess.verbose_internals() {
@@ -550,22 +443,29 @@ macro_rules! define_callbacks {
                 }
                 match self {
                     $(
-                        TaggedQueryKey::$name(key) => crate::query::QueryKey::default_span(key, tcx),
+                        TaggedQueryKey::$name(key) =>
+                            $crate::query::QueryKey::default_span(key, tcx),
                     )*
                 }
             }
 
             pub fn def_kind(&self, tcx: TyCtxt<'tcx>) -> Option<DefKind> {
                 // This is used to reduce code generation as it
                 // can be reused for queries with the same key type.
-                fn inner<'tcx>(key: &impl crate::query::QueryKey, tcx: TyCtxt<'tcx>) -> Option<DefKind> {
-                    key.key_as_def_id().and_then(|def_id| def_id.as_local()).map(|def_id| tcx.def_kind(def_id))
+                fn inner<'tcx>(key: &impl $crate::query::QueryKey, tcx: TyCtxt<'tcx>)
+                    -> Option<DefKind>
+                {
+                    key
+                        .key_as_def_id()
+                        .and_then(|def_id| def_id.as_local())
+                        .map(|def_id| tcx.def_kind(def_id))
                 }
 
                 if let TaggedQueryKey::def_kind(..) = self {
                     // Try to avoid infinite recursion.
                     return None
                 }
+
                 match self {
                     $(
                         TaggedQueryKey::$name(key) => inner(key, tcx),
@@ -577,7 +477,19 @@ macro_rules! define_callbacks {
         /// Holds a `QueryVTable` for each query.
         pub struct QueryVTables<'tcx> {
             $(
-                pub $name: crate::query::QueryVTable<'tcx, $name::Cache<'tcx>>,
+                pub $name: $crate::query::QueryVTable<'tcx, $name::Cache<'tcx>>,
+            )*
+        }
+
+        /// Holds per-query arenas for queries with the `arena_cache` modifier.
+        #[derive(Default)]
+        pub struct QueryArenas<'tcx> {
+            $(
+                // Use the `ArenaCached` helper trait to determine the arena's value type.
+                #[cfg($arena_cache)]
+                pub $name: TypedArena<
+                    <$V as $crate::query::arena_cached::ArenaCached<'tcx>>::Allocated,
+                >,
             )*
         }
 
@@ -637,12 +549,108 @@ macro_rules! define_callbacks {
         impl Clone for ExternProviders {
             fn clone(&self) -> Self { *self }
         }
+
+        impl<'tcx> TyCtxt<'tcx> {
+            $(
+                $(#[$attr])*
+                #[inline(always)]
+                #[must_use]
+                pub fn $name(self, key: maybe_into_query_key!($($K)*)) -> $V {
+                    self.at(DUMMY_SP).$name(key)
+                }
+            )*
+        }
+
+        impl<'tcx> $crate::query::TyCtxtAt<'tcx> {
+            $(
+                $(#[$attr])*
+                #[inline(always)]
+                pub fn $name(self, key: maybe_into_query_key!($($K)*)) -> $V {
+                    use $crate::query::{erase, inner};
+
+                    erase::restore_val::<$V>(inner::query_get_at(
+                        self.tcx,
+                        self.span,
+                        &self.tcx.query_system.query_vtables.$name,
+                        $crate::query::IntoQueryKey::into_query_key(key),
+                    ))
+                }
+            )*
+        }
+
+        impl<'tcx> $crate::query::TyCtxtEnsureOk<'tcx> {
+            $(
+                $(#[$attr])*
+                #[inline(always)]
+                pub fn $name(self, key: maybe_into_query_key!($($K)*)) {
+                    $crate::query::inner::query_ensure_ok_or_done(
+                        self.tcx,
+                        &self.tcx.query_system.query_vtables.$name,
+                        $crate::query::IntoQueryKey::into_query_key(key),
+                        $crate::query::EnsureMode::Ok,
+                    )
+                }
+            )*
+        }
+
+        // Only defined when the `returns_error_guaranteed` modifier is present.
+        impl<'tcx> $crate::query::TyCtxtEnsureResult<'tcx> {
+            $(
+                #[cfg($returns_error_guaranteed)]
+                $(#[$attr])*
+                #[inline(always)]
+                pub fn $name(
+                    self,
+                    key: maybe_into_query_key!($($K)*),
+                ) -> Result<(), rustc_errors::ErrorGuaranteed> {
+                    $crate::query::inner::query_ensure_result(
+                        self.tcx,
+                        &self.tcx.query_system.query_vtables.$name,
+                        $crate::query::IntoQueryKey::into_query_key(key),
+                    )
+                }
+            )*
+        }
+
+        impl<'tcx> $crate::query::TyCtxtEnsureDone<'tcx> {
+            $(
+                $(#[$attr])*
+                #[inline(always)]
+                pub fn $name(self, key: maybe_into_query_key!($($K)*)) {
+                    $crate::query::inner::query_ensure_ok_or_done(
+                        self.tcx,
+                        &self.tcx.query_system.query_vtables.$name,
+                        $crate::query::IntoQueryKey::into_query_key(key),
+                        $crate::query::EnsureMode::Done,
+                    );
+                }
+            )*
+        }
+
+        $(
+            // Only defined when the `feedable` modifier is present.
+            #[cfg($feedable)]
+            impl<'tcx, K: $crate::query::IntoQueryKey<$name::Key<'tcx>> + Copy>
+                TyCtxtFeed<'tcx, K>
+            {
+                $(#[$attr])*
+                #[inline(always)]
+                pub fn $name(self, value: $name::ProvidedValue<'tcx>) {
+                    $crate::query::inner::query_feed(
+                        self.tcx,
+                        &self.tcx.query_system.query_vtables.$name,
+                        self.key().into_query_key(),
+                        $name::provided_to_erased(self.tcx, value),
+                    );
+                }
+            }
+        )*
     };
 }
 
 // Re-export `macro_rules!` macros as normal items, so that they can be imported normally.
 pub(crate) use define_callbacks;
-pub(crate) use query_helper_param_ty;
+pub(crate) use maybe_into_query_key;
 
 #[cold]
 pub(crate) fn default_query(name: &str, key: &dyn std::fmt::Debug) -> ! {