rustc_query_system/query/
caches.rs

1use std::fmt::Debug;
2use std::hash::Hash;
3
4use rustc_data_structures::fx::FxHashMap;
5use rustc_data_structures::sharded::{self, Sharded};
6use rustc_data_structures::sync::OnceLock;
7pub use rustc_data_structures::vec_cache::VecCache;
8use rustc_hir::def_id::LOCAL_CRATE;
9use rustc_index::Idx;
10use rustc_span::def_id::{DefId, DefIndex};
11
12use crate::dep_graph::DepNodeIndex;
13
14/// Trait for types that serve as an in-memory cache for query results,
15/// for a given key (argument) type and value (return) type.
16///
17/// Types implementing this trait are associated with actual key/value types
18/// by the `Cache` associated type of the `rustc_middle::query::Key` trait.
19pub trait QueryCache: Sized {
20    type Key: Hash + Eq + Copy + Debug;
21    type Value: Copy;
22
23    /// Returns the cached value (and other information) associated with the
24    /// given key, if it is present in the cache.
25    fn lookup(&self, key: &Self::Key) -> Option<(Self::Value, DepNodeIndex)>;
26
27    /// Adds a key/value entry to this cache.
28    ///
29    /// Called by some part of the query system, after having obtained the
30    /// value by executing the query or loading a cached value from disk.
31    fn complete(&self, key: Self::Key, value: Self::Value, index: DepNodeIndex);
32
33    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex));
34}
35
36/// In-memory cache for queries whose keys aren't suitable for any of the
37/// more specialized kinds of cache. Backed by a sharded hashmap.
38pub struct DefaultCache<K, V> {
39    cache: Sharded<FxHashMap<K, (V, DepNodeIndex)>>,
40}
41
42impl<K, V> Default for DefaultCache<K, V> {
43    fn default() -> Self {
44        DefaultCache { cache: Default::default() }
45    }
46}
47
48impl<K, V> QueryCache for DefaultCache<K, V>
49where
50    K: Eq + Hash + Copy + Debug,
51    V: Copy,
52{
53    type Key = K;
54    type Value = V;
55
56    #[inline(always)]
57    fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
58        let key_hash = sharded::make_hash(key);
59        let lock = self.cache.lock_shard_by_hash(key_hash);
60        let result = lock.raw_entry().from_key_hashed_nocheck(key_hash, key);
61
62        if let Some((_, value)) = result { Some(*value) } else { None }
63    }
64
65    #[inline]
66    fn complete(&self, key: K, value: V, index: DepNodeIndex) {
67        let mut lock = self.cache.lock_shard_by_value(&key);
68        // We may be overwriting another value. This is all right, since the dep-graph
69        // will check that the fingerprint matches.
70        lock.insert(key, (value, index));
71    }
72
73    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
74        for shard in self.cache.lock_shards() {
75            for (k, v) in shard.iter() {
76                f(k, &v.0, v.1);
77            }
78        }
79    }
80}
81
82/// In-memory cache for queries whose key type only has one value (e.g. `()`).
83/// The cache therefore only needs to store one query return value.
84pub struct SingleCache<V> {
85    cache: OnceLock<(V, DepNodeIndex)>,
86}
87
88impl<V> Default for SingleCache<V> {
89    fn default() -> Self {
90        SingleCache { cache: OnceLock::new() }
91    }
92}
93
94impl<V> QueryCache for SingleCache<V>
95where
96    V: Copy,
97{
98    type Key = ();
99    type Value = V;
100
101    #[inline(always)]
102    fn lookup(&self, _key: &()) -> Option<(V, DepNodeIndex)> {
103        self.cache.get().copied()
104    }
105
106    #[inline]
107    fn complete(&self, _key: (), value: V, index: DepNodeIndex) {
108        self.cache.set((value, index)).ok();
109    }
110
111    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
112        if let Some(value) = self.cache.get() {
113            f(&(), &value.0, value.1)
114        }
115    }
116}
117
118/// In-memory cache for queries whose key is a [`DefId`].
119///
120/// Selects between one of two internal caches, depending on whether the key
121/// is a local ID or foreign-crate ID.
122pub struct DefIdCache<V> {
123    /// Stores the local DefIds in a dense map. Local queries are much more often dense, so this is
124    /// a win over hashing query keys at marginal memory cost (~5% at most) compared to FxHashMap.
125    local: VecCache<DefIndex, V, DepNodeIndex>,
126    foreign: DefaultCache<DefId, V>,
127}
128
129impl<V> Default for DefIdCache<V> {
130    fn default() -> Self {
131        DefIdCache { local: Default::default(), foreign: Default::default() }
132    }
133}
134
135impl<V> QueryCache for DefIdCache<V>
136where
137    V: Copy,
138{
139    type Key = DefId;
140    type Value = V;
141
142    #[inline(always)]
143    fn lookup(&self, key: &DefId) -> Option<(V, DepNodeIndex)> {
144        if key.krate == LOCAL_CRATE {
145            self.local.lookup(&key.index)
146        } else {
147            self.foreign.lookup(key)
148        }
149    }
150
151    #[inline]
152    fn complete(&self, key: DefId, value: V, index: DepNodeIndex) {
153        if key.krate == LOCAL_CRATE {
154            self.local.complete(key.index, value, index)
155        } else {
156            self.foreign.complete(key, value, index)
157        }
158    }
159
160    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
161        self.local.iter(&mut |key, value, index| {
162            f(&DefId { krate: LOCAL_CRATE, index: *key }, value, index);
163        });
164        self.foreign.iter(f);
165    }
166}
167
168impl<K, V> QueryCache for VecCache<K, V, DepNodeIndex>
169where
170    K: Idx + Eq + Hash + Copy + Debug,
171    V: Copy,
172{
173    type Key = K;
174    type Value = V;
175
176    #[inline(always)]
177    fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
178        self.lookup(key)
179    }
180
181    #[inline]
182    fn complete(&self, key: K, value: V, index: DepNodeIndex) {
183        self.complete(key, value, index)
184    }
185
186    fn iter(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
187        self.iter(f)
188    }
189}
rustc_query_system/query/caches.rs

rustc_query_system/query/
caches.rs
