Introduce ArenaStorage.

3 files changed, 91 insertions, 1 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 7b9d75037c2..d921bcdcdd3 100644
--- a/src/librustc_query_system/query/caches.rs
+++ b/src/librustc_query_system/query/caches.rs
@@ -2,8 +2,10 @@ 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;
@@ -133,3 +135,88 @@ impl<K: Eq + Hash, V: Clone> QueryCache for DefaultCache<K, V> {
         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<&'tcx TypedArena<(V, DepNodeIndex)>>,
+    phantom: PhantomData<K>,
+}
+
+impl<'tcx, K, V> Default for ArenaCache<'tcx, K, V> {
+    fn default() -> Self {
+        ArenaCache {
+            arena: WorkerLocal::new(|_| &*Box::leak(Box::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;
+
+    fn store_nocache(&self, value: Self::Value) -> Self::Stored {
+        let value = self.arena.alloc((value, DepNodeIndex::INVALID));
+        &value.0
+    }
+}
+
+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));
+        lock_sharded_storage.insert(key, value);
+        &value.0
+    }
+
+    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))
+    }
+}
diff --git a/src/librustc_query_system/query/mod.rs b/src/librustc_query_system/query/mod.rs
index 83513b7e11d..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, QueryStorage};
+pub use self::caches::{
+    ArenaCacheSelector, CacheSelector, DefaultCacheSelector, QueryCache, QueryStorage,
+};
 
 mod config;
 pub use self::config::{QueryAccessors, QueryConfig, QueryDescription};