rustc_query_system/query/
caches.rs

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

rustc_query_system/query/
caches.rs
