Factor query arena allocation out from query caches

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

compiler/rustc_middle/src/ty/context.rs

@@ -479,7 +479,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.
@@ -664,7 +664,7 @@ impl<'tcx> TyCtxt<'tcx> {
             untracked,
             on_disk_cache,
             queries,
-            query_caches: query::QueryCaches::default(),
+            query_system: Default::default(),
             query_kinds,
             ty_rcache: Default::default(),
             pred_rcache: Default::default(),

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;
@@ -34,13 +36,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;
@@ -59,13 +63,20 @@ use rustc_span::symbol::Symbol;
 use rustc_span::{Span, DUMMY_SP};
 use rustc_target::abi;
 use rustc_target::spec::PanicStrategy;
+use std::mem;
 use std::ops::Deref;
 use std::path::PathBuf;
 use std::sync::Arc;
 
 pub(crate) use rustc_query_system::query::QueryJobId;
 use rustc_query_system::query::*;
 
+#[derive(Default)]
+pub struct QuerySystem<'tcx> {
+    pub arenas: QueryArenas<'tcx>,
+    pub caches: QueryCaches<'tcx>,
+}
+
 #[derive(Copy, Clone)]
 pub struct TyCtxtAt<'tcx> {
     pub tcx: TyCtxt<'tcx>,
@@ -112,10 +123,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)*) => {
@@ -131,7 +142,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)*])
@@ -183,30 +194,77 @@ 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 {
+
+        /// This module specifies the type returned from query providers and the type used for
+        /// decoding. For regular queries this is the declared returned type `V`, but
+        /// `arena_cache` will use `<V as Deref>::Target` instead.
+        #[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));
             )*
         }
 
+        /// This module has a function per query which takes a `query_provided` value and coverts
+        /// it to a regular `V` value by allocating it on an arena if the query has the
+        /// `arena_cache` modifier. This will happen when computing the query using a provider or
+        /// decoding a stored result.
         #[allow(nonstandard_style, unused_lifetimes)]
-        pub mod query_stored {
+        pub mod query_provided_to_value {
             use super::*;
 
-            $(pub type $name<'tcx> = $V;)*
+            $(
+                #[inline(always)]
+                pub fn $name<'tcx>(
+                    _tcx: TyCtxt<'tcx>,
+                    value: query_provided::$name<'tcx>,
+                ) -> query_values::$name<'tcx> {
+                    query_if_arena!([$($modifiers)*]
+                        {
+                            if mem::needs_drop::<query_provided::$name<'tcx>>() {
+                                &*_tcx.query_system.arenas.$name.alloc(value)
+                            } else {
+                                &*_tcx.arena.dropless.alloc(value)
+                            }
+                        }
+                        (value)
+                    )
+                }
+            )*
+        }
+        #[allow(nonstandard_style, unused_lifetimes)]
+        pub mod query_storage {
+            use super::*;
+
+            $(
+                pub type $name<'tcx> = <<$($K)* as Key>::CacheSelector as CacheSelector<'tcx, $V>>::Cache;
+            )*
+        }
+
+        pub struct QueryArenas<'tcx> {
+            $($(#[$attr])* pub $name: query_if_arena!([$($modifiers)*]
+                (WorkerLocal<TypedArena<<$V as Deref>::Target>>)
+                ()
+            ),)*
+        }
+
+        impl Default for QueryArenas<'_> {
+            fn default() -> Self {
+                Self {
+                    $($name: query_if_arena!([$($modifiers)*]
+                        (WorkerLocal::new(|_| Default::default()))
+                        ()
+                    ),)*
+                }
+            }
         }
 
         #[derive(Default)]
@@ -221,7 +279,7 @@ macro_rules! define_callbacks {
                 let key = key.into_query_param();
                 opt_remap_env_constness!([$($modifiers)*][key]);
 
-                match try_get_cached(self.tcx, &self.tcx.query_caches.$name, &key) {
+                match try_get_cached(self.tcx, &self.tcx.query_system.caches.$name, &key) {
                     Some(_) => return,
                     None => self.tcx.queries.$name(self.tcx, DUMMY_SP, key, QueryMode::Ensure),
                 };
@@ -246,7 +304,7 @@ macro_rules! define_callbacks {
                 let key = key.into_query_param();
                 opt_remap_env_constness!([$($modifiers)*][key]);
 
-                match try_get_cached(self.tcx, &self.tcx.query_caches.$name, &key) {
+                match try_get_cached(self.tcx, &self.tcx.query_system.caches.$name, &key) {
                     Some(value) => value,
                     None => self.tcx.queries.$name(self.tcx, self.span, key, QueryMode::Get).unwrap(),
                 }
@@ -257,7 +315,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 {
@@ -334,12 +392,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();
                 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) => {
@@ -357,7 +416,8 @@ macro_rules! define_feedable {
                             &value,
                             hash_result!([$($modifiers)*]),
                         );
-                        cache.complete(key, value, dep_node_index)
+                        cache.complete(key, value, dep_node_index);
+                        value
                     }
                 }
             }

compiler/rustc_query_impl/src/lib.rs

@@ -21,7 +21,9 @@ 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::{
+    query_keys, query_provided, query_provided_to_value, query_storage, query_values,
+};
 use rustc_middle::ty::query::{ExternProviders, Providers, QueryEngine};
 use rustc_middle::ty::TyCtxt;
 use rustc_span::Span;

compiler/rustc_query_impl/src/plumbing.rs

@@ -293,14 +293,14 @@ macro_rules! get_provider {
 }
 
 macro_rules! should_ever_cache_on_disk {
-    ([]) => {{
-        None
+    ([]$yes:tt $no:tt) => {{
+        $no
     }};
-    ([(cache) $($rest:tt)*]) => {{
-        Some($crate::plumbing::try_load_from_disk::<Self::Value>)
+    ([(cache) $($rest:tt)*]$yes:tt $no:tt) => {{
+        $yes
     }};
-    ([$other:tt $($modifiers:tt)*]) => {
-        should_ever_cache_on_disk!([$($modifiers)*])
+    ([$other:tt $($modifiers:tt)*]$yes:tt $no:tt) => {
+        should_ever_cache_on_disk!([$($modifiers)*]$yes $no)
     };
 }
 
@@ -472,7 +472,6 @@ macro_rules! define_queries {
         $(impl<'tcx> QueryConfig<QueryCtxt<'tcx>> for queries::$name<'tcx> {
             type Key = query_keys::$name<'tcx>;
             type Value = query_values::$name<'tcx>;
-            type Stored = query_stored::$name<'tcx>;
             const NAME: &'static str = stringify!($name);
 
             #[inline]
@@ -493,24 +492,39 @@ macro_rules! define_queries {
             fn query_cache<'a>(tcx: QueryCtxt<'tcx>) -> &'a Self::Cache
                 where 'tcx:'a
             {
-                &tcx.query_caches.$name
+                &tcx.query_system.caches.$name
             }
 
-            fn execute_query(tcx: TyCtxt<'tcx>, key: Self::Key) -> Self::Stored {
+            fn execute_query(tcx: TyCtxt<'tcx>, key: Self::Key) -> Self::Value {
                 tcx.$name(key)
             }
 
             #[inline]
-            // key is only sometimes used
             #[allow(unused_variables)]
-            fn compute(qcx: QueryCtxt<'tcx>, key: &Self::Key) -> fn(TyCtxt<'tcx>, Self::Key) -> Self::Value {
-                get_provider!([$($modifiers)*][qcx, $name, key])
+            fn compute(qcx: QueryCtxt<'tcx>, key: Self::Key) -> Self::Value {
+                query_provided_to_value::$name(
+                    qcx.tcx,
+                    get_provider!([$($modifiers)*][qcx, $name, key])(qcx.tcx, key)
+                )
             }
 
             #[inline]
-            fn try_load_from_disk(qcx: QueryCtxt<'tcx>, key: &Self::Key) -> rustc_query_system::query::TryLoadFromDisk<QueryCtxt<'tcx>, Self> {
-                let cache_on_disk = Self::cache_on_disk(qcx.tcx, key);
-                if cache_on_disk { should_ever_cache_on_disk!([$($modifiers)*]) } else { None }
+            fn try_load_from_disk(_qcx: QueryCtxt<'tcx>, _key: &Self::Key) -> rustc_query_system::query::TryLoadFromDisk<QueryCtxt<'tcx>, Self> {
+                should_ever_cache_on_disk!([$($modifiers)*] {
+                    if Self::cache_on_disk(_qcx.tcx, _key) {
+                        Some(|qcx: QueryCtxt<'tcx>, dep_node| {
+                            let value = $crate::plumbing::try_load_from_disk::<query_provided::$name<'tcx>>(
+                                qcx,
+                                dep_node
+                            );
+                            value.map(|value| query_provided_to_value::$name(qcx.tcx, value))
+                        })
+                    } else {
+                        None
+                    }
+                } {
+                    None
+                })
             }
 
             const ANON: bool = is_anon!([$($modifiers)*]);
@@ -633,7 +647,7 @@ macro_rules! define_queries {
                     $crate::profiling_support::alloc_self_profile_query_strings_for_query_cache(
                         tcx,
                         stringify!($name),
-                        &tcx.query_caches.$name,
+                        &tcx.query_system.caches.$name,
                         string_cache,
                     )
                 },
@@ -725,7 +739,7 @@ macro_rules! define_queries_struct {
                 span: Span,
                 key: <queries::$name<'tcx> as QueryConfig<QueryCtxt<'tcx>>>::Key,
                 mode: QueryMode,
-            ) -> Option<query_stored::$name<'tcx>> {
+            ) -> Option<query_values::$name<'tcx>> {
                 let qcx = QueryCtxt { tcx, queries: self };
                 get_query::<queries::$name<'tcx>, _, rustc_middle::dep_graph::DepKind>(qcx, span, key, mode)
             })*