Auto merge of #152835 - nnethercote:fix-query-state-query-cache, r=Zalathar

Improve how `QueryCache`/`QueryState`/`QueryEngine` are stored

*[View all comments](https://triagebot.infra.rust-lang.org/gh-comments/rust-lang/rust/pull/152835)*

`QueryVTable` has two fields `query_state` and `query_cache` that are *byte offsets* of fields in other structs. They are used in `unsafe` combination with `byte_add` and `cast` to access said fields. I don't like this at all and this PR replaces them with something sensible.

r? @Zalathar

Author: bors

compiler/rustc_middle/src/query/inner.rs

@@ -33,15 +33,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,
@@ -52,14 +52,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,
     ensure_mode: EnsureMode,
 ) where
-    Cache: QueryCache,
+    C: QueryCache,
 {
     if try_get_cached(tcx, query_cache, &key).is_none() {
         execute_query(tcx, DUMMY_SP, key, QueryMode::Ensure { ensure_mode });
@@ -69,17 +69,17 @@ 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,
     // This arg is needed to match the signature of `query_ensure`,
     // but should always be `EnsureMode::Ok`.
     ensure_mode: EnsureMode,
 ) -> Result<(), ErrorGuaranteed>
 where
-    Cache: QueryCache<Value = Erased<Result<T, ErrorGuaranteed>>>,
+    C: QueryCache<Value = Erased<Result<T, ErrorGuaranteed>>>,
     Result<T, ErrorGuaranteed>: Erasable,
 {
     assert_matches!(ensure_mode, EnsureMode::Ok);
@@ -101,21 +101,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.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,
@@ -158,7 +157,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);
         }
     }
 }

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, Providers, QueryArenas, QueryVTables};
 use crate::query::on_disk_cache::{CacheEncoder, EncodedDepNodeIndex, OnDiskCache};
 use crate::query::stack::{QueryStackDeferred, QueryStackFrame, QueryStackFrameExtra};
 use crate::query::{QueryCache, QueryInfo, QueryJob};
@@ -113,7 +111,7 @@ pub enum EnsureMode {
     Done,
 }
 
-/// 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,
 
@@ -129,10 +127,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
@@ -161,6 +157,8 @@ pub struct QueryVTable<'tcx, C: QueryCache> {
     ///
     /// Used when reporting query cycle errors and similar problems.
     pub description_fn: fn(TyCtxt<'tcx>, C::Key) -> String,
+
+    pub execute_query_fn: fn(TyCtxt<'tcx>, Span, C::Key, QueryMode) -> Option<C::Value>,
 }
 
 impl<'tcx, C: QueryCache> fmt::Debug for QueryVTable<'tcx, C> {
@@ -181,30 +179,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,7 +217,6 @@ impl<'tcx, C: QueryCache> QueryVTable<'tcx, C> {
 }
 
 pub struct QuerySystemFns {
-    pub engine: QueryEngine,
     pub local_providers: Providers,
     pub extern_providers: ExternProviders,
     pub encode_query_results: for<'tcx> fn(
@@ -255,10 +228,8 @@ 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>,
+    pub query_vtables: QueryVTables<'tcx>,
 
     /// This provides access to the incremental compilation on-disk cache for query results.
     /// Do not access this directly. It is only meant to be used by
@@ -521,13 +492,6 @@ macro_rules! define_callbacks {
             )*
         }
 
-        #[derive(Default)]
-        pub struct QueryCaches<'tcx> {
-            $(
-                pub $name: $name::Storage<'tcx>,
-            )*
-        }
-
         impl<'tcx> $crate::query::TyCtxtEnsureOk<'tcx> {
             $(
                 $(#[$attr])*
@@ -546,8 +510,8 @@ 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.execute_query_fn,
+                        &self.tcx.query_system.query_vtables.$name.cache,
                         $crate::query::IntoQueryParam::into_query_param(key),
                         $crate::query::EnsureMode::Ok,
                     )
@@ -562,8 +526,8 @@ macro_rules! define_callbacks {
                 pub fn $name(self, key: query_helper_param_ty!($($K)*)) {
                     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.execute_query_fn,
+                        &self.tcx.query_system.query_vtables.$name.cache,
                         $crate::query::IntoQueryParam::into_query_param(key),
                         $crate::query::EnsureMode::Done,
                     );
@@ -591,8 +555,8 @@ 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.execute_query_fn,
+                        &self.tcx.query_system.query_vtables.$name.cache,
                         self.span,
                         $crate::query::IntoQueryParam::into_query_param(key),
                     ))
@@ -615,7 +579,6 @@ macro_rules! define_callbacks {
                             self.tcx,
                             dep_graph::DepKind::$name,
                             &self.tcx.query_system.query_vtables.$name,
-                            &self.tcx.query_system.caches.$name,
                             key,
                             erased_value,
                         );
@@ -625,21 +588,12 @@ macro_rules! define_callbacks {
         )*
 
         /// Holds a `QueryVTable` for each query.
-        ///
-        /// ("Per" just makes this pluralized name more visually distinct.)
-        pub struct PerQueryVTables<'tcx> {
+        pub struct QueryVTables<'tcx> {
             $(
                 pub $name: ::rustc_middle::query::plumbing::QueryVTable<'tcx, $name::Storage<'tcx>>,
             )*
         }
 
-        #[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
@@ -699,17 +653,6 @@ macro_rules! define_callbacks {
         impl Clone for ExternProviders {
             fn clone(&self) -> Self { *self }
         }
-
-        pub struct QueryEngine {
-            $(
-                pub $name: for<'tcx> fn(
-                    TyCtxt<'tcx>,
-                    Span,
-                    $name::Key<'tcx>,
-                    $crate::query::QueryMode,
-                ) -> Option<$crate::query::erase::Erased<$V>>,
-            )*
-        }
     };
 }
 

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(),
@@ -280,9 +280,8 @@ fn try_execute_query<'tcx, C: QueryCache, const INCR: bool>(
     // query+key, which we should reuse instead of creating a new one.
     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
@@ -291,7 +290,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));
         }
@@ -310,7 +309,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 {
@@ -342,15 +341,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);
 
@@ -361,7 +359,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
@@ -706,7 +704,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;
     }

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};
 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};
@@ -58,13 +56,10 @@ pub fn query_system<'tcx>(
     incremental: bool,
 ) -> QuerySystem<'tcx> {
     QuerySystem {
-        states: Default::default(),
         arenas: Default::default(),
-        caches: Default::default(),
-        query_vtables: make_query_vtables(),
+        query_vtables: make_query_vtables(incremental),
         on_disk_cache,
         fns: QuerySystemFns {
-            engine: engine(incremental),
             local_providers,
             extern_providers,
             encode_query_results: encode_all_query_results,