Factor query arena allocation out from query caches

rust-lang · Feb 9, 2023 · 65f0ce0 · 65f0ce0
1 parent 97befd4
commit 65f0ce0
Show file tree

Hide file tree

Showing 10 changed files with 160 additions and 255 deletions.
diff --git a/compiler/rustc_data_structures/src/sync/worker_local.rs b/compiler/rustc_data_structures/src/sync/worker_local.rs
@@ -164,6 +164,12 @@ impl<T> WorkerLocal<Vec<T>> {
     }
 }
 
+impl<T: Default> Default for WorkerLocal<T> {
+    fn default() -> Self {
+        WorkerLocal::new(|_| T::default())
+    }
+}
+
 impl<T> Deref for WorkerLocal<T> {
     type Target = T;
 

diff --git a/compiler/rustc_interface/src/passes.rs b/compiler/rustc_interface/src/passes.rs
@@ -781,9 +781,7 @@ pub fn create_global_ctxt<'tcx>(
         callback(sess, &mut local_providers, &mut extern_providers);
     }
 
-    let queries = queries.get_or_init(|| {
-        TcxQueries::new(local_providers, extern_providers, query_result_on_disk_cache)
-    });
+    let queries = queries.get_or_init(|| TcxQueries::new(query_result_on_disk_cache));
 
     let gcx = sess.time("setup_global_ctxt", || {
         global_ctxt.get_or_init(move || {
@@ -795,6 +793,8 @@ pub fn create_global_ctxt<'tcx>(
                 untracked,
                 dep_graph,
                 queries.on_disk_cache.as_ref().map(OnDiskCache::as_dyn),
+                local_providers,
+                extern_providers,
                 queries.as_dyn(),
                 rustc_query_impl::query_callbacks(arena),
             )

diff --git a/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs b/compiler/rustc_metadata/src/rmeta/decoder/cstore_impl.rs
@@ -114,7 +114,7 @@ macro_rules! provide_one {
         fn $name<'tcx>(
             $tcx: TyCtxt<'tcx>,
             def_id_arg: ty::query::query_keys::$name<'tcx>,
-        ) -> ty::query::query_values::$name<'tcx> {
+        ) -> ty::query::query_provided::$name<'tcx> {
             let _prof_timer =
                 $tcx.prof.generic_activity(concat!("metadata_decode_entry_", stringify!($name)));
 

diff --git a/compiler/rustc_middle/src/ty/context.rs b/compiler/rustc_middle/src/ty/context.rs
@@ -18,6 +18,8 @@ use crate::mir::{
 use crate::thir::Thir;
 use crate::traits;
 use crate::traits::solve::{ExternalConstraints, ExternalConstraintsData};
+use crate::ty::query::ExternProviders;
+use crate::ty::query::Providers;
 use crate::ty::query::{self, TyCtxtAt};
 use crate::ty::{
     self, AdtDef, AdtDefData, AdtKind, Binder, Const, ConstData, DefIdTree, FloatTy, FloatVar,
@@ -445,7 +447,7 @@ pub struct GlobalCtxt<'tcx> {
     pub on_disk_cache: Option<&'tcx dyn OnDiskCache<'tcx>>,
 
     pub queries: &'tcx dyn query::QueryEngine<'tcx>,
-    pub query_caches: query::QueryCaches<'tcx>,
+    pub query_system: query::QuerySystem<'tcx>,
     pub(crate) query_kinds: &'tcx [DepKindStruct<'tcx>],
 
     // Internal caches for metadata decoding. No need to track deps on this.
@@ -593,6 +595,8 @@ impl<'tcx> TyCtxt<'tcx> {
         untracked: Untracked,
         dep_graph: DepGraph,
         on_disk_cache: Option<&'tcx dyn OnDiskCache<'tcx>>,
+        local_providers: Providers,
+        extern_providers: ExternProviders,
         queries: &'tcx dyn query::QueryEngine<'tcx>,
         query_kinds: &'tcx [DepKindStruct<'tcx>],
     ) -> GlobalCtxt<'tcx> {
@@ -618,7 +622,7 @@ impl<'tcx> TyCtxt<'tcx> {
             untracked,
             on_disk_cache,
             queries,
-            query_caches: query::QueryCaches::default(),
+            query_system: query::QuerySystem::new(local_providers, extern_providers),
             query_kinds,
             ty_rcache: Default::default(),
             pred_rcache: Default::default(),

diff --git a/compiler/rustc_middle/src/ty/query.rs b/compiler/rustc_middle/src/ty/query.rs
@@ -1,3 +1,5 @@
+#![allow(unused_parens)]
+
 use crate::dep_graph;
 use crate::infer::canonical::{self, Canonical};
 use crate::lint::LintExpectation;
@@ -35,13 +37,15 @@ use crate::ty::subst::{GenericArg, SubstsRef};
 use crate::ty::util::AlwaysRequiresDrop;
 use crate::ty::GeneratorDiagnosticData;
 use crate::ty::{self, CrateInherentImpls, ParamEnvAnd, Ty, TyCtxt, UnusedGenericParams};
+use rustc_arena::TypedArena;
 use rustc_ast as ast;
 use rustc_ast::expand::allocator::AllocatorKind;
 use rustc_attr as attr;
 use rustc_data_structures::fx::{FxHashMap, FxHashSet, FxIndexMap, FxIndexSet};
 use rustc_data_structures::steal::Steal;
 use rustc_data_structures::svh::Svh;
 use rustc_data_structures::sync::Lrc;
+use rustc_data_structures::sync::WorkerLocal;
 use rustc_data_structures::unord::UnordSet;
 use rustc_errors::ErrorGuaranteed;
 use rustc_hir as hir;
@@ -67,6 +71,24 @@ use std::sync::Arc;
 pub(crate) use rustc_query_system::query::QueryJobId;
 use rustc_query_system::query::*;
 
+pub struct QuerySystem<'tcx> {
+    pub local_providers: Box<Providers>,
+    pub extern_providers: Box<ExternProviders>,
+    pub arenas: QueryArenas<'tcx>,
+    pub caches: QueryCaches<'tcx>,
+}
+
+impl<'tcx> QuerySystem<'tcx> {
+    pub fn new(local_providers: Providers, extern_providers: ExternProviders) -> Self {
+        QuerySystem {
+            local_providers: Box::new(local_providers),
+            extern_providers: Box::new(extern_providers),
+            arenas: Default::default(),
+            caches: Default::default(),
+        }
+    }
+}
+
 #[derive(Copy, Clone)]
 pub struct TyCtxtAt<'tcx> {
     pub tcx: TyCtxt<'tcx>,
@@ -115,14 +137,14 @@ fn query_get_at<'tcx, Cache, K>(
         Span,
         Cache::Key,
         QueryMode,
-    ) -> Option<Cache::Stored>,
+    ) -> Option<Cache::Value>,
     query_cache: &Cache,
     span: Span,
     key: K,
-) -> Cache::Stored
+) -> Cache::Value
 where
     K: IntoQueryParam<Cache::Key>,
-    Cache::Stored: Copy,
+    Cache::Value: Copy,
     Cache: QueryCache,
 {
     let key = key.into_query_param();
@@ -141,12 +163,12 @@ fn query_ensure<'tcx, Cache, K>(
         Span,
         Cache::Key,
         QueryMode,
-    ) -> Option<Cache::Stored>,
+    ) -> Option<Cache::Value>,
     query_cache: &Cache,
     key: K,
 ) where
     K: IntoQueryParam<Cache::Key>,
-    Cache::Stored: Copy,
+    Cache::Value: Copy,
     Cache: QueryCache,
 {
     let key = key.into_query_param();
@@ -162,10 +184,10 @@ macro_rules! query_helper_param_ty {
 }
 
 macro_rules! query_if_arena {
-    ([] $arena:ty, $no_arena:ty) => {
+    ([] $arena:tt $no_arena:tt) => {
         $no_arena
     };
-    ([(arena_cache) $($rest:tt)*] $arena:ty, $no_arena:ty) => {
+    ([(arena_cache) $($rest:tt)*] $arena:tt $no_arena:tt) => {
         $arena
     };
     ([$other:tt $($modifiers:tt)*]$($args:tt)*) => {
@@ -181,7 +203,7 @@ macro_rules! separate_provide_extern_decl {
         for<'tcx> fn(
             TyCtxt<'tcx>,
             query_keys::$name<'tcx>,
-        ) -> query_values::$name<'tcx>
+        ) -> query_provided::$name<'tcx>
     };
     ([$other:tt $($modifiers:tt)*][$($args:tt)*]) => {
         separate_provide_extern_decl!([$($modifiers)*][$($args)*])
@@ -233,30 +255,52 @@ macro_rules! define_callbacks {
 
             $(pub type $name<'tcx> = $($K)*;)*
         }
-        #[allow(nonstandard_style, unused_lifetimes, unused_parens)]
+        #[allow(nonstandard_style, unused_lifetimes)]
         pub mod query_values {
             use super::*;
 
-            $(pub type $name<'tcx> = query_if_arena!([$($modifiers)*] <$V as Deref>::Target, $V);)*
+            $(pub type $name<'tcx> = $V;)*
         }
-        #[allow(nonstandard_style, unused_lifetimes, unused_parens)]
-        pub mod query_storage {
+        #[allow(nonstandard_style, unused_lifetimes)]
+        pub mod query_provided {
             use super::*;
 
             $(
-                pub type $name<'tcx> = query_if_arena!([$($modifiers)*]
-                    <<$($K)* as Key>::CacheSelector
-                        as CacheSelector<'tcx, <$V as Deref>::Target>>::ArenaCache,
-                    <<$($K)* as Key>::CacheSelector as CacheSelector<'tcx, $V>>::Cache
-                );
+                pub type $name<'tcx> = query_if_arena!([$($modifiers)*] (<$V as Deref>::Target) ($V));
             )*
         }
+        #[allow(nonstandard_style, unused_lifetimes)]
+        pub mod query_provided_to_value {
+            use super::*;
 
+            $(
+                #[inline]
+                pub fn $name<'tcx>(
+                    _tcx: TyCtxt<'tcx>,
+                    value: query_provided::$name<'tcx>,
+                ) -> query_values::$name<'tcx> {
+                    query_if_arena!([$($modifiers)*]
+                        (&*_tcx.query_system.arenas.$name.alloc(value))
+                        (value)
+                    )
+                }
+            )*
+        }
         #[allow(nonstandard_style, unused_lifetimes)]
-        pub mod query_stored {
+        pub mod query_storage {
             use super::*;
 
-            $(pub type $name<'tcx> = $V;)*
+            $(
+                pub type $name<'tcx> = <<$($K)* as Key>::CacheSelector as CacheSelector<'tcx, $V>>::Cache;
+            )*
+        }
+
+        #[derive(Default)]
+        pub struct QueryArenas<'tcx> {
+            $($(#[$attr])* pub $name: query_if_arena!([$($modifiers)*]
+                (WorkerLocal<TypedArena<<$V as Deref>::Target>>)
+                ()
+            ),)*
         }
 
         #[derive(Default)]
@@ -271,7 +315,7 @@ macro_rules! define_callbacks {
                 query_ensure(
                     self.tcx,
                     QueryEngine::$name,
-                    &self.tcx.query_caches.$name,
+                    &self.tcx.query_system.caches.$name,
                     opt_remap_env_constness!([$($modifiers)*][key]),
                 );
             })*
@@ -286,7 +330,7 @@ macro_rules! define_callbacks {
                 query_get_at(
                     self,
                     QueryEngine::$name,
-                    &self.query_caches.$name,
+                    &self.query_system.caches.$name,
                     DUMMY_SP,
                     opt_remap_env_constness!([$($modifiers)*][key]),
                 )
@@ -301,7 +345,7 @@ macro_rules! define_callbacks {
                 query_get_at(
                     self.tcx,
                     QueryEngine::$name,
-                    &self.tcx.query_caches.$name,
+                    &self.tcx.query_system.caches.$name,
                     self.span,
                     opt_remap_env_constness!([$($modifiers)*][key]),
                 )
@@ -312,7 +356,7 @@ macro_rules! define_callbacks {
             $(pub $name: for<'tcx> fn(
                 TyCtxt<'tcx>,
                 query_keys::$name<'tcx>,
-            ) -> query_values::$name<'tcx>,)*
+            ) -> query_provided::$name<'tcx>,)*
         }
 
         pub struct ExternProviders {
@@ -389,12 +433,13 @@ macro_rules! define_feedable {
         $(impl<'tcx, K: IntoQueryParam<$($K)*> + Copy> TyCtxtFeed<'tcx, K> {
             $(#[$attr])*
             #[inline(always)]
-            pub fn $name(self, value: query_values::$name<'tcx>) -> $V {
+            pub fn $name(self, value: query_provided::$name<'tcx>) -> $V {
                 let key = self.key().into_query_param();
                 let key = opt_remap_env_constness!([$($modifiers)*][key]);
 
                 let tcx = self.tcx;
-                let cache = &tcx.query_caches.$name;
+                let value = query_provided_to_value::$name(tcx, value);
+                let cache = &tcx.query_system.caches.$name;
 
                 match try_get_cached(tcx, cache, &key) {
                     Some(old) => {
@@ -412,7 +457,8 @@ macro_rules! define_feedable {
                             &value,
                             hash_result!([$($modifiers)*]),
                         );
-                        cache.complete(key, value, dep_node_index)
+                        cache.complete(key, value, dep_node_index);
+                        value
                     }
                 }
             }

diff --git a/compiler/rustc_query_impl/src/lib.rs b/compiler/rustc_query_impl/src/lib.rs
@@ -21,8 +21,10 @@ use rustc_data_structures::sync::AtomicU64;
 use rustc_middle::arena::Arena;
 use rustc_middle::dep_graph::{self, DepKindStruct};
 use rustc_middle::query::Key;
-use rustc_middle::ty::query::{query_keys, query_storage, query_stored, query_values};
-use rustc_middle::ty::query::{ExternProviders, Providers, QueryEngine};
+use rustc_middle::ty::query::QueryEngine;
+use rustc_middle::ty::query::{
+    query_keys, query_provided, query_provided_to_value, query_storage, query_values,
+};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::Span;