Improve how QueryCaches and QueryStates are stored.

nnethercote · nnethercote · commit 54314d71c1c0 · 2026-02-24T13:01:21.000+11:00
`QuerySystem` has these fields:
```
pub states: QueryStates&lt;'tcx&gt;,
pub caches: QueryCaches&lt;'tcx&gt;,
pub query_vtables: PerQueryVTables&lt;'tcx&gt;,
```
Each one has an entry for each query.

Some methods that take a query-specific `QueryVTable` need to access the
corresponding query-specific states and/or caches. So `QueryVTable`
stores the *byte offset* to the relevant fields within `states` and
`caches`, and uses that to (with `unsafe`) access the fields. This is
bizarre, and I hope it made sense in the past when the code was
structured differently.

This commit moves `QueryState` and `QueryCache` inside `QueryVTable`. As
a result, the query-specific states and/or caches are directly
accessible, and no unsafe offset computations are required. Much simpler
and more normal. Also, `QueryCaches` and `QueryStates` are no longer
needed, which makes `define_callbacks!` a bit simpler.

The commit also rename the fields and their getters to `states` and
`caches`.
diff --git a/compiler/rustc_middle/src/query/inner.rs b/compiler/rustc_middle/src/query/inner.rs
@@ -100,7 +100,6 @@ pub(crate) fn query_feed<'tcx, Cache>(
     tcx: TyCtxt<'tcx>,
     dep_kind: DepKind,
     query_vtable: &QueryVTable<'tcx, Cache>,
-    cache: &Cache,
     key: Cache::Key,
     value: Cache::Value,
 ) where
@@ -110,7 +109,7 @@ pub(crate) fn query_feed<'tcx, Cache>(
     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.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.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, DepNode, 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};
@@ -101,7 +99,7 @@ pub enum QueryMode {
     Ensure { check_cache: bool },
 }
 
-/// Stores function pointers and other metadata for a particular query.
+/// Stores data and metadata (e.g. function pointers) for a particular query.
 pub struct QueryVTable<'tcx, C: QueryCache> {
     pub name: &'static str,
 
@@ -117,10 +115,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 state: QueryState<'tcx, C::Key>,
+    pub 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
@@ -169,30 +165,6 @@ impl<'tcx, C: QueryCache> QueryVTable<'tcx, C> {
         self.will_cache_on_disk_for_key_fn.map_or(false, |f| f(tcx, key))
     }
 
-    // Don't use this method to access query results, instead use the methods on TyCtxt.
-    #[inline(always)]
-    pub 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.query_state)
-                .cast::<QueryState<'tcx, C::Key>>()
-        }
-    }
-
-    // Don't use this method to access query results, instead use the methods on TyCtxt.
-    #[inline(always)]
-    pub 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.query_cache)
-                .cast::<C>()
-        }
-    }
-
     #[inline(always)]
     pub fn try_load_from_disk(
         &self,
@@ -243,9 +215,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.
@@ -498,13 +468,6 @@ macro_rules! define_callbacks {
             )*
         }
 
-        #[derive(Default)]
-        pub struct QueryCaches<'tcx> {
-            $(
-                pub $name: $name::Storage<'tcx>,
-            )*
-        }
-
         impl<'tcx> $crate::query::TyCtxtEnsureOk<'tcx> {
             $(
                 $(#[$attr])*
@@ -524,7 +487,7 @@ macro_rules! define_callbacks {
                     )(
                         self.tcx,
                         self.tcx.query_system.fns.engine.$name,
-                        &self.tcx.query_system.caches.$name,
+                        &self.tcx.query_system.query_vtables.$name.cache,
                         $crate::query::IntoQueryParam::into_query_param(key),
                         false,
                     )
@@ -540,7 +503,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.cache,
                         $crate::query::IntoQueryParam::into_query_param(key),
                         true,
                     );
@@ -569,7 +532,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.cache,
                         self.span,
                         $crate::query::IntoQueryParam::into_query_param(key),
                     ))
@@ -586,13 +549,6 @@ macro_rules! define_callbacks {
             )*
         }
 
-        #[derive(Default)]
-        pub struct QueryStates<'tcx> {
-            $(
-                pub $name: $crate::query::QueryState<'tcx, $name::Key<'tcx>>,
-            )*
-        }
-
         pub struct Providers {
             $(
                 /// This is the provider for the query. Use `Find references` on this to
@@ -689,7 +645,6 @@ macro_rules! define_feedable {
                         tcx,
                         dep_graph::DepKind::$name,
                         &tcx.query_system.query_vtables.$name,
-                        &tcx.query_system.caches.$name,
                         key,
                         erased_value,
                     );
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>(
 
     match result {
         Ok(()) => {
-            let Some((v, index)) = query.query_cache(tcx).lookup(&key) else {
+            let Some((v, index)) = 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.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 INCR: bool>(
     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.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 INCR: bool>(
     // 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.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 INCR: bool>(
             // Drop the lock before we start executing the query
             drop(state_lock);
 
-            execute_job::<C, INCR>(query, tcx, state, key, key_hash, id, dep_node)
+            execute_job::<C, 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 INCR: bool>(
 fn execute_job<'tcx, C: QueryCache, const INCR: bool>(
     query: &'tcx QueryVTable<'tcx, C>,
     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.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 INCR: bool>(
         execute_job_non_incr(query, tcx, key, id)
     };
 
-    let cache = query.query_cache(tcx);
+    let cache = &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
@@ -681,7 +679,7 @@ pub(crate) fn force_query<'tcx, C: QueryCache>(
 ) {
     // 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.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
@@ -10,9 +10,7 @@
 
 use rustc_data_structures::sync::AtomicU64;
 use rustc_middle::dep_graph;
-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::{QuerySystem, QuerySystemFns, QueryVTable};
 use rustc_middle::query::{AsLocalKey, QueryCache, QueryMode};
@@ -62,9 +60,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
@@ -314,6 +314,7 @@ pub(crate) fn create_deferred_query_stack_frame<'tcx, Cache>(
 where
     Cache: QueryCache,
     Cache::Key: Key + DynSend + DynSync,
+    QueryVTable<'tcx, Cache>: DynSync,
 {
     let kind = vtable.dep_kind;
 
@@ -335,9 +336,8 @@ pub(crate) fn encode_query_results<'a, 'tcx, Q, C: QueryCache>(
 {
     let _timer = tcx.prof.generic_activity_with_arg("encode_query_results_for", query.name);
 
-    assert!(all_inactive(query.query_state(tcx)));
-    let cache = query.query_cache(tcx);
-    cache.iter(&mut |key, value, dep_node| {
+    assert!(all_inactive(&query.state));
+    query.cache.iter(&mut |key, value, dep_node| {
         if query.will_cache_on_disk_for_key(tcx, key) {
             let dep_node = SerializedDepNodeIndex::new(dep_node.index());
 
@@ -357,7 +357,7 @@ pub(crate) fn query_key_hash_verify<'tcx, C: QueryCache>(
 ) {
     let _timer = tcx.prof.generic_activity_with_arg("query_key_hash_verify_for", query.name);
 
-    let cache = query.query_cache(tcx);
+    let cache = &query.cache;
     let mut map = UnordMap::with_capacity(cache.len());
     cache.iter(&mut |key, _, _| {
         let node = DepNode::construct(tcx, query.dep_kind, key);
@@ -560,8 +560,8 @@ macro_rules! define_queries {
                     feedable: is_feedable!([$($modifiers)*]),
                     dep_kind: dep_graph::DepKind::$name,
                     cycle_error_handling: cycle_error_handling!([$($modifiers)*]),
-                    query_state: std::mem::offset_of!(QueryStates<'tcx>, $name),
-                    query_cache: std::mem::offset_of!(QueryCaches<'tcx>, $name),
+                    state: Default::default(),
+                    cache: Default::default(),
                     will_cache_on_disk_for_key_fn: if_cache_on_disk!([$($modifiers)*] {
                         Some(::rustc_middle::queries::_cache_on_disk_if_fns::$name)
                     } {
@@ -646,7 +646,8 @@ macro_rules! define_queries {
                 };
 
                 // Call `gather_active_jobs_inner` to do the actual work.
-                let res = crate::execution::gather_active_jobs_inner(&tcx.query_system.states.$name,
+                let res = crate::execution::gather_active_jobs_inner(
+                    &tcx.query_system.query_vtables.$name.state,
                     tcx,
                     make_frame,
                     require_complete,
@@ -672,7 +673,7 @@ macro_rules! define_queries {
                 $crate::profiling_support::alloc_self_profile_query_strings_for_query_cache(
                     tcx,
                     stringify!($name),
-                    &tcx.query_system.caches.$name,
+                    &tcx.query_system.query_vtables.$name.cache,
                     string_cache,
                 )
             }
