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    /// Calls a closure on each entry in this cache.
32    fn for_each(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex));
33
34    /// Returns the number of entries currently in this cache.
35    ///
36    /// Useful for reserving capacity in data structures that will hold the
37    /// output of a call to [`Self::for_each`].
38    fn len(&self) -> usize;
39}
40
41/// In-memory cache for queries whose keys aren't suitable for any of the
42/// more specialized kinds of cache. Backed by a sharded hashmap.
43pub struct DefaultCache<K, V> {
44    cache: ShardedHashMap<K, (V, DepNodeIndex)>,
45}
46
47impl<K, V> Default for DefaultCache<K, V> {
48    fn default() -> Self {
49        DefaultCache { cache: Default::default() }
50    }
51}
52
53impl<K, V> QueryCache for DefaultCache<K, V>
54where
55    K: QueryKey,
56    V: Copy,
57{
58    type Key = K;
59    type Value = V;
60
61    #[inline(always)]
62    fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
63        self.cache.get(key)
64    }
65
66    #[inline]
67    fn complete(&self, key: K, value: V, index: DepNodeIndex) {
68        self.cache.insert_unique(key, (value, index));
69    }
70
71    fn for_each(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
72        for shard in self.cache.lock_shards() {
73            for (k, v) in shard.iter() {
74                f(k, &v.0, v.1);
75            }
76        }
77    }
78
79    fn len(&self) -> usize {
80        self.cache.len()
81    }
82}
83
84/// In-memory cache for queries whose key type only has one value (e.g. `()`).
85/// The cache therefore only needs to store one query return value.
86pub struct SingleCache<V> {
87    cache: OnceLock<(V, DepNodeIndex)>,
88}
89
90impl<V> Default for SingleCache<V> {
91    fn default() -> Self {
92        SingleCache { cache: OnceLock::new() }
93    }
94}
95
96impl<V> QueryCache for SingleCache<V>
97where
98    V: Copy,
99{
100    type Key = ();
101    type Value = V;
102
103    #[inline(always)]
104    fn lookup(&self, _key: &()) -> Option<(V, DepNodeIndex)> {
105        self.cache.get().copied()
106    }
107
108    #[inline]
109    fn complete(&self, _key: (), value: V, index: DepNodeIndex) {
110        self.cache.set((value, index)).ok();
111    }
112
113    fn for_each(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
114        if let Some(value) = self.cache.get() {
115            f(&(), &value.0, value.1)
116        }
117    }
118
119    fn len(&self) -> usize {
120        self.cache.get().is_some().into()
121    }
122}
123
124/// In-memory cache for queries whose key is a [`DefId`].
125///
126/// Selects between one of two internal caches, depending on whether the key
127/// is a local ID or foreign-crate ID.
128pub struct DefIdCache<V> {
129    /// Stores the local DefIds in a dense map. Local queries are much more often dense, so this is
130    /// a win over hashing query keys at marginal memory cost (~5% at most) compared to FxHashMap.
131    local: VecCache<DefIndex, V, DepNodeIndex>,
132    foreign: DefaultCache<DefId, V>,
133}
134
135impl<V> Default for DefIdCache<V> {
136    fn default() -> Self {
137        DefIdCache { local: Default::default(), foreign: Default::default() }
138    }
139}
140
141impl<V> QueryCache for DefIdCache<V>
142where
143    V: Copy,
144{
145    type Key = DefId;
146    type Value = V;
147
148    #[inline(always)]
149    fn lookup(&self, key: &DefId) -> Option<(V, DepNodeIndex)> {
150        if key.krate == LOCAL_CRATE {
151            self.local.lookup(&key.index)
152        } else {
153            self.foreign.lookup(key)
154        }
155    }
156
157    #[inline]
158    fn complete(&self, key: DefId, value: V, index: DepNodeIndex) {
159        if key.krate == LOCAL_CRATE {
160            self.local.complete(key.index, value, index)
161        } else {
162            self.foreign.complete(key, value, index)
163        }
164    }
165
166    fn for_each(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
167        self.local.for_each(&mut |key, value, index| {
168            f(&DefId { krate: LOCAL_CRATE, index: *key }, value, index);
169        });
170        self.foreign.for_each(f);
171    }
172
173    fn len(&self) -> usize {
174        self.local.len() + self.foreign.len()
175    }
176}
177
178impl<K, V> QueryCache for VecCache<K, V, DepNodeIndex>
179where
180    K: Idx + QueryKey,
181    V: Copy,
182{
183    type Key = K;
184    type Value = V;
185
186    #[inline(always)]
187    fn lookup(&self, key: &K) -> Option<(V, DepNodeIndex)> {
188        self.lookup(key)
189    }
190
191    #[inline]
192    fn complete(&self, key: K, value: V, index: DepNodeIndex) {
193        self.complete(key, value, index)
194    }
195
196    fn for_each(&self, f: &mut dyn FnMut(&Self::Key, &Self::Value, DepNodeIndex)) {
197        self.for_each(f)
198    }
199
200    fn len(&self) -> usize {
201        self.len()
202    }
203}
rustc_middle/query/caches.rs

rustc_middle/query/
caches.rs
