diff --git a/src/librustc/query/mod.rs b/src/librustc/query/mod.rs index ff3a82e53639e..9ccd0946ba3a4 100644 --- a/src/librustc/query/mod.rs +++ b/src/librustc/query/mod.rs @@ -75,6 +75,7 @@ rustc_queries! { // This can be conveniently accessed by methods on `tcx.hir()`. // Avoid calling this query directly. query hir_owner(key: DefId) -> &'tcx HirOwner<'tcx> { + storage(caches::LocalDenseDefIdCacheSelector<&'tcx HirOwner<'tcx>>) eval_always } @@ -83,6 +84,7 @@ rustc_queries! { // This can be conveniently accessed by methods on `tcx.hir()`. // Avoid calling this query directly. query hir_owner_items(key: DefId) -> &'tcx HirOwnerItems<'tcx> { + storage(caches::LocalDenseDefIdCacheSelector<&'tcx HirOwnerItems<'tcx>>) eval_always } diff --git a/src/librustc/ty/query/caches.rs b/src/librustc/ty/query/caches.rs index efc2804bd4d59..ca0e3c5345c66 100644 --- a/src/librustc/ty/query/caches.rs +++ b/src/librustc/ty/query/caches.rs @@ -5,8 +5,12 @@ use crate::ty::TyCtxt; use rustc_data_structures::fx::FxHashMap; use rustc_data_structures::sharded::Sharded; +use rustc_hir::def_id::{DefId, DefIndex, LOCAL_CRATE}; +use rustc_index::vec::IndexVec; +use std::cell::RefCell; use std::default::Default; use std::hash::Hash; +use std::marker::PhantomData; pub(crate) trait CacheSelector { type Cache: QueryCache; @@ -54,13 +58,18 @@ pub(crate) trait QueryCache: Default { pub struct DefaultCacheSelector; impl CacheSelector for DefaultCacheSelector { - type Cache = DefaultCache; + type Cache = DefaultCache<()>; } -#[derive(Default)] -pub struct DefaultCache; +pub struct DefaultCache(PhantomData); -impl QueryCache for DefaultCache { +impl Default for DefaultCache { + fn default() -> Self { + DefaultCache(PhantomData) + } +} + +impl QueryCache for DefaultCache { type Sharded = FxHashMap; #[inline(always)] @@ -110,3 +119,88 @@ impl QueryCache for DefaultCache { f(Box::new(results)) } } + +#[cfg(parallel_compiler)] +pub type LocalDenseDefIdCacheSelector = DefaultCache; +#[cfg(not(parallel_compiler))] +pub type LocalDenseDefIdCacheSelector = LocalDenseDefIdCache; + +pub struct LocalDenseDefIdCache { + local: RefCell>>, + other: DefaultCache<()>, +} + +impl Default for LocalDenseDefIdCache { + fn default() -> Self { + LocalDenseDefIdCache { local: RefCell::new(IndexVec::new()), other: Default::default() } + } +} + +impl QueryCache for LocalDenseDefIdCache { + type Sharded = as QueryCache>::Sharded; + + #[inline(always)] + fn lookup<'tcx, R, GetCache, OnHit, OnMiss, Q>( + &self, + state: &'tcx QueryState<'tcx, Q>, + get_cache: GetCache, + key: DefId, + on_hit: OnHit, + on_miss: OnMiss, + ) -> R + where + Q: QueryAccessors<'tcx>, + GetCache: for<'a> Fn(&'a mut QueryStateShard<'tcx, Q>) -> &'a mut Self::Sharded, + OnHit: FnOnce(&V, DepNodeIndex) -> R, + OnMiss: FnOnce(DefId, QueryLookup<'tcx, Q>) -> R, + { + if key.krate == LOCAL_CRATE { + let local = self.local.borrow(); + if let Some(result) = local.get(key.index).and_then(|v| v.as_ref()) { + on_hit(&result.0, result.1) + } else { + drop(local); + let lookup = state.get_lookup(&key); + on_miss(key, lookup) + } + } else { + self.other.lookup(state, get_cache, key, on_hit, on_miss) + } + } + + #[inline] + fn complete( + &self, + tcx: TyCtxt<'tcx>, + lock_sharded_storage: &mut Self::Sharded, + key: DefId, + value: V, + index: DepNodeIndex, + ) { + if key.krate == LOCAL_CRATE { + let mut local = self.local.borrow_mut(); + if local.raw.capacity() == 0 { + *local = IndexVec::from_elem_n(None, tcx.hir().definitions().def_index_count()); + } + local[key.index] = Some((value, index)); + } else { + self.other.complete(tcx, lock_sharded_storage, key, value, index); + } + } + + fn iter( + &self, + shards: &Sharded, + get_shard: impl Fn(&mut L) -> &mut Self::Sharded, + f: impl for<'a> FnOnce(Box + 'a>) -> R, + ) -> R { + let local = self.local.borrow(); + let local: Vec<(DefId, &V, DepNodeIndex)> = local + .iter_enumerated() + .filter_map(|(i, e)| e.as_ref().map(|e| (DefId::local(i), &e.0, e.1))) + .collect(); + self.other.iter(shards, get_shard, |results| { + f(Box::new(results.chain(local.iter().map(|(id, v, i)| (id, *v, *i))))) + }) + } +} diff --git a/src/librustc_macros/src/query.rs b/src/librustc_macros/src/query.rs index 56b7be2f7e2d5..d7de97006de17 100644 --- a/src/librustc_macros/src/query.rs +++ b/src/librustc_macros/src/query.rs @@ -107,7 +107,9 @@ impl Parse for QueryModifier { let block = input.parse()?; Ok(QueryModifier::LoadCached(tcx, id, block)) } else if modifier == "storage" { - let ty = input.parse()?; + let args; + parenthesized!(args in input); + let ty = args.parse()?; Ok(QueryModifier::Storage(ty)) } else if modifier == "fatal_cycle" { Ok(QueryModifier::FatalCycle)