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 { + query mir_keys(_: CrateNum) -> FxHashSet { + 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 { +fn mir_keys(tcx: TyCtxt<'_>, krate: CrateNum) -> FxHashSet { assert_eq!(krate, LOCAL_CRATE); let mut set = FxHashSet::default(); @@ -97,7 +97,7 @@ fn mir_keys(tcx: TyCtxt<'_>, krate: CrateNum) -> &FxHashSet { .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 QueryCache for DefaultCache { f(Box::new(results)) } } + +pub struct ArenaCacheSelector<'tcx>(PhantomData<&'tcx ()>); + +impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector for ArenaCacheSelector<'tcx> { + type Cache = ArenaCache<'tcx, K, V>; +} + +pub struct ArenaCache<'tcx, K, V> { + arena: WorkerLocal<&'tcx TypedArena<(V, DepNodeIndex)>>, + phantom: PhantomData, +} + +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; + + #[inline(always)] + fn lookup( + &self, + state: &QueryState, + 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( + &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( + &self, + shards: &Sharded, + get_shard: impl Fn(&mut L) -> &mut Self::Sharded, + f: impl for<'a> FnOnce(Box + '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};