1#![warn(missing_docs)]
33
34use dashmap::DashMap;
35use parking_lot::{Mutex, RwLock};
36use rustc_hash::FxHasher;
37use std::any::{Any, TypeId};
38use std::fmt::Debug;
39use std::hash::{BuildHasherDefault, Hash};
40use std::sync::atomic::{AtomicU64, Ordering};
41use std::sync::Arc;
42
43#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
45pub struct Revision(u64);
46
47impl Revision {
48 pub const INITIAL: Self = Self(0);
50
51 #[must_use]
53 pub const fn new(value: u64) -> Self {
54 Self(value)
55 }
56
57 #[must_use]
59 pub const fn next(self) -> Self {
60 Self(self.0 + 1)
61 }
62
63 #[must_use]
65 pub const fn as_u64(self) -> u64 {
66 self.0
67 }
68}
69
70#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
72pub struct QueryId {
73 pub query_type: TypeId,
75 pub key_hash: u64,
77}
78
79#[derive(Clone, Copy, Debug, PartialEq, Eq)]
81pub enum QueryState {
82 NotComputed,
84 InProgress,
86 Memoized(Revision),
88 Stale,
90}
91
92#[derive(Debug, Clone)]
94pub enum QueryError {
95 Cycle(Vec<QueryId>),
97 Panic(String),
99 Cancelled,
101}
102
103impl std::fmt::Display for QueryError {
104 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
105 match self {
106 Self::Cycle(ids) => write!(f, "query cycle detected: {} queries involved", ids.len()),
107 Self::Panic(msg) => write!(f, "query panicked: {msg}"),
108 Self::Cancelled => write!(f, "query was cancelled"),
109 }
110 }
111}
112
113impl std::error::Error for QueryError {}
114
115pub type QueryResult<T> = Result<T, QueryError>;
117
118#[derive(Debug)]
120pub struct MemoizedResult<V> {
121 pub value: V,
123 pub computed_at: Revision,
125 pub verified_at: Revision,
127 pub dependencies: Vec<QueryId>,
129}
130
131pub trait QueryKey: Clone + Eq + Hash + Debug + Send + Sync + 'static {}
133
134impl<T: Clone + Eq + Hash + Debug + Send + Sync + 'static> QueryKey for T {}
135
136pub trait QueryValue: Clone + Debug + Send + Sync + 'static {}
138
139impl<T: Clone + Debug + Send + Sync + 'static> QueryValue for T {}
140
141pub trait Query: Send + Sync + 'static {
143 type Key: QueryKey;
145 type Value: QueryValue;
147
148 fn compute(&self, db: &dyn QueryDatabase, key: &Self::Key) -> Self::Value;
150
151 fn name(&self) -> &'static str;
153}
154
155pub trait QueryDatabase: Send + Sync {
157 fn current_revision(&self) -> Revision;
159
160 fn increment_revision(&self) -> Revision;
162
163 fn mark_in_progress(&self, id: QueryId);
165
166 fn unmark_in_progress(&self, id: QueryId);
168
169 fn is_in_progress(&self, id: QueryId) -> bool;
171
172 fn record_dependency(&self, from: QueryId, to: QueryId);
174}
175
176pub struct QueryStorage<Q: Query> {
178 query: Q,
180 results: DashMap<Q::Key, MemoizedResult<Q::Value>, BuildHasherDefault<FxHasher>>,
182}
183
184impl<Q: Query> QueryStorage<Q> {
185 pub fn new(query: Q) -> Self {
187 Self {
188 query,
189 results: DashMap::with_hasher(BuildHasherDefault::default()),
190 }
191 }
192
193 pub fn get(&self, db: &dyn QueryDatabase, key: &Q::Key) -> QueryResult<Q::Value> {
195 if let Some(entry) = self.results.get(key) {
197 if entry.computed_at >= db.current_revision() {
198 tracing::trace!(
199 query = Q::name(&self.query),
200 "cache hit"
201 );
202 return Ok(entry.value.clone());
203 }
204 }
205
206 let query_id = self.make_query_id(key);
208
209 if db.is_in_progress(query_id) {
211 return Err(QueryError::Cycle(vec![query_id]));
212 }
213
214 db.mark_in_progress(query_id);
215 let value = self.query.compute(db, key);
216 db.unmark_in_progress(query_id);
217
218 let current_rev = db.current_revision();
220 self.results.insert(
221 key.clone(),
222 MemoizedResult {
223 value: value.clone(),
224 computed_at: current_rev,
225 verified_at: current_rev,
226 dependencies: Vec::new(),
227 },
228 );
229
230 tracing::trace!(
231 query = Q::name(&self.query),
232 "computed and cached"
233 );
234
235 Ok(value)
236 }
237
238 pub fn invalidate_all(&self) {
240 self.results.clear();
241 }
242
243 pub fn invalidate(&self, key: &Q::Key) {
245 self.results.remove(key);
246 }
247
248 fn make_query_id(&self, key: &Q::Key) -> QueryId {
249 use std::hash::Hasher;
250 let mut hasher = FxHasher::default();
251 key.hash(&mut hasher);
252 QueryId {
253 query_type: TypeId::of::<Q>(),
254 key_hash: hasher.finish(),
255 }
256 }
257}
258
259pub struct SimpleDatabase {
261 revision: AtomicU64,
262 in_progress: DashMap<QueryId, (), BuildHasherDefault<FxHasher>>,
263 dependencies: DashMap<QueryId, Vec<QueryId>, BuildHasherDefault<FxHasher>>,
264}
265
266impl SimpleDatabase {
267 #[must_use]
269 pub fn new() -> Self {
270 Self {
271 revision: AtomicU64::new(0),
272 in_progress: DashMap::with_hasher(BuildHasherDefault::default()),
273 dependencies: DashMap::with_hasher(BuildHasherDefault::default()),
274 }
275 }
276}
277
278impl Default for SimpleDatabase {
279 fn default() -> Self {
280 Self::new()
281 }
282}
283
284impl QueryDatabase for SimpleDatabase {
285 fn current_revision(&self) -> Revision {
286 Revision(self.revision.load(Ordering::Acquire))
287 }
288
289 fn increment_revision(&self) -> Revision {
290 Revision(self.revision.fetch_add(1, Ordering::AcqRel) + 1)
291 }
292
293 fn mark_in_progress(&self, id: QueryId) {
294 self.in_progress.insert(id, ());
295 }
296
297 fn unmark_in_progress(&self, id: QueryId) {
298 self.in_progress.remove(&id);
299 }
300
301 fn is_in_progress(&self, id: QueryId) -> bool {
302 self.in_progress.contains_key(&id)
303 }
304
305 fn record_dependency(&self, from: QueryId, to: QueryId) {
306 self.dependencies
307 .entry(from)
308 .or_insert_with(Vec::new)
309 .push(to);
310 }
311}
312
313pub struct QueryRuntime<DB> {
315 pub db: DB,
317}
318
319impl<DB: QueryDatabase> QueryRuntime<DB> {
320 pub fn new(db: DB) -> Self {
322 Self { db }
323 }
324
325 pub fn query<Q: Query>(&self, storage: &QueryStorage<Q>, key: &Q::Key) -> QueryResult<Q::Value> {
327 storage.get(&self.db, key)
328 }
329}
330
331pub trait InputQuery: Query + Sized {
333 fn set(storage: &QueryStorage<Self>, key: Self::Key, value: Self::Value);
335}
336
337#[cfg(test)]
338mod tests {
339 use super::*;
340
341 struct TestQuery;
342
343 impl Query for TestQuery {
344 type Key = u32;
345 type Value = u64;
346
347 fn compute(&self, _db: &dyn QueryDatabase, key: &u32) -> u64 {
348 (*key as u64) * 2
349 }
350
351 fn name(&self) -> &'static str {
352 "test_query"
353 }
354 }
355
356 #[test]
357 fn test_basic_query() {
358 let db = SimpleDatabase::new();
359 let storage = QueryStorage::new(TestQuery);
360 let runtime = QueryRuntime::new(db);
361
362 let result = runtime.query(&storage, &21).unwrap();
363 assert_eq!(result, 42);
364
365 let result2 = runtime.query(&storage, &21).unwrap();
367 assert_eq!(result2, 42);
368 }
369
370 #[test]
371 fn test_revision_tracking() {
372 let db = SimpleDatabase::new();
373 assert_eq!(db.current_revision(), Revision::INITIAL);
374
375 let rev1 = db.increment_revision();
376 assert_eq!(rev1, Revision::new(1));
377
378 let rev2 = db.increment_revision();
379 assert_eq!(rev2, Revision::new(2));
380 }
381}