]>
Commit | Line | Data |
---|---|---|
74b04a01 | 1 | use crate::dep_graph::DepNodeIndex; |
ba9703b0 XL |
2 | use crate::query::plumbing::{QueryLookup, QueryState}; |
3 | use crate::query::QueryContext; | |
74b04a01 | 4 | |
f035d41b | 5 | use rustc_arena::TypedArena; |
74b04a01 XL |
6 | use rustc_data_structures::fx::FxHashMap; |
7 | use rustc_data_structures::sharded::Sharded; | |
f9f354fc | 8 | use rustc_data_structures::sync::WorkerLocal; |
74b04a01 XL |
9 | use std::default::Default; |
10 | use std::hash::Hash; | |
ba9703b0 | 11 | use std::marker::PhantomData; |
74b04a01 | 12 | |
f9f354fc XL |
13 | pub trait CacheSelector<K, V> { |
14 | type Cache; | |
74b04a01 XL |
15 | } |
16 | ||
f9f354fc | 17 | pub trait QueryStorage: Default { |
ba9703b0 | 18 | type Value; |
f9f354fc XL |
19 | type Stored: Clone; |
20 | ||
21 | /// Store a value without putting it in the cache. | |
22 | /// This is meant to be used with cycle errors. | |
23 | fn store_nocache(&self, value: Self::Value) -> Self::Stored; | |
24 | } | |
25 | ||
26 | pub trait QueryCache: QueryStorage { | |
27 | type Key: Hash; | |
74b04a01 XL |
28 | type Sharded: Default; |
29 | ||
30 | /// Checks if the query is already computed and in the cache. | |
31 | /// It returns the shard index and a lock guard to the shard, | |
32 | /// which will be used if the query is not in the cache and we need | |
33 | /// to compute it. | |
ba9703b0 | 34 | fn lookup<CTX: QueryContext, R, OnHit, OnMiss>( |
74b04a01 | 35 | &self, |
ba9703b0 XL |
36 | state: &QueryState<CTX, Self>, |
37 | key: Self::Key, | |
74b04a01 XL |
38 | // `on_hit` can be called while holding a lock to the query state shard. |
39 | on_hit: OnHit, | |
40 | on_miss: OnMiss, | |
41 | ) -> R | |
42 | where | |
f9f354fc | 43 | OnHit: FnOnce(&Self::Stored, DepNodeIndex) -> R, |
ba9703b0 | 44 | OnMiss: FnOnce(Self::Key, QueryLookup<'_, CTX, Self::Key, Self::Sharded>) -> R; |
74b04a01 | 45 | |
3dfed10e | 46 | fn complete( |
74b04a01 | 47 | &self, |
74b04a01 | 48 | lock_sharded_storage: &mut Self::Sharded, |
ba9703b0 XL |
49 | key: Self::Key, |
50 | value: Self::Value, | |
74b04a01 | 51 | index: DepNodeIndex, |
f9f354fc | 52 | ) -> Self::Stored; |
74b04a01 XL |
53 | |
54 | fn iter<R, L>( | |
55 | &self, | |
56 | shards: &Sharded<L>, | |
57 | get_shard: impl Fn(&mut L) -> &mut Self::Sharded, | |
ba9703b0 XL |
58 | f: impl for<'a> FnOnce( |
59 | Box<dyn Iterator<Item = (&'a Self::Key, &'a Self::Value, DepNodeIndex)> + 'a>, | |
60 | ) -> R, | |
74b04a01 XL |
61 | ) -> R; |
62 | } | |
63 | ||
64 | pub struct DefaultCacheSelector; | |
65 | ||
66 | impl<K: Eq + Hash, V: Clone> CacheSelector<K, V> for DefaultCacheSelector { | |
ba9703b0 | 67 | type Cache = DefaultCache<K, V>; |
74b04a01 XL |
68 | } |
69 | ||
ba9703b0 XL |
70 | pub struct DefaultCache<K, V>(PhantomData<(K, V)>); |
71 | ||
72 | impl<K, V> Default for DefaultCache<K, V> { | |
73 | fn default() -> Self { | |
74 | DefaultCache(PhantomData) | |
75 | } | |
76 | } | |
74b04a01 | 77 | |
f9f354fc XL |
78 | impl<K: Eq + Hash, V: Clone> QueryStorage for DefaultCache<K, V> { |
79 | type Value = V; | |
80 | type Stored = V; | |
81 | ||
82 | #[inline] | |
83 | fn store_nocache(&self, value: Self::Value) -> Self::Stored { | |
84 | // We have no dedicated storage | |
85 | value | |
86 | } | |
87 | } | |
88 | ||
ba9703b0 XL |
89 | impl<K: Eq + Hash, V: Clone> QueryCache for DefaultCache<K, V> { |
90 | type Key = K; | |
74b04a01 XL |
91 | type Sharded = FxHashMap<K, (V, DepNodeIndex)>; |
92 | ||
93 | #[inline(always)] | |
ba9703b0 | 94 | fn lookup<CTX: QueryContext, R, OnHit, OnMiss>( |
74b04a01 | 95 | &self, |
ba9703b0 | 96 | state: &QueryState<CTX, Self>, |
74b04a01 XL |
97 | key: K, |
98 | on_hit: OnHit, | |
99 | on_miss: OnMiss, | |
100 | ) -> R | |
101 | where | |
74b04a01 | 102 | OnHit: FnOnce(&V, DepNodeIndex) -> R, |
ba9703b0 | 103 | OnMiss: FnOnce(K, QueryLookup<'_, CTX, K, Self::Sharded>) -> R, |
74b04a01 XL |
104 | { |
105 | let mut lookup = state.get_lookup(&key); | |
106 | let lock = &mut *lookup.lock; | |
107 | ||
ba9703b0 | 108 | let result = lock.cache.raw_entry().from_key_hashed_nocheck(lookup.key_hash, &key); |
74b04a01 XL |
109 | |
110 | if let Some((_, value)) = result { on_hit(&value.0, value.1) } else { on_miss(key, lookup) } | |
111 | } | |
112 | ||
113 | #[inline] | |
3dfed10e | 114 | fn complete( |
74b04a01 | 115 | &self, |
74b04a01 XL |
116 | lock_sharded_storage: &mut Self::Sharded, |
117 | key: K, | |
118 | value: V, | |
119 | index: DepNodeIndex, | |
f9f354fc XL |
120 | ) -> Self::Stored { |
121 | lock_sharded_storage.insert(key, (value.clone(), index)); | |
122 | value | |
123 | } | |
124 | ||
125 | fn iter<R, L>( | |
126 | &self, | |
127 | shards: &Sharded<L>, | |
128 | get_shard: impl Fn(&mut L) -> &mut Self::Sharded, | |
129 | f: impl for<'a> FnOnce(Box<dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)> + 'a>) -> R, | |
130 | ) -> R { | |
131 | let mut shards = shards.lock_shards(); | |
132 | let mut shards: Vec<_> = shards.iter_mut().map(|shard| get_shard(shard)).collect(); | |
133 | let results = shards.iter_mut().flat_map(|shard| shard.iter()).map(|(k, v)| (k, &v.0, v.1)); | |
134 | f(Box::new(results)) | |
135 | } | |
136 | } | |
137 | ||
138 | pub struct ArenaCacheSelector<'tcx>(PhantomData<&'tcx ()>); | |
139 | ||
140 | impl<'tcx, K: Eq + Hash, V: 'tcx> CacheSelector<K, V> for ArenaCacheSelector<'tcx> { | |
141 | type Cache = ArenaCache<'tcx, K, V>; | |
142 | } | |
143 | ||
144 | pub struct ArenaCache<'tcx, K, V> { | |
145 | arena: WorkerLocal<TypedArena<(V, DepNodeIndex)>>, | |
146 | phantom: PhantomData<(K, &'tcx V)>, | |
147 | } | |
148 | ||
149 | impl<'tcx, K, V> Default for ArenaCache<'tcx, K, V> { | |
150 | fn default() -> Self { | |
151 | ArenaCache { arena: WorkerLocal::new(|_| TypedArena::default()), phantom: PhantomData } | |
152 | } | |
153 | } | |
154 | ||
155 | impl<'tcx, K: Eq + Hash, V: 'tcx> QueryStorage for ArenaCache<'tcx, K, V> { | |
156 | type Value = V; | |
157 | type Stored = &'tcx V; | |
158 | ||
159 | #[inline] | |
160 | fn store_nocache(&self, value: Self::Value) -> Self::Stored { | |
161 | let value = self.arena.alloc((value, DepNodeIndex::INVALID)); | |
162 | let value = unsafe { &*(&value.0 as *const _) }; | |
163 | &value | |
164 | } | |
165 | } | |
166 | ||
167 | impl<'tcx, K: Eq + Hash, V: 'tcx> QueryCache for ArenaCache<'tcx, K, V> { | |
168 | type Key = K; | |
169 | type Sharded = FxHashMap<K, &'tcx (V, DepNodeIndex)>; | |
170 | ||
171 | #[inline(always)] | |
172 | fn lookup<CTX: QueryContext, R, OnHit, OnMiss>( | |
173 | &self, | |
174 | state: &QueryState<CTX, Self>, | |
175 | key: K, | |
176 | on_hit: OnHit, | |
177 | on_miss: OnMiss, | |
178 | ) -> R | |
179 | where | |
180 | OnHit: FnOnce(&&'tcx V, DepNodeIndex) -> R, | |
181 | OnMiss: FnOnce(K, QueryLookup<'_, CTX, K, Self::Sharded>) -> R, | |
182 | { | |
183 | let mut lookup = state.get_lookup(&key); | |
184 | let lock = &mut *lookup.lock; | |
185 | ||
186 | let result = lock.cache.raw_entry().from_key_hashed_nocheck(lookup.key_hash, &key); | |
187 | ||
188 | if let Some((_, value)) = result { | |
189 | on_hit(&&value.0, value.1) | |
190 | } else { | |
191 | on_miss(key, lookup) | |
192 | } | |
193 | } | |
194 | ||
195 | #[inline] | |
3dfed10e | 196 | fn complete( |
f9f354fc | 197 | &self, |
f9f354fc XL |
198 | lock_sharded_storage: &mut Self::Sharded, |
199 | key: K, | |
200 | value: V, | |
201 | index: DepNodeIndex, | |
202 | ) -> Self::Stored { | |
203 | let value = self.arena.alloc((value, index)); | |
204 | let value = unsafe { &*(value as *const _) }; | |
205 | lock_sharded_storage.insert(key, value); | |
206 | &value.0 | |
74b04a01 XL |
207 | } |
208 | ||
209 | fn iter<R, L>( | |
210 | &self, | |
211 | shards: &Sharded<L>, | |
212 | get_shard: impl Fn(&mut L) -> &mut Self::Sharded, | |
213 | f: impl for<'a> FnOnce(Box<dyn Iterator<Item = (&'a K, &'a V, DepNodeIndex)> + 'a>) -> R, | |
214 | ) -> R { | |
215 | let mut shards = shards.lock_shards(); | |
216 | let mut shards: Vec<_> = shards.iter_mut().map(|shard| get_shard(shard)).collect(); | |
217 | let results = shards.iter_mut().flat_map(|shard| shard.iter()).map(|(k, v)| (k, &v.0, v.1)); | |
218 | f(Box::new(results)) | |
219 | } | |
220 | } |