rustc_middle/query/
caches.rs

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

rustc_middle/query/
caches.rs
