Auto merge of #152835 - nnethercote:fix-query-state-query-cache, r=<try>

bors · bors · commit 6efacd91cd7c · 2026-02-19T03:50:59.000Z
Improve how `QueryCache`s and `QueryState`s are stored
diff --git a/compiler/rustc_middle/src/query/inner.rs b/compiler/rustc_middle/src/query/inner.rs
@@ -32,15 +32,15 @@ where
 /// Shared implementation of `tcx.$query(..)` and `tcx.at(span).$query(..)`
 /// for all queries.
 #[inline(always)]
-pub(crate) fn query_get_at<'tcx, Cache>(
+pub(crate) fn query_get_at<'tcx, C>(
     tcx: TyCtxt<'tcx>,
-    execute_query: fn(TyCtxt<'tcx>, Span, Cache::Key, QueryMode) -> Option<Cache::Value>,
-    query_cache: &Cache,
+    execute_query: fn(TyCtxt<'tcx>, Span, C::Key, QueryMode) -> Option<C::Value>,
+    query_cache: &C,
     span: Span,
-    key: Cache::Key,
-) -> Cache::Value
+    key: C::Key,
+) -> C::Value
 where
-    Cache: QueryCache,
+    C: QueryCache,
 {
     match try_get_cached(tcx, query_cache, &key) {
         Some(value) => value,
@@ -51,14 +51,14 @@ where
 /// Shared implementation of `tcx.ensure_ok().$query(..)` for most queries,
 /// and `tcx.ensure_done().$query(..)` for all queries.
 #[inline]
-pub(crate) fn query_ensure<'tcx, Cache>(
+pub(crate) fn query_ensure<'tcx, C>(
     tcx: TyCtxt<'tcx>,
-    execute_query: fn(TyCtxt<'tcx>, Span, Cache::Key, QueryMode) -> Option<Cache::Value>,
-    query_cache: &Cache,
-    key: Cache::Key,
+    execute_query: fn(TyCtxt<'tcx>, Span, C::Key, QueryMode) -> Option<C::Value>,
+    query_cache: &C,
+    key: C::Key,
     check_cache: bool,
 ) where
-    Cache: QueryCache,
+    C: QueryCache,
 {
     if try_get_cached(tcx, query_cache, &key).is_none() {
         execute_query(tcx, DUMMY_SP, key, QueryMode::Ensure { check_cache });
@@ -68,15 +68,15 @@ pub(crate) fn query_ensure<'tcx, Cache>(
 /// Shared implementation of `tcx.ensure_ok().$query(..)` for queries that
 /// have the `return_result_from_ensure_ok` modifier.
 #[inline]
-pub(crate) fn query_ensure_error_guaranteed<'tcx, Cache, T>(
+pub(crate) fn query_ensure_error_guaranteed<'tcx, C, T>(
     tcx: TyCtxt<'tcx>,
-    execute_query: fn(TyCtxt<'tcx>, Span, Cache::Key, QueryMode) -> Option<Cache::Value>,
-    query_cache: &Cache,
-    key: Cache::Key,
+    execute_query: fn(TyCtxt<'tcx>, Span, C::Key, QueryMode) -> Option<C::Value>,
+    query_cache: &C,
+    key: C::Key,
     check_cache: bool,
 ) -> Result<(), ErrorGuaranteed>
 where
-    Cache: QueryCache<Value = Erased<Result<T, ErrorGuaranteed>>>,
+    C: QueryCache<Value = Erased<Result<T, ErrorGuaranteed>>>,
     Result<T, ErrorGuaranteed>: Erasable,
 {
     if let Some(res) = try_get_cached(tcx, query_cache, &key) {
@@ -96,21 +96,20 @@ where
 }
 
 /// Common implementation of query feeding, used by `define_feedable!`.
-pub(crate) fn query_feed<'tcx, Cache>(
+pub(crate) fn query_feed<'tcx, C>(
     tcx: TyCtxt<'tcx>,
     dep_kind: DepKind,
-    query_vtable: &QueryVTable<'tcx, Cache>,
-    cache: &Cache,
-    key: Cache::Key,
-    value: Cache::Value,
+    query_vtable: &QueryVTable<'tcx, C>,
+    key: C::Key,
+    value: C::Value,
 ) where
-    Cache: QueryCache,
-    Cache::Key: DepNodeKey<'tcx>,
+    C: QueryCache,
+    C::Key: DepNodeKey<'tcx>,
 {
     let format_value = query_vtable.format_value;
 
     // Check whether the in-memory cache already has a value for this key.
-    match try_get_cached(tcx, cache, &key) {
+    match try_get_cached(tcx, &query_vtable.query_cache, &key) {
         Some(old) => {
             // The query already has a cached value for this key.
             // That's OK if both values are the same, i.e. they have the same hash,
@@ -153,7 +152,7 @@ pub(crate) fn query_feed<'tcx, Cache>(
                 query_vtable.hash_result,
                 query_vtable.format_value,
             );
-            cache.complete(key, value, dep_node_index);
+            query_vtable.query_cache.complete(key, value, dep_node_index);
         }
     }
 }
diff --git a/compiler/rustc_middle/src/query/plumbing.rs b/compiler/rustc_middle/src/query/plumbing.rs
@@ -14,9 +14,7 @@ pub use sealed::IntoQueryParam;
 use crate::dep_graph;
 use crate::dep_graph::{DepKind, DepNodeIndex, SerializedDepNodeIndex};
 use crate::ich::StableHashingContext;
-use crate::queries::{
-    ExternProviders, PerQueryVTables, Providers, QueryArenas, QueryCaches, QueryEngine, QueryStates,
-};
+use crate::queries::{ExternProviders, PerQueryVTables, Providers, QueryArenas, QueryEngine};
 use crate::query::on_disk_cache::{CacheEncoder, EncodedDepNodeIndex, OnDiskCache};
 use crate::query::stack::{QueryStackDeferred, QueryStackFrame, QueryStackFrameExtra};
 use crate::query::{QueryCache, QueryInfo, QueryJob};
@@ -108,10 +106,8 @@ pub struct QueryVTable<'tcx, C: QueryCache> {
     pub dep_kind: DepKind,
     /// How this query deals with query cycle errors.
     pub cycle_error_handling: CycleErrorHandling,
-    // Offset of this query's state field in the QueryStates struct
-    pub query_state: usize,
-    // Offset of this query's cache field in the QueryCaches struct
-    pub query_cache: usize,
+    pub query_state: QueryState<'tcx, C::Key>,
+    pub query_cache: C,
     pub will_cache_on_disk_for_key_fn: Option<WillCacheOnDiskForKeyFn<'tcx, C::Key>>,
 
     /// Function pointer that calls `tcx.$query(key)` for this query and
@@ -155,9 +151,7 @@ pub struct QuerySystemFns {
 }
 
 pub struct QuerySystem<'tcx> {
-    pub states: QueryStates<'tcx>,
     pub arenas: WorkerLocal<QueryArenas<'tcx>>,
-    pub caches: QueryCaches<'tcx>,
     pub query_vtables: PerQueryVTables<'tcx>,
 
     /// This provides access to the incremental compilation on-disk cache for query results.
@@ -445,11 +439,6 @@ macro_rules! define_callbacks {
             )*
         }
 
-        #[derive(Default)]
-        pub struct QueryCaches<'tcx> {
-            $($(#[$attr])* pub $name: $name::Storage<'tcx>,)*
-        }
-
         impl<'tcx> $crate::query::TyCtxtEnsureOk<'tcx> {
             $($(#[$attr])*
             #[inline(always)]
@@ -461,7 +450,7 @@ macro_rules! define_callbacks {
                     [$($modifiers)*]
                     self.tcx,
                     self.tcx.query_system.fns.engine.$name,
-                    &self.tcx.query_system.caches.$name,
+                    &self.tcx.query_system.query_vtables.$name.query_cache,
                     $crate::query::IntoQueryParam::into_query_param(key),
                     false,
                 )
@@ -475,7 +464,7 @@ macro_rules! define_callbacks {
                 crate::query::inner::query_ensure(
                     self.tcx,
                     self.tcx.query_system.fns.engine.$name,
-                    &self.tcx.query_system.caches.$name,
+                    &self.tcx.query_system.query_vtables.$name.query_cache,
                     $crate::query::IntoQueryParam::into_query_param(key),
                     true,
                 );
@@ -502,7 +491,7 @@ macro_rules! define_callbacks {
                 erase::restore_val::<$V>(inner::query_get_at(
                     self.tcx,
                     self.tcx.query_system.fns.engine.$name,
-                    &self.tcx.query_system.caches.$name,
+                    &self.tcx.query_system.query_vtables.$name.query_cache,
                     self.span,
                     $crate::query::IntoQueryParam::into_query_param(key),
                 ))
@@ -518,13 +507,6 @@ macro_rules! define_callbacks {
             )*
         }
 
-        #[derive(Default)]
-        pub struct QueryStates<'tcx> {
-            $(
-                pub $name: $crate::query::QueryState<'tcx, $($K)*>,
-            )*
-        }
-
         pub struct Providers {
             $(
                 /// This is the provider for the query. Use `Find references` on this to
@@ -594,7 +576,6 @@ macro_rules! define_feedable {
                     tcx,
                     dep_kind,
                     &tcx.query_system.query_vtables.$name,
-                    &tcx.query_system.caches.$name,
                     key,
                     erased_value,
                 );
diff --git a/compiler/rustc_query_impl/src/dep_kind_vtables.rs b/compiler/rustc_query_impl/src/dep_kind_vtables.rs
@@ -110,18 +110,18 @@ mod non_query {
 
 /// Shared implementation of the [`DepKindVTable`] constructor for queries.
 /// Called from macro-generated code for each query.
-pub(crate) fn make_dep_kind_vtable_for_query<'tcx, Q, Cache, const FLAGS: QueryFlags>(
+pub(crate) fn make_dep_kind_vtable_for_query<'tcx, Q, C, const FLAGS: QueryFlags>(
     is_eval_always: bool,
 ) -> DepKindVTable<'tcx>
 where
-    Q: QueryDispatcherUnerased<'tcx, Cache, FLAGS>,
-    Cache: QueryCache + 'tcx,
+    Q: QueryDispatcherUnerased<'tcx, C, FLAGS>,
+    C: QueryCache + 'tcx,
 {
     let is_anon = FLAGS.is_anon;
     let key_fingerprint_style = if is_anon {
         KeyFingerprintStyle::Opaque
     } else {
-        <Cache::Key as DepNodeKey<'tcx>>::key_fingerprint_style()
+        <C::Key as DepNodeKey<'tcx>>::key_fingerprint_style()
     };
 
     if is_anon || !key_fingerprint_style.reconstructible() {
diff --git a/compiler/rustc_query_impl/src/execution.rs b/compiler/rustc_query_impl/src/execution.rs
@@ -244,12 +244,12 @@ fn wait_for_query<'tcx, C: QueryCache, const FLAGS: QueryFlags>(
 
     match result {
         Ok(()) => {
-            let Some((v, index)) = query.query_cache(tcx).lookup(&key) else {
+            let Some((v, index)) = query.query_cache().lookup(&key) else {
                 outline(|| {
                     // We didn't find the query result in the query cache. Check if it was
                     // poisoned due to a panic instead.
                     let key_hash = sharded::make_hash(&key);
-                    let shard = query.query_state(tcx).active.lock_shard_by_hash(key_hash);
+                    let shard = query.query_state().active.lock_shard_by_hash(key_hash);
                     match shard.find(key_hash, equivalent_key(&key)) {
                         // The query we waited on panicked. Continue unwinding here.
                         Some((_, ActiveKeyStatus::Poisoned)) => FatalError.raise(),
@@ -278,9 +278,8 @@ fn try_execute_query<'tcx, C: QueryCache, const FLAGS: QueryFlags, const INCR: b
     key: C::Key,
     dep_node: Option<DepNode>,
 ) -> (C::Value, Option<DepNodeIndex>) {
-    let state = query.query_state(tcx);
     let key_hash = sharded::make_hash(&key);
-    let mut state_lock = state.active.lock_shard_by_hash(key_hash);
+    let mut state_lock = query.query_state().active.lock_shard_by_hash(key_hash);
 
     // For the parallel compiler we need to check both the query cache and query state structures
     // while holding the state lock to ensure that 1) the query has not yet completed and 2) the
@@ -289,7 +288,7 @@ fn try_execute_query<'tcx, C: QueryCache, const FLAGS: QueryFlags, const INCR: b
     // executing, but another thread may have already completed the query and stores it result
     // in the query cache.
     if tcx.sess.threads() > 1 {
-        if let Some((value, index)) = query.query_cache(tcx).lookup(&key) {
+        if let Some((value, index)) = query.query_cache().lookup(&key) {
             tcx.prof.query_cache_hit(index.into());
             return (value, Some(index));
         }
@@ -308,7 +307,7 @@ fn try_execute_query<'tcx, C: QueryCache, const FLAGS: QueryFlags, const INCR: b
             // Drop the lock before we start executing the query
             drop(state_lock);
 
-            execute_job::<C, FLAGS, INCR>(query, tcx, state, key, key_hash, id, dep_node)
+            execute_job::<C, FLAGS, INCR>(query, tcx, key, key_hash, id, dep_node)
         }
         Entry::Occupied(mut entry) => {
             match &mut entry.get_mut().1 {
@@ -340,15 +339,14 @@ fn try_execute_query<'tcx, C: QueryCache, const FLAGS: QueryFlags, const INCR: b
 fn execute_job<'tcx, C: QueryCache, const FLAGS: QueryFlags, const INCR: bool>(
     query: SemiDynamicQueryDispatcher<'tcx, C, FLAGS>,
     tcx: TyCtxt<'tcx>,
-    state: &'tcx QueryState<'tcx, C::Key>,
     key: C::Key,
     key_hash: u64,
     id: QueryJobId,
     dep_node: Option<DepNode>,
 ) -> (C::Value, Option<DepNodeIndex>) {
     // Set up a guard object that will automatically poison the query if a
     // panic occurs while executing the query (or any intermediate plumbing).
-    let job_guard = ActiveJobGuard { state, key, key_hash };
+    let job_guard = ActiveJobGuard { state: query.query_state(), key, key_hash };
 
     debug_assert_eq!(tcx.dep_graph.is_fully_enabled(), INCR);
 
@@ -359,7 +357,7 @@ fn execute_job<'tcx, C: QueryCache, const FLAGS: QueryFlags, const INCR: bool>(
         execute_job_non_incr(query, tcx, key, id)
     };
 
-    let cache = query.query_cache(tcx);
+    let cache = query.query_cache();
     if query.feedable() {
         // We should not compute queries that also got a value via feeding.
         // This can't happen, as query feeding adds the very dependencies to the fed query
@@ -680,7 +678,7 @@ pub(crate) fn force_query<'tcx, C: QueryCache, const FLAGS: QueryFlags>(
 ) {
     // We may be concurrently trying both execute and force a query.
     // Ensure that only one of them runs the query.
-    if let Some((_, index)) = query.query_cache(tcx).lookup(&key) {
+    if let Some((_, index)) = query.query_cache().lookup(&key) {
         tcx.prof.query_cache_hit(index.into());
         return;
     }
diff --git a/compiler/rustc_query_impl/src/lib.rs b/compiler/rustc_query_impl/src/lib.rs
@@ -14,9 +14,7 @@ use std::marker::ConstParamTy;
 
 use rustc_data_structures::sync::AtomicU64;
 use rustc_middle::dep_graph::{self, DepKind, DepNode, DepNodeIndex, SerializedDepNodeIndex};
-use rustc_middle::queries::{
-    self, ExternProviders, Providers, QueryCaches, QueryEngine, QueryStates,
-};
+use rustc_middle::queries::{self, ExternProviders, Providers, QueryEngine};
 use rustc_middle::query::on_disk_cache::{CacheEncoder, EncodedDepNodeIndex, OnDiskCache};
 use rustc_middle::query::plumbing::{
     HashResult, QueryState, QuerySystem, QuerySystemFns, QueryVTable,
@@ -104,26 +102,14 @@ impl<'tcx, C: QueryCache, const FLAGS: QueryFlags> SemiDynamicQueryDispatcher<'t
 
     // Don't use this method to access query results, instead use the methods on TyCtxt.
     #[inline(always)]
-    fn query_state(self, tcx: TyCtxt<'tcx>) -> &'tcx QueryState<'tcx, C::Key> {
-        // Safety:
-        // This is just manually doing the subfield referencing through pointer math.
-        unsafe {
-            &*(&tcx.query_system.states as *const QueryStates<'tcx>)
-                .byte_add(self.vtable.query_state)
-                .cast::<QueryState<'tcx, C::Key>>()
-        }
+    fn query_state(self) -> &'tcx QueryState<'tcx, C::Key> {
+        &self.vtable.query_state
     }
 
     // Don't use this method to access query results, instead use the methods on TyCtxt.
     #[inline(always)]
-    fn query_cache(self, tcx: TyCtxt<'tcx>) -> &'tcx C {
-        // Safety:
-        // This is just manually doing the subfield referencing through pointer math.
-        unsafe {
-            &*(&tcx.query_system.caches as *const QueryCaches<'tcx>)
-                .byte_add(self.vtable.query_cache)
-                .cast::<C>()
-        }
+    fn query_cache(self) -> &'tcx C {
+        &self.vtable.query_cache
     }
 
     /// Calls `tcx.$query(key)` for this query, and discards the returned value.
@@ -245,9 +231,7 @@ pub fn query_system<'tcx>(
     incremental: bool,
 ) -> QuerySystem<'tcx> {
     QuerySystem {
-        states: Default::default(),
         arenas: Default::default(),
-        caches: Default::default(),
         query_vtables: make_query_vtables(),
         on_disk_cache,
         fns: QuerySystemFns {
diff --git a/compiler/rustc_query_impl/src/plumbing.rs b/compiler/rustc_query_impl/src/plumbing.rs
