Auto merge of #107833 - Zoxc:arena-query-clean, r=cjgillot

Factor query arena allocation out from query caches

This moves the logic for arena allocation out from the query caches into conditional code in the query system. The specialized arena caches are removed. A new `QuerySystem` type is added in `rustc_middle` which contains the arenas, providers and query caches.

Performance seems to be slightly regressed:
<table><tr><td rowspan="2">Benchmark</td><td colspan="1"><b>Before</b></th><td colspan="2"><b>After</b></th></tr><tr><td align="right">Time</td><td align="right">Time</td><td align="right">%</th></tr><tr><td>🟣 <b>clap</b>:check</td><td align="right">1.8053s</td><td align="right">1.8109s</td><td align="right"> 0.31%</td></tr><tr><td>🟣 <b>hyper</b>:check</td><td align="right">0.2600s</td><td align="right">0.2597s</td><td align="right"> -0.10%</td></tr><tr><td>🟣 <b>regex</b>:check</td><td align="right">0.9973s</td><td align="right">1.0006s</td><td align="right"> 0.34%</td></tr><tr><td>🟣 <b>syn</b>:check</td><td align="right">1.6048s</td><td align="right">1.6051s</td><td align="right"> 0.02%</td></tr><tr><td>🟣 <b>syntex_syntax</b>:check</td><td align="right">6.2992s</td><td align="right">6.3159s</td><td align="right"> 0.26%</td></tr><tr><td>Total</td><td align="right">10.9664s</td><td align="right">10.9922s</td><td align="right"> 0.23%</td></tr><tr><td>Summary</td><td align="right">1.0000s</td><td align="right">1.0017s</td><td align="right"> 0.17%</td></tr></table>

Incremental performance is a bit worse:
<table><tr><td rowspan="2">Benchmark</td><td colspan="1"><b>Before</b></th><td colspan="2"><b>After</b></th></tr><tr><td align="right">Time</td><td align="right">Time</td><td align="right">%</th></tr><tr><td>🟣 <b>clap</b>:check:initial</td><td align="right">2.2103s</td><td align="right">2.2247s</td><td align="right"> 0.65%</td></tr><tr><td>🟣 <b>hyper</b>:check:initial</td><td align="right">0.3335s</td><td align="right">0.3349s</td><td align="right"> 0.41%</td></tr><tr><td>🟣 <b>regex</b>:check:initial</td><td align="right">1.2597s</td><td align="right">1.2650s</td><td align="right"> 0.42%</td></tr><tr><td>🟣 <b>syn</b>:check:initial</td><td align="right">2.0521s</td><td align="right">2.0613s</td><td align="right"> 0.45%</td></tr><tr><td>🟣 <b>syntex_syntax</b>:check:initial</td><td align="right">7.8275s</td><td align="right">7.8583s</td><td align="right"> 0.39%</td></tr><tr><td>Total</td><td align="right">13.6832s</td><td align="right">13.7442s</td><td align="right"> 0.45%</td></tr><tr><td>Summary</td><td align="right">1.0000s</td><td align="right">1.0046s</td><td align="right"> 0.46%</td></tr></table>

It does seem like LLVM optimizers struggle a bit with the current state of the query system.

Based on top of #107782 and #107802.

r? `@cjgillot`

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

@@ -520,7 +520,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.
@@ -705,7 +705,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)
             })*