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