Introduce ArenaStorage.

7 files changed, 97 insertions, 5 deletions

diff --git a/Cargo.lock b/Cargo.lock
index 816d65cb2ce..e92d298105c 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4125,6 +4125,7 @@ dependencies = [
 name = "rustc_query_system"
 version = "0.0.0"
 dependencies = [
+ "arena",
  "log",
  "parking_lot 0.10.2",
  "rustc-rayon-core",
diff --git a/src/librustc_middle/query/mod.rs b/src/librustc_middle/query/mod.rs
index c3e023f4029..f2a131a5ab2 100644
--- a/src/librustc_middle/query/mod.rs
+++ b/src/librustc_middle/query/mod.rs
@@ -156,7 +156,8 @@ rustc_queries! {
         /// Set of all the `DefId`s in this crate that have MIR associated with
         /// them. This includes all the body owners, but also things like struct
         /// constructors.
-        query mir_keys(_: CrateNum) -> &'tcx FxHashSet<LocalDefId> {
+        query mir_keys(_: CrateNum) -> FxHashSet<LocalDefId> {
+            storage(ArenaCacheSelector<'tcx>)
             desc { "getting a list of all mir_keys" }
         }
 
diff --git a/src/librustc_middle/ty/query/plumbing.rs b/src/librustc_middle/ty/query/plumbing.rs
index a9367dc40b7..e4237df5923 100644
--- a/src/librustc_middle/ty/query/plumbing.rs
+++ b/src/librustc_middle/ty/query/plumbing.rs
@@ -213,7 +213,7 @@ macro_rules! query_storage {
         <<$K as Key>::CacheSelector as CacheSelector<$K, $V>>::Cache
     };
     ([storage($ty:ty) $($rest:tt)*][$K:ty, $V:ty]) => {
-        $ty
+        <$ty as CacheSelector<$K, $V>>::Cache
     };
     ([$other:ident $(($($other_args:tt)*))* $(, $($modifiers:tt)*)*][$($args:tt)*]) => {
         query_storage!([$($($modifiers)*)*][$($args)*])
diff --git a/src/librustc_mir/transform/mod.rs b/src/librustc_mir/transform/mod.rs
index 9f63340065f..5cf74bc7b07 100644
--- a/src/librustc_mir/transform/mod.rs
+++ b/src/librustc_mir/transform/mod.rs
@@ -60,7 +60,7 @@ fn is_mir_available(tcx: TyCtxt<'_>, def_id: DefId) -> bool {
 
 /// Finds the full set of `DefId`s within the current crate that have
 /// MIR associated with them.
-fn mir_keys(tcx: TyCtxt<'_>, krate: CrateNum) -> &FxHashSet<LocalDefId> {
+fn mir_keys(tcx: TyCtxt<'_>, krate: CrateNum) -> FxHashSet<LocalDefId> {
     assert_eq!(krate, LOCAL_CRATE);
 
     let mut set = FxHashSet::default();
@@ -97,7 +97,7 @@ fn mir_keys(tcx: TyCtxt<'_>, krate: CrateNum) -> &FxHashSet<LocalDefId> {
         .krate()
         .visit_all_item_likes(&mut GatherCtors { tcx, set: &mut set }.as_deep_visitor());
 
-    tcx.arena.alloc(set)
+    set
 }
 
 /// Where a specific `mir::Body` comes from.
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};