Auto merge of #70674 - cjgillot:query-arena-all, r=matthewjasper

Have the per-query caches store the results on arenas

This PR leverages the cache for each query to serve as storage area for the query results.

It introduces a new cache `ArenaCache`, which moves the result to an arena,
and only stores the reference in the hash map.
This allows to remove a sizeable part of the usage of the global `TyCtxt` arena.

I only migrated queries that already used arenas before.

5 files changed, 142 insertions, 36 deletions

diff --git a/src/librustc_query_system/Cargo.toml b/src/librustc_query_system/Cargo.toml
index 6629d841fc8..392e19e1f44 100644
--- a/src/librustc_query_system/Cargo.toml
+++ b/src/librustc_query_system/Cargo.toml
@@ -10,6 +10,7 @@ path = "lib.rs"
 doctest = false
 
 [dependencies]
+arena = { path = "../libarena" }
 log = { version = "0.4", features = ["release_max_level_info", "std"] }
 rustc-rayon-core = "0.3.0"
 rustc_data_structures = { path = "../librustc_data_structures" }
diff --git a/src/librustc_query_system/query/caches.rs b/src/librustc_query_system/query/caches.rs
index 0c0335ba04f..6a47abc5b46 100644
--- a/src/librustc_query_system/query/caches.rs
+++ b/src/librustc_query_system/query/caches.rs
@@ -2,19 +2,29 @@ use crate::dep_graph::DepNodeIndex;
 use crate::query::plumbing::{QueryLookup, QueryState};
 use crate::query::QueryContext;
 
+use arena::TypedArena;
 use rustc_data_structures::fx::FxHashMap;
 use rustc_data_structures::sharded::Sharded;
+use rustc_data_structures::sync::WorkerLocal;
 use std::default::Default;
 use std::hash::Hash;
 use std::marker::PhantomData;
 
-pub trait CacheSelector<K: Hash, V> {
-    type Cache: QueryCache<Key = K, Value = V>;
+pub trait CacheSelector<K, V> {
+    type Cache;
 }
 
-pub trait QueryCache: Default {
-    type Key: Hash;
+pub trait QueryStorage: Default {
     type Value;
+    type Stored: Clone;
+
+    /// Store a value without putting it in the cache.
+    /// This is meant to be used with cycle errors.
+    fn store_nocache(&self, value: Self::Value) -> Self::Stored;
+}
+
+pub trait QueryCache: QueryStorage {
+    type Key: Hash;
     type Sharded: Default;
 
     /// Checks if the query is already computed and in the cache.
@@ -30,7 +40,7 @@ pub trait QueryCache: Default {
         on_miss: OnMiss,
     ) -> R
     where
-        OnHit: FnOnce(&Self::Value, DepNodeIndex) -> R,
+        OnHit: FnOnce(&Self::Stored, DepNodeIndex) -> R,
         OnMiss: FnOnce(Self::Key, QueryLookup<'_, CTX, Self::Key, Self::Sharded>) -> R;
 
     fn complete<CTX: QueryContext>(
@@ -40,7 +50,7 @@ pub trait QueryCache: Default {
         key: Self::Key,
         value: Self::Value,
         index: DepNodeIndex,
-    );
+    ) -> Self::Stored;
 
     fn iter<R, L>(
         &self,
@@ -66,9 +76,19 @@ impl<K, V> Default for DefaultCache<K, V> {
     }
 }
 
+impl<K: Eq + Hash, V: Clone> QueryStorage for DefaultCache<K, V> {
+    type Value = V;
+    type Stored = V;
+
+    #[inline]
+    fn store_nocache(&self, value: Self::Value) -> Self::Stored {
+        // We have no dedicated storage
+        value
+    }
+}
+
 impl<K: Eq + Hash, V: Clone> QueryCache for DefaultCache<K, V> {
     type Key = K;
-    type Value = V;
     type Sharded = FxHashMap<K, (V, DepNodeIndex)>;
 
     #[inline(always)]
@@ -99,8 +119,94 @@ impl<K: Eq + Hash, V: Clone> QueryCache for DefaultCache<K, V> {
         key: K,
         value: V,
         index: DepNodeIndex,
-    ) {
-        lock_sharded_storage.insert(key, (value, index));
+    ) -> Self::Stored {
+        lock_sharded_storage.insert(key, (value.clone(), index));
+        value
+    }
+
+    fn iter<R, L>(
+        &self,
+        shards: &Sharded<L>,
+        get_shard: impl Fn(&mut L) -> &mut Self::Sharded,
+        f: impl for<'a> FnOnce(Box<dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)> + 'a>) -> R,
+    ) -> R {
+        let mut shards = shards.lock_shards();
+        let mut shards: Vec<_> = shards.iter_mut().map(|shard| get_shard(shard)).collect();
+        let results = shards.iter_mut().flat_map(|shard| shard.iter()).map(|(k, v)| (k, &v.0, v.1));
+        f(Box::new(results))
+    }
+}
+
+pub struct ArenaCacheSelector<'tcx>(PhantomData<&'tcx ()>);
+
+impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<K, V> for ArenaCacheSelector<'tcx> {
+    type Cache = ArenaCache<'tcx, K, V>;
+}
+
+pub struct ArenaCache<'tcx, K, V> {
+    arena: WorkerLocal<TypedArena<(V, DepNodeIndex)>>,
+    phantom: PhantomData<(K, &'tcx V)>,
+}
+
+impl<'tcx, K, V> Default for ArenaCache<'tcx, K, V> {
+    fn default() -> Self {
+        ArenaCache { arena: WorkerLocal::new(|_| TypedArena::default()), phantom: PhantomData }
+    }
+}
+
+impl<'tcx, K: Eq + Hash, V: 'tcx> QueryStorage for ArenaCache<'tcx, K, V> {
+    type Value = V;
+    type Stored = &'tcx V;
+
+    #[inline]
+    fn store_nocache(&self, value: Self::Value) -> Self::Stored {
+        let value = self.arena.alloc((value, DepNodeIndex::INVALID));
+        let value = unsafe { &*(&value.0 as *const _) };
+        &value
+    }
+}
+
+impl<'tcx, K: Eq + Hash, V: 'tcx> QueryCache for ArenaCache<'tcx, K, V> {
+    type Key = K;
+    type Sharded = FxHashMap<K, &'tcx (V, DepNodeIndex)>;
+
+    #[inline(always)]
+    fn lookup<CTX: QueryContext, R, OnHit, OnMiss>(
+        &self,
+        state: &QueryState<CTX, Self>,
+        key: K,
+        on_hit: OnHit,
+        on_miss: OnMiss,
+    ) -> R
+    where
+        OnHit: FnOnce(&&'tcx V, DepNodeIndex) -> R,
+        OnMiss: FnOnce(K, QueryLookup<'_, CTX, K, Self::Sharded>) -> R,
+    {
+        let mut lookup = state.get_lookup(&key);
+        let lock = &mut *lookup.lock;
+
+        let result = lock.cache.raw_entry().from_key_hashed_nocheck(lookup.key_hash, &key);
+
+        if let Some((_, value)) = result {
+            on_hit(&&value.0, value.1)
+        } else {
+            on_miss(key, lookup)
+        }
+    }
+
+    #[inline]
+    fn complete<CTX: QueryContext>(
+        &self,
+        _: CTX,
+        lock_sharded_storage: &mut Self::Sharded,
+        key: K,
+        value: V,
+        index: DepNodeIndex,
+    ) -> Self::Stored {
+        let value = self.arena.alloc((value, index));
+        let value = unsafe { &*(value as *const _) };
+        lock_sharded_storage.insert(key, value);
+        &value.0
     }
 
     fn iter<R, L>(
diff --git a/src/librustc_query_system/query/config.rs b/src/librustc_query_system/query/config.rs
index 06e3302b263..710ec3bfb0b 100644
--- a/src/librustc_query_system/query/config.rs
+++ b/src/librustc_query_system/query/config.rs
@@ -20,7 +20,8 @@ pub trait QueryConfig<CTX> {
     const CATEGORY: ProfileCategory;
 
     type Key: Eq + Hash + Clone + Debug;
-    type Value: Clone;
+    type Value;
+    type Stored: Clone;
 }
 
 pub trait QueryAccessors<CTX: QueryContext>: QueryConfig<CTX> {
@@ -28,7 +29,7 @@ pub trait QueryAccessors<CTX: QueryContext>: QueryConfig<CTX> {
     const EVAL_ALWAYS: bool;
     const DEP_KIND: CTX::DepKind;
 
-    type Cache: QueryCache<Key = Self::Key, Value = Self::Value>;
+    type Cache: QueryCache<Key = Self::Key, Stored = Self::Stored, Value = Self::Value>;
 
     // Don't use this method to access query results, instead use the methods on TyCtxt
     fn query_state<'a>(tcx: CTX) -> &'a QueryState<CTX, Self::Cache>;
diff --git a/src/librustc_query_system/query/mod.rs b/src/librustc_query_system/query/mod.rs
index b1677c5c93d..49097725bc9 100644
--- a/src/librustc_query_system/query/mod.rs
+++ b/src/librustc_query_system/query/mod.rs
@@ -7,7 +7,9 @@ pub use self::job::deadlock;
 pub use self::job::{QueryInfo, QueryJob, QueryJobId, QueryJobInfo};
 
 mod caches;
-pub use self::caches::{CacheSelector, DefaultCacheSelector, QueryCache};
+pub use self::caches::{
+    ArenaCacheSelector, CacheSelector, DefaultCacheSelector, QueryCache, QueryStorage,
+};
 
 mod config;
 pub use self::config::{QueryAccessors, QueryConfig, QueryDescription};
diff --git a/src/librustc_query_system/query/plumbing.rs b/src/librustc_query_system/query/plumbing.rs
index 9da13f23664..0aeec269e61 100644
--- a/src/librustc_query_system/query/plumbing.rs
+++ b/src/librustc_query_system/query/plumbing.rs
@@ -148,7 +148,6 @@ struct JobOwner<'tcx, CTX: QueryContext, C>
 where
     C: QueryCache,
     C::Key: Eq + Hash + Clone + Debug,
-    C::Value: Clone,
 {
     state: &'tcx QueryState<CTX, C>,
     key: C::Key,
@@ -159,7 +158,6 @@ impl<'tcx, CTX: QueryContext, C> JobOwner<'tcx, CTX, C>
 where
     C: QueryCache,
     C::Key: Eq + Hash + Clone + Debug,
-    C::Value: Clone,
 {
     /// Either gets a `JobOwner` corresponding the query, allowing us to
     /// start executing the query, or returns with the result of the query.
@@ -177,7 +175,7 @@ where
         mut lookup: QueryLookup<'a, CTX, C::Key, C::Sharded>,
     ) -> TryGetJob<'b, CTX, C>
     where
-        Q: QueryDescription<CTX, Key = C::Key, Value = C::Value, Cache = C>,
+        Q: QueryDescription<CTX, Key = C::Key, Stored = C::Stored, Value = C::Value, Cache = C>,
         CTX: QueryContext,
     {
         let lock = &mut *lookup.lock;
@@ -229,7 +227,8 @@ where
         // so we just return the error.
         #[cfg(not(parallel_compiler))]
         return TryGetJob::Cycle(cold_path(|| {
-            Q::handle_cycle_error(tcx, latch.find_cycle_in_stack(tcx, span))
+            let value = Q::handle_cycle_error(tcx, latch.find_cycle_in_stack(tcx, span));
+            Q::query_state(tcx).cache.store_nocache(value)
         }));
 
         // With parallel queries we might just have to wait on some other
@@ -239,7 +238,9 @@ where
             let result = latch.wait_on(tcx, span);
 
             if let Err(cycle) = result {
-                return TryGetJob::Cycle(Q::handle_cycle_error(tcx, cycle));
+                let value = Q::handle_cycle_error(tcx, cycle);
+                let value = Q::query_state(tcx).cache.store_nocache(value);
+                return TryGetJob::Cycle(value);
             }
 
             let cached = try_get_cached(
@@ -261,7 +262,7 @@ where
     /// Completes the query by updating the query cache with the `result`,
     /// signals the waiter and forgets the JobOwner, so it won't poison the query
     #[inline(always)]
-    fn complete(self, tcx: CTX, result: &C::Value, dep_node_index: DepNodeIndex) {
+    fn complete(self, tcx: CTX, result: C::Value, dep_node_index: DepNodeIndex) -> C::Stored {
         // We can move out of `self` here because we `mem::forget` it below
         let key = unsafe { ptr::read(&self.key) };
         let state = self.state;
@@ -269,18 +270,18 @@ where
         // Forget ourself so our destructor won't poison the query
         mem::forget(self);
 
-        let job = {
-            let result = result.clone();
+        let (job, result) = {
             let mut lock = state.shards.get_shard_by_value(&key).lock();
             let job = match lock.active.remove(&key).unwrap() {
                 QueryResult::Started(job) => job,
                 QueryResult::Poisoned => panic!(),
             };
-            state.cache.complete(tcx, &mut lock.cache, key, result, dep_node_index);
-            job
+            let result = state.cache.complete(tcx, &mut lock.cache, key, result, dep_node_index);
+            (job, result)
         };
 
         job.signal_complete();
+        result
     }
 }
 
@@ -297,7 +298,6 @@ where
 impl<'tcx, CTX: QueryContext, C: QueryCache> Drop for JobOwner<'tcx, CTX, C>
 where
     C::Key: Eq + Hash + Clone + Debug,
-    C::Value: Clone,
 {
     #[inline(never)]
     #[cold]
@@ -331,7 +331,6 @@ pub struct CycleError<Q> {
 enum TryGetJob<'tcx, CTX: QueryContext, C: QueryCache>
 where
     C::Key: Eq + Hash + Clone + Debug,
-    C::Value: Clone,
 {
     /// The query is not yet started. Contains a guard to the cache eventually used to start it.
     NotYetStarted(JobOwner<'tcx, CTX, C>),
@@ -340,10 +339,10 @@ where
     /// Returns the result of the query and its dep-node index
     /// if it succeeded or a cycle error if it failed.
     #[cfg(parallel_compiler)]
-    JobCompleted((C::Value, DepNodeIndex)),
+    JobCompleted((C::Stored, DepNodeIndex)),
 
     /// Trying to execute the query resulted in a cycle.
-    Cycle(C::Value),
+    Cycle(C::Stored),
 }
 
 /// Checks if the query is already computed and in the cache.
@@ -362,7 +361,7 @@ fn try_get_cached<CTX, C, R, OnHit, OnMiss>(
 where
     C: QueryCache,
     CTX: QueryContext,
-    OnHit: FnOnce(&C::Value, DepNodeIndex) -> R,
+    OnHit: FnOnce(&C::Stored, DepNodeIndex) -> R,
     OnMiss: FnOnce(C::Key, QueryLookup<'_, CTX, C::Key, C::Sharded>) -> R,
 {
     state.cache.lookup(
@@ -388,7 +387,7 @@ fn try_execute_query<Q, CTX>(
     span: Span,
     key: Q::Key,
     lookup: QueryLookup<'_, CTX, Q::Key, <Q::Cache as QueryCache>::Sharded>,
-) -> Q::Value
+) -> Q::Stored
 where
     Q: QueryDescription<CTX>,
     CTX: QueryContext,
@@ -427,9 +426,7 @@ where
             tcx.store_diagnostics_for_anon_node(dep_node_index, diagnostics);
         }
 
-        job.complete(tcx, &result, dep_node_index);
-
-        return result;
+        return job.complete(tcx, result, dep_node_index);
     }
 
     let dep_node = Q::to_dep_node(tcx, &key);
@@ -454,8 +451,7 @@ where
             })
         });
         if let Some((result, dep_node_index)) = loaded {
-            job.complete(tcx, &result, dep_node_index);
-            return result;
+            return job.complete(tcx, result, dep_node_index);
         }
     }
 
@@ -558,7 +554,7 @@ fn force_query_with_job<Q, CTX>(
     key: Q::Key,
     job: JobOwner<'_, CTX, Q::Cache>,
     dep_node: DepNode<CTX::DepKind>,
-) -> (Q::Value, DepNodeIndex)
+) -> (Q::Stored, DepNodeIndex)
 where
     Q: QueryDescription<CTX>,
     CTX: QueryContext,
@@ -603,13 +599,13 @@ where
         }
     }
 
-    job.complete(tcx, &result, dep_node_index);
+    let result = job.complete(tcx, result, dep_node_index);
 
     (result, dep_node_index)
 }
 
 #[inline(never)]
-pub fn get_query<Q, CTX>(tcx: CTX, span: Span, key: Q::Key) -> Q::Value
+pub fn get_query<Q, CTX>(tcx: CTX, span: Span, key: Q::Key) -> Q::Stored
 where
     Q: QueryDescription<CTX>,
     CTX: QueryContext,