rustc_query_system/query/
caches.rsuse std::fmt::Debug;
use std::hash::Hash;
use rustc_data_structures::fx::FxHashMap;
use rustc_data_structures::sharded::{self, Sharded};
use rustc_data_structures::sync::OnceLock;
pub use rustc_data_structures::vec_cache::VecCache;
use rustc_hir::def_id::LOCAL_CRATE;
use rustc_index::Idx;
use rustc_span::def_id::{DefId, DefIndex};
use crate::dep_graph::DepNodeIndex;
pub trait QueryCache: Sized {
type Key: Hash + Eq + Copy + Debug;
type Value: Copy;
fn lookup(&self, key: &Self::Key) -> Option<(Self::Value, DepNodeIndex)>;
fn complete(&self, key: Self::Key, value: Self::Value, index: DepNodeIndex);
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex));
}
pub struct DefaultCache<K, V> {
cache: Sharded<FxHashMap<K, (V, DepNodeIndex)>>,
}
impl<K, V> Default for DefaultCache<K, V> {
fn default() -> Self {
DefaultCache { cache: Default::default() }
}
}
impl<K, V> QueryCache for DefaultCache<K, V>
where
K: Eq + Hash + Copy + Debug,
V: Copy,
{
type Key = K;
type Value = V;
#[inline(always)]
fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
let key_hash = sharded::make_hash(key);
let lock = self.cache.lock_shard_by_hash(key_hash);
let result = lock.raw_entry().from_key_hashed_nocheck(key_hash, key);
if let Some((_, value)) = result { Some(*value) } else { None }
}
#[inline]
fn complete(&self, key: K, value: V, index: DepNodeIndex) {
let mut lock = self.cache.lock_shard_by_value(&key);
lock.insert(key, (value, index));
}
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
for shard in self.cache.lock_shards() {
for (k, v) in shard.iter() {
f(k, &v.0, v.1);
}
}
}
}
pub struct SingleCache<V> {
cache: OnceLock<(V, DepNodeIndex)>,
}
impl<V> Default for SingleCache<V> {
fn default() -> Self {
SingleCache { cache: OnceLock::new() }
}
}
impl<V> QueryCache for SingleCache<V>
where
V: Copy,
{
type Key = ();
type Value = V;
#[inline(always)]
fn lookup(&self, _key: &()) -> Option<(V, DepNodeIndex)> {
self.cache.get().copied()
}
#[inline]
fn complete(&self, _key: (), value: V, index: DepNodeIndex) {
self.cache.set((value, index)).ok();
}
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
if let Some(value) = self.cache.get() {
f(&(), &value.0, value.1)
}
}
}
pub struct DefIdCache<V> {
local: VecCache<DefIndex, V, DepNodeIndex>,
foreign: DefaultCache<DefId, V>,
}
impl<V> Default for DefIdCache<V> {
fn default() -> Self {
DefIdCache { local: Default::default(), foreign: Default::default() }
}
}
impl<V> QueryCache for DefIdCache<V>
where
V: Copy,
{
type Key = DefId;
type Value = V;
#[inline(always)]
fn lookup(&self, key: &DefId) -> Option<(V, DepNodeIndex)> {
if key.krate == LOCAL_CRATE {
self.local.lookup(&key.index)
} else {
self.foreign.lookup(key)
}
}
#[inline]
fn complete(&self, key: DefId, value: V, index: DepNodeIndex) {
if key.krate == LOCAL_CRATE {
self.local.complete(key.index, value, index)
} else {
self.foreign.complete(key, value, index)
}
}
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
self.local.iter(&mut |key, value, index| {
f(&DefId { krate: LOCAL_CRATE, index: *key }, value, index);
});
self.foreign.iter(f);
}
}
impl<K, V> QueryCache for VecCache<K, V, DepNodeIndex>
where
K: Idx + Eq + Hash + Copy + Debug,
V: Copy,
{
type Key = K;
type Value = V;
#[inline(always)]
fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
self.lookup(key)
}
#[inline]
fn complete(&self, key: K, value: V, index: DepNodeIndex) {
self.complete(key, value, index)
}
fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
self.iter(f)
}
}