From 37849643c670f7f0809948986329b04ce5c90405 Mon Sep 17 00:00:00 2001 From: Mark Rousskov Date: Sun, 14 Jan 2024 18:08:02 -0500 Subject: [PATCH] Cache local DefId-keyed queries without hashing Foreign maps are used to cache external DefIds, typically backed by metadata decoding. In the future we might skip caching `V` there (since loading from metadata usually is already cheap enough), but for now this cuts down on the impact to memory usage and time to None-init a bunch of memory. Foreign data is usually much sparser, since we're not usually loading *all* entries from the foreign crate(s). --- compiler/rustc_middle/src/query/keys.rs | 3 +- .../rustc_query_system/src/query/caches.rs | 79 ++++++++++++++++++- compiler/rustc_query_system/src/query/mod.rs | 3 +- 3 files changed, 82 insertions(+), 3 deletions(-) diff --git a/compiler/rustc_middle/src/query/keys.rs b/compiler/rustc_middle/src/query/keys.rs index 945f17d5df2..69d3974184d 100644 --- a/compiler/rustc_middle/src/query/keys.rs +++ b/compiler/rustc_middle/src/query/keys.rs @@ -9,6 +9,7 @@ use crate::ty::{GenericArg, GenericArgsRef}; use rustc_hir::def_id::{CrateNum, DefId, LocalDefId, LocalModDefId, ModDefId, LOCAL_CRATE}; use rustc_hir::hir_id::{HirId, OwnerId}; +use rustc_query_system::query::DefIdCacheSelector; use rustc_query_system::query::{DefaultCacheSelector, SingleCacheSelector, VecCacheSelector}; use rustc_span::symbol::{Ident, Symbol}; use rustc_span::{Span, DUMMY_SP}; @@ -152,7 +153,7 @@ fn key_as_def_id(&self) -> Option { } impl Key for DefId { - type CacheSelector = DefaultCacheSelector; + type CacheSelector = DefIdCacheSelector; fn default_span(&self, tcx: TyCtxt<'_>) -> Span { tcx.def_span(*self) diff --git a/compiler/rustc_query_system/src/query/caches.rs b/compiler/rustc_query_system/src/query/caches.rs index 0240f012da0..674a0984ae9 100644 --- a/compiler/rustc_query_system/src/query/caches.rs +++ b/compiler/rustc_query_system/src/query/caches.rs @@ -2,8 +2,11 @@ use rustc_data_structures::fx::FxHashMap; use rustc_data_structures::sharded::{self, Sharded}; -use rustc_data_structures::sync::OnceLock; +use rustc_data_structures::sync::{Lock, OnceLock}; +use rustc_hir::def_id::LOCAL_CRATE; use rustc_index::{Idx, IndexVec}; +use rustc_span::def_id::DefId; +use rustc_span::def_id::DefIndex; use std::fmt::Debug; use std::hash::Hash; use std::marker::PhantomData; @@ -148,6 +151,8 @@ impl QueryCache for VecCache #[inline(always)] fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> { + // FIXME: lock_shard_by_hash will use high bits which are usually zero in the index() passed + // here. This makes sharding essentially useless, always selecting the zero'th shard. let lock = self.cache.lock_shard_by_hash(key.index() as u64); if let Some(Some(value)) = lock.get(*key) { Some(*value) } else { None } } @@ -168,3 +173,75 @@ fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) { } } } + +pub struct DefIdCacheSelector; + +impl<'tcx, V: 'tcx> CacheSelector<'tcx, V> for DefIdCacheSelector { + type Cache = DefIdCache + where + V: Copy; +} + +pub struct DefIdCache { + /// Stores the local DefIds in a dense map. Local queries are much more often dense, so this is + /// a win over hashing query keys at marginal memory cost (~5% at most) compared to FxHashMap. + /// + /// The second element of the tuple is the set of keys actually present in the IndexVec, used + /// for faster iteration in `iter()`. + // FIXME: This may want to be sharded, like VecCache. However *how* to shard an IndexVec isn't + // super clear; VecCache is effectively not sharded today (see FIXME there). For now just omit + // that complexity here. + local: Lock<(IndexVec>, Vec)>, + foreign: DefaultCache, +} + +impl Default for DefIdCache { + fn default() -> Self { + DefIdCache { local: Default::default(), foreign: Default::default() } + } +} + +impl QueryCache for DefIdCache +where + V: Copy, +{ + type Key = DefId; + type Value = V; + + #[inline(always)] + fn lookup(&self, key: &DefId) -> Option<(V, DepNodeIndex)> { + if key.krate == LOCAL_CRATE { + let cache = self.local.lock(); + cache.0.get(key.index).and_then(|v| *v) + } else { + self.foreign.lookup(key) + } + } + + #[inline] + fn complete(&self, key: DefId, value: V, index: DepNodeIndex) { + if key.krate == LOCAL_CRATE { + let mut cache = self.local.lock(); + let (cache, present) = &mut *cache; + let slot = cache.ensure_contains_elem(key.index, Default::default); + if slot.is_none() { + // FIXME: Only store the present set when running in incremental mode. `iter` is not + // used outside of saving caches to disk and self-profile. + present.push(key.index); + } + *slot = Some((value, index)); + } else { + self.foreign.complete(key, value, index) + } + } + + fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) { + let guard = self.local.lock(); + let (cache, present) = &*guard; + for &idx in present.iter() { + let value = cache[idx].unwrap(); + f(&DefId { krate: LOCAL_CRATE, index: idx }, &value.0, value.1); + } + self.foreign.iter(f); + } +} diff --git a/compiler/rustc_query_system/src/query/mod.rs b/compiler/rustc_query_system/src/query/mod.rs index 9ff04c4e910..eecbf86c173 100644 --- a/compiler/rustc_query_system/src/query/mod.rs +++ b/compiler/rustc_query_system/src/query/mod.rs @@ -10,7 +10,8 @@ mod caches; pub use self::caches::{ - CacheSelector, DefaultCacheSelector, QueryCache, SingleCacheSelector, VecCacheSelector, + CacheSelector, DefIdCacheSelector, DefaultCacheSelector, QueryCache, SingleCacheSelector, + VecCacheSelector, }; mod config;