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!(query = Q::name(&self.query), "cache hit");
199 return Ok(entry.value.clone());
200 }
201 }
202
203 let query_id = self.make_query_id(key);
205
206 if db.is_in_progress(query_id) {
208 return Err(QueryError::Cycle(vec![query_id]));
209 }
210
211 db.mark_in_progress(query_id);
212 let value = self.query.compute(db, key);
213 db.unmark_in_progress(query_id);
214
215 let current_rev = db.current_revision();
217 self.results.insert(
218 key.clone(),
219 MemoizedResult {
220 value: value.clone(),
221 computed_at: current_rev,
222 verified_at: current_rev,
223 dependencies: Vec::new(),
224 },
225 );
226
227 tracing::trace!(query = Q::name(&self.query), "computed and cached");
228
229 Ok(value)
230 }
231
232 pub fn invalidate_all(&self) {
234 self.results.clear();
235 }
236
237 pub fn invalidate(&self, key: &Q::Key) {
239 self.results.remove(key);
240 }
241
242 fn make_query_id(&self, key: &Q::Key) -> QueryId {
243 use std::hash::Hasher;
244 let mut hasher = FxHasher::default();
245 key.hash(&mut hasher);
246 QueryId {
247 query_type: TypeId::of::<Q>(),
248 key_hash: hasher.finish(),
249 }
250 }
251}
252
253pub struct SimpleDatabase {
255 revision: AtomicU64,
256 in_progress: DashMap<QueryId, (), BuildHasherDefault<FxHasher>>,
257 dependencies: DashMap<QueryId, Vec<QueryId>, BuildHasherDefault<FxHasher>>,
258}
259
260impl SimpleDatabase {
261 #[must_use]
263 pub fn new() -> Self {
264 Self {
265 revision: AtomicU64::new(0),
266 in_progress: DashMap::with_hasher(BuildHasherDefault::default()),
267 dependencies: DashMap::with_hasher(BuildHasherDefault::default()),
268 }
269 }
270}
271
272impl Default for SimpleDatabase {
273 fn default() -> Self {
274 Self::new()
275 }
276}
277
278impl QueryDatabase for SimpleDatabase {
279 fn current_revision(&self) -> Revision {
280 Revision(self.revision.load(Ordering::Acquire))
281 }
282
283 fn increment_revision(&self) -> Revision {
284 Revision(self.revision.fetch_add(1, Ordering::AcqRel) + 1)
285 }
286
287 fn mark_in_progress(&self, id: QueryId) {
288 self.in_progress.insert(id, ());
289 }
290
291 fn unmark_in_progress(&self, id: QueryId) {
292 self.in_progress.remove(&id);
293 }
294
295 fn is_in_progress(&self, id: QueryId) -> bool {
296 self.in_progress.contains_key(&id)
297 }
298
299 fn record_dependency(&self, from: QueryId, to: QueryId) {
300 self.dependencies
301 .entry(from)
302 .or_insert_with(Vec::new)
303 .push(to);
304 }
305}
306
307pub struct QueryRuntime<DB> {
309 pub db: DB,
311}
312
313impl<DB: QueryDatabase> QueryRuntime<DB> {
314 pub fn new(db: DB) -> Self {
316 Self { db }
317 }
318
319 pub fn query<Q: Query>(
321 &self,
322 storage: &QueryStorage<Q>,
323 key: &Q::Key,
324 ) -> QueryResult<Q::Value> {
325 storage.get(&self.db, key)
326 }
327}
328
329pub trait InputQuery: Query + Sized {
331 fn set(storage: &QueryStorage<Self>, key: Self::Key, value: Self::Value);
333}
334
335#[cfg(test)]
336mod tests {
337 use super::*;
338
339 struct TestQuery;
340
341 impl Query for TestQuery {
342 type Key = u32;
343 type Value = u64;
344
345 fn compute(&self, _db: &dyn QueryDatabase, key: &u32) -> u64 {
346 (*key as u64) * 2
347 }
348
349 fn name(&self) -> &'static str {
350 "test_query"
351 }
352 }
353
354 #[test]
355 fn test_basic_query() {
356 let db = SimpleDatabase::new();
357 let storage = QueryStorage::new(TestQuery);
358 let runtime = QueryRuntime::new(db);
359
360 let result = runtime.query(&storage, &21).unwrap();
361 assert_eq!(result, 42);
362
363 let result2 = runtime.query(&storage, &21).unwrap();
365 assert_eq!(result2, 42);
366 }
367
368 #[test]
369 fn test_revision_tracking() {
370 let db = SimpleDatabase::new();
371 assert_eq!(db.current_revision(), Revision::INITIAL);
372
373 let rev1 = db.increment_revision();
374 assert_eq!(rev1, Revision::new(1));
375
376 let rev2 = db.increment_revision();
377 assert_eq!(rev2, Revision::new(2));
378 }
379}