Improve how QueryCaches and QueryStates are stored

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.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);
         }
     }
 }

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,
                 );

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;
     }

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 {

compiler/rustc_query_impl/src/plumbing.rs

@@ -354,9 +354,8 @@ pub(crate) fn encode_query_results<'a, 'tcx, Q, C: QueryCache, const FLAGS: Quer
 {
     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.query_state()));
+    query.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());
 
@@ -376,7 +375,7 @@ pub(crate) fn query_key_hash_verify<'tcx, C: QueryCache, const FLAGS: QueryFlags
 ) {
     let _timer = tcx.prof.generic_activity_with_arg("query_key_hash_verify_for", query.name());
 
-    let cache = query.query_cache(tcx);
+    let cache = query.query_cache();
     let mut map = UnordMap::with_capacity(cache.len());
     cache.iter(&mut |key, _, _| {
         let node = DepNode::construct(tcx, query.dep_kind(), key);
@@ -577,8 +576,8 @@ macro_rules! define_queries {
                     eval_always: is_eval_always!([$($modifiers)*]),
                     dep_kind: dep_graph::dep_kinds::$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),
+                    query_state: Default::default(),
+                    query_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)
                     } {
@@ -673,7 +672,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.query_state,
                     tcx,
                     make_frame,
                     require_complete,
@@ -699,7 +699,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.query_cache,
                     string_cache,
                 )
             }