Rollup merge of #154591 - zetanumbers:refactor_query_fns_and_traits, r=cjgillot

JonathanBrouwer · web-flow · commit e8659d7988cc · 2026-05-30T11:40:13.000+02:00
Remove `will_cache_on_disk_for_key_fn` Refactors `QueryVtables` and removes `will_cache_on_disk_for_key_fn` as unnecessary indirection. Based and blocked on #154576. Expecting perf improvements.
diff --git a/compiler/rustc_middle/src/query/keys.rs b/compiler/rustc_middle/src/query/keys.rs
@@ -36,6 +36,8 @@ pub trait QueryKey: Sized + QueryKeyBounds {
     /// [`QueryCache`]: rustc_middle::query::QueryCache
     type Cache<V> = DefaultCache<Self, V>;
 
+    type LocalQueryKey = !;
+
     /// In the event that a cycle occurs, if no explicit span has been
     /// given for a query with key `self`, what span should we use?
     fn default_span(&self, tcx: TyCtxt<'_>) -> Span;
@@ -45,14 +47,12 @@ pub trait QueryKey: Sized + QueryKeyBounds {
     fn key_as_def_id(&self) -> Option<DefId> {
         None
     }
-}
-
-pub trait AsLocalQueryKey: QueryKey {
-    type LocalQueryKey;
 
     /// Given an instance of this key, what crate is it referring to?
     /// This is used to find the provider.
-    fn as_local_key(&self) -> Option<Self::LocalQueryKey>;
+    fn as_local_key(&self) -> Option<Self::LocalQueryKey> {
+        None
+    }
 }
 
 impl QueryKey for () {
@@ -96,13 +96,11 @@ impl<'tcx> QueryKey for ty::LitToConstInput<'tcx> {
 impl QueryKey for CrateNum {
     type Cache<V> = VecCache<Self, V, DepNodeIndex>;
 
+    type LocalQueryKey = LocalCrate;
+
     fn default_span(&self, _: TyCtxt<'_>) -> Span {
         DUMMY_SP
     }
-}
-
-impl AsLocalQueryKey for CrateNum {
-    type LocalQueryKey = LocalCrate;
 
     #[inline(always)]
     fn as_local_key(&self) -> Option<Self::LocalQueryKey> {
@@ -136,6 +134,7 @@ impl QueryKey for LocalDefId {
 
 impl QueryKey for DefId {
     type Cache<V> = DefIdCache<V>;
+    type LocalQueryKey = LocalDefId;
 
     fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
         tcx.def_span(*self)
@@ -145,10 +144,6 @@ impl QueryKey for DefId {
     fn key_as_def_id(&self) -> Option<DefId> {
         Some(*self)
     }
-}
-
-impl AsLocalQueryKey for DefId {
-    type LocalQueryKey = LocalDefId;
 
     #[inline(always)]
     fn as_local_key(&self) -> Option<Self::LocalQueryKey> {
@@ -197,13 +192,11 @@ impl QueryKey for (LocalDefId, LocalDefId, Ident) {
 }
 
 impl QueryKey for (CrateNum, DefId) {
+    type LocalQueryKey = DefId;
+
     fn default_span(&self, tcx: TyCtxt<'_>) -> Span {
         self.1.default_span(tcx)
     }
-}
-
-impl AsLocalQueryKey for (CrateNum, DefId) {
-    type LocalQueryKey = DefId;
 
     #[inline(always)]
     fn as_local_key(&self) -> Option<Self::LocalQueryKey> {
@@ -212,13 +205,11 @@ impl AsLocalQueryKey for (CrateNum, DefId) {
 }
 
 impl QueryKey for (CrateNum, SimplifiedType) {
+    type LocalQueryKey = SimplifiedType;
+
     fn default_span(&self, _: TyCtxt<'_>) -> Span {
         DUMMY_SP
     }
-}
-
-impl AsLocalQueryKey for (CrateNum, SimplifiedType) {
-    type LocalQueryKey = SimplifiedType;
 
     #[inline(always)]
     fn as_local_key(&self) -> Option<Self::LocalQueryKey> {
diff --git a/compiler/rustc_middle/src/query/mod.rs b/compiler/rustc_middle/src/query/mod.rs
@@ -3,7 +3,7 @@ use rustc_hir::def_id::LocalDefId;
 pub use self::caches::{DefIdCache, DefaultCache, QueryCache, SingleCache, VecCache};
 pub use self::into_query_key::IntoQueryKey;
 pub use self::job::{QueryJob, QueryJobId, QueryLatch, QueryWaiter};
-pub use self::keys::{AsLocalQueryKey, LocalCrate, QueryKey};
+pub use self::keys::{LocalCrate, QueryKey};
 pub use self::plumbing::{
     ActiveKeyStatus, Cycle, EnsureMode, QueryMode, QueryState, QuerySystem, QueryVTable, TyCtxtAt,
     TyCtxtEnsureDone, TyCtxtEnsureOk, TyCtxtEnsureResult,
diff --git a/compiler/rustc_middle/src/query/plumbing.rs b/compiler/rustc_middle/src/query/plumbing.rs
@@ -14,7 +14,7 @@ use crate::dep_graph::{DepKind, DepNodeIndex, QuerySideEffect, SerializedDepNode
 use crate::ich::StableHashState;
 use crate::queries::{ExternProviders, Providers, QueryArenas, QueryVTables, TaggedQueryKey};
 use crate::query::on_disk_cache::OnDiskCache;
-use crate::query::{IntoQueryKey, QueryCache, QueryJob, QueryStackFrame};
+use crate::query::{IntoQueryKey, QueryCache, QueryJob, QueryKey, QueryStackFrame};
 use crate::ty::{self, TyCtxt};
 
 /// For a particular query, keeps track of "active" keys, i.e. keys whose
@@ -86,6 +86,9 @@ pub struct QueryVTable<'tcx, C: QueryCache> {
     /// True if this query has the `feedable` modifier.
     pub feedable: bool,
 
+    pub cache_on_disk_local: bool,
+    pub separate_provide_extern: bool,
+
     pub dep_kind: DepKind,
     pub state: QueryState<'tcx, C::Key>,
     pub cache: C,
@@ -97,11 +100,9 @@ pub struct QueryVTable<'tcx, C: QueryCache> {
     /// This should be the only code that calls the provider function.
     pub invoke_provider_fn: fn(tcx: TyCtxt<'tcx>, key: C::Key) -> C::Value,
 
-    pub will_cache_on_disk_for_key_fn: fn(key: C::Key) -> bool,
-
     /// Function pointer that tries to load a query value from disk.
     ///
-    /// This should only be called after a successful check of `will_cache_on_disk_for_key_fn`.
+    /// This should only be called after a successful check of [`Self::will_cache_on_disk_for_key`].
     pub try_load_from_disk_fn:
         fn(tcx: TyCtxt<'tcx>, prev_index: SerializedDepNodeIndex) -> Option<C::Value>,
 
@@ -133,6 +134,12 @@ pub struct QueryVTable<'tcx, C: QueryCache> {
     pub execute_query_fn: fn(TyCtxt<'tcx>, Span, C::Key, QueryMode) -> Option<C::Value>,
 }
 
+impl<'tcx, C: QueryCache> QueryVTable<'tcx, C> {
+    pub fn will_cache_on_disk_for_key(&self, key: C::Key) -> bool {
+        self.cache_on_disk_local && (!self.separate_provide_extern || key.as_local_key().is_some())
+    }
+}
+
 impl<'tcx, C: QueryCache> fmt::Debug for QueryVTable<'tcx, C> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         // When debug-printing a query vtable (e.g. for ICE or tracing),
@@ -328,7 +335,7 @@ macro_rules! define_callbacks {
                 /// This query has the `separate_provide_extern` modifier.
                 #[cfg($separate_provide_extern)]
                 pub type LocalKey<'tcx> =
-                    <Key<'tcx> as $crate::query::AsLocalQueryKey>::LocalQueryKey;
+                    <Key<'tcx> as $crate::query::QueryKey>::LocalQueryKey;
                 /// Key type used by provider functions in `local_providers`.
                 #[cfg(not($separate_provide_extern))]
                 pub type LocalKey<'tcx> = Key<'tcx>;
diff --git a/compiler/rustc_query_impl/src/execution.rs b/compiler/rustc_query_impl/src/execution.rs
@@ -498,7 +498,7 @@ fn load_from_disk_or_invoke_provider_green<'tcx, C: QueryCache>(
     debug_assert!(dep_graph_data.is_index_green(prev_index));
 
     // First try to load the result from the on-disk cache. Some things are never cached on disk.
-    let try_value = if (query.will_cache_on_disk_for_key_fn)(key) {
+    let try_value = if query.will_cache_on_disk_for_key(key) {
         let prof_timer = tcx.prof.incr_cache_loading();
         let value = (query.try_load_from_disk_fn)(tcx, prev_index);
         prof_timer.finish_with_query_invocation_id(dep_node_index.into());
@@ -603,7 +603,7 @@ fn ensure_can_skip_execution<'tcx, C: QueryCache>(
                 // needed, which guarantees the query provider will never run
                 // for this key.
                 EnsureMode::Done => {
-                    (query.will_cache_on_disk_for_key_fn)(key)
+                    query.will_cache_on_disk_for_key(key)
                         && loadable_from_disk(tcx, serialized_dep_node_index)
                 }
             }
diff --git a/compiler/rustc_query_impl/src/plumbing.rs b/compiler/rustc_query_impl/src/plumbing.rs
@@ -93,7 +93,7 @@ fn encode_query_values_inner<'a, 'tcx, C, V>(
 
     assert!(all_inactive(&query.state));
     query.cache.for_each(&mut |key, value, dep_node| {
-        if (query.will_cache_on_disk_for_key_fn)(*key) {
+        if query.will_cache_on_disk_for_key(*key) {
             encoder.encode_query_value::<V>(dep_node, &erase::restore_val::<V>(*value));
         }
     });
@@ -152,7 +152,7 @@ pub(crate) fn promote_from_disk_inner<'tcx, C: QueryCache>(
 
     // If the recovered key isn't eligible for cache-on-disk, then there's no
     // value on disk to promote.
-    if !(query.will_cache_on_disk_for_key_fn)(key) {
+    if !query.will_cache_on_disk_for_key(key) {
         return;
     }
 
diff --git a/compiler/rustc_query_impl/src/query_impl.rs b/compiler/rustc_query_impl/src/query_impl.rs
@@ -1,6 +1,6 @@
 use rustc_middle::queries::TaggedQueryKey;
 use rustc_middle::query::erase::{self, Erased};
-use rustc_middle::query::{AsLocalQueryKey, QueryMode, QueryVTable};
+use rustc_middle::query::{QueryKey, QueryMode, QueryVTable};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::Span;
 
@@ -127,20 +127,6 @@ macro_rules! define_queries {
                     }
                 }
 
-                fn will_cache_on_disk_for_key<'tcx>(
-                    _key: rustc_middle::queries::$name::Key<'tcx>,
-                ) -> bool {
-                    cfg_select! {
-                        // If a query has both `cache_on_disk` and `separate_provide_extern`, only
-                        // disk-cache values for "local" keys, i.e. things in the current crate.
-                        all($cache_on_disk, $separate_provide_extern) => {
-                            AsLocalQueryKey::as_local_key(&_key).is_some()
-                        }
-                        all($cache_on_disk, not($separate_provide_extern)) => true,
-                        not($cache_on_disk) => false,
-                    }
-                }
-
                 pub(crate) fn make_query_vtable<'tcx>(incremental: bool)
                     -> QueryVTable<'tcx, rustc_middle::queries::$name::Cache<'tcx>>
                 {
@@ -153,12 +139,11 @@ macro_rules! define_queries {
                         dep_kind: rustc_middle::dep_graph::DepKind::$name,
                         state: Default::default(),
                         cache: Default::default(),
+                        cache_on_disk_local: $cache_on_disk,
+                        separate_provide_extern: $separate_provide_extern,
 
                         invoke_provider_fn: self::invoke_provider_fn::__rust_begin_short_backtrace,
 
-                        will_cache_on_disk_for_key_fn:
-                            $crate::query_impl::$name::will_cache_on_disk_for_key,
-
                         #[cfg($cache_on_disk)]
                         try_load_from_disk_fn: |tcx, prev_index| {
                             use rustc_middle::queries::$name::{ProvidedValue, provided_to_erased};

Original file line number	Diff line number	Diff line change
`@@ -93,7 +93,7 @@ fn encode_query_values_inner<'a, 'tcx, C, V>(`
`93`	`93`
`94`	`94`	`assert!(all_inactive(&query.state));`
`95`	`95`	`query.cache.for_each(&mut \|key, value, dep_node\| {`
`96`		`- if (query.will_cache_on_disk_for_key_fn)(*key) {`
	`96`	`+ if query.will_cache_on_disk_for_key(*key) {`
`97`	`97`	`encoder.encode_query_value::<V>(dep_node, &erase::restore_val::<V>(*value));`
`98`	`98`	`}`
`99`	`99`	`});`
`@@ -152,7 +152,7 @@ pub(crate) fn promote_from_disk_inner<'tcx, C: QueryCache>(`
`152`	`152`
`153`	`153`	`// If the recovered key isn't eligible for cache-on-disk, then there's no`
`154`	`154`	`// value on disk to promote.`
`155`		`- if !(query.will_cache_on_disk_for_key_fn)(key) {`
	`155`	`+ if !query.will_cache_on_disk_for_key(key) {`
`156`	`156`	`return;`
`157`	`157`	`}`
`158`	`158`