Skip to main content

bhc_query/
lib.rs

1//! Query-based compilation system for BHC.
2//!
3//! This crate provides the infrastructure for incremental, demand-driven
4//! compilation using a query-based architecture similar to rustc's query
5//! system and salsa.
6//!
7//! # Overview
8//!
9//! The query system allows the compiler to:
10//!
11//! - **Incrementally recompute** only what has changed between compilations
12//! - **Demand-driven evaluation** where queries are only computed when needed
13//! - **Automatic memoization** of query results
14//! - **Cycle detection** for recursive queries
15//!
16//! # Architecture
17//!
18//! Queries are defined as traits that can be implemented by a database.
19//! The database stores memoized results and tracks dependencies between queries.
20//!
21//! ```text
22//! ┌─────────────────────────────────────────────────────┐
23//! │                    QueryDatabase                     │
24//! │  ┌─────────────┐  ┌─────────────┐  ┌─────────────┐ │
25//! │  │  Query A    │  │  Query B    │  │  Query C    │ │
26//! │  │  (cached)   │──│  (cached)   │──│  (pending)  │ │
27//! │  └─────────────┘  └─────────────┘  └─────────────┘ │
28//! │                    Dependencies                      │
29//! └─────────────────────────────────────────────────────┘
30//! ```
31
32#![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/// A revision number tracking database state changes.
44#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
45pub struct Revision(u64);
46
47impl Revision {
48    /// The initial revision.
49    pub const INITIAL: Self = Self(0);
50
51    /// Create a new revision with the given value.
52    #[must_use]
53    pub const fn new(value: u64) -> Self {
54        Self(value)
55    }
56
57    /// Get the next revision.
58    #[must_use]
59    pub const fn next(self) -> Self {
60        Self(self.0 + 1)
61    }
62
63    /// Get the raw revision value.
64    #[must_use]
65    pub const fn as_u64(self) -> u64 {
66        self.0
67    }
68}
69
70/// A unique identifier for a query invocation.
71#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
72pub struct QueryId {
73    /// The type of query.
74    pub query_type: TypeId,
75    /// A hash of the query key.
76    pub key_hash: u64,
77}
78
79/// The state of a query computation.
80#[derive(Clone, Copy, Debug, PartialEq, Eq)]
81pub enum QueryState {
82    /// Query has not been computed yet.
83    NotComputed,
84    /// Query is currently being computed (cycle detection).
85    InProgress,
86    /// Query has been computed and memoized.
87    Memoized(Revision),
88    /// Query result is stale and needs recomputation.
89    Stale,
90}
91
92/// Error types for the query system.
93#[derive(Debug, Clone)]
94pub enum QueryError {
95    /// A cycle was detected in query dependencies.
96    Cycle(Vec<QueryId>),
97    /// The query panicked during computation.
98    Panic(String),
99    /// The query was cancelled.
100    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
115/// Result type for query computations.
116pub type QueryResult<T> = Result<T, QueryError>;
117
118/// A memoized query result with revision tracking.
119#[derive(Debug)]
120pub struct MemoizedResult<V> {
121    /// The cached value.
122    pub value: V,
123    /// The revision when this was computed.
124    pub computed_at: Revision,
125    /// The revision when this was last verified.
126    pub verified_at: Revision,
127    /// Dependencies of this query (for invalidation).
128    pub dependencies: Vec<QueryId>,
129}
130
131/// Trait for types that can be used as query keys.
132pub trait QueryKey: Clone + Eq + Hash + Debug + Send + Sync + 'static {}
133
134impl<T: Clone + Eq + Hash + Debug + Send + Sync + 'static> QueryKey for T {}
135
136/// Trait for types that can be used as query values.
137pub trait QueryValue: Clone + Debug + Send + Sync + 'static {}
138
139impl<T: Clone + Debug + Send + Sync + 'static> QueryValue for T {}
140
141/// A query definition with its computation function.
142pub trait Query: Send + Sync + 'static {
143    /// The key type for this query.
144    type Key: QueryKey;
145    /// The value type for this query.
146    type Value: QueryValue;
147
148    /// Compute the query result for the given key.
149    fn compute(&self, db: &dyn QueryDatabase, key: &Self::Key) -> Self::Value;
150
151    /// Get a human-readable name for this query.
152    fn name(&self) -> &'static str;
153}
154
155/// The central database trait that all query databases must implement.
156pub trait QueryDatabase: Send + Sync {
157    /// Get the current revision of the database.
158    fn current_revision(&self) -> Revision;
159
160    /// Increment the revision (called when inputs change).
161    fn increment_revision(&self) -> Revision;
162
163    /// Mark a query as in-progress (for cycle detection).
164    fn mark_in_progress(&self, id: QueryId);
165
166    /// Unmark a query as in-progress.
167    fn unmark_in_progress(&self, id: QueryId);
168
169    /// Check if a query is in progress (cycle detection).
170    fn is_in_progress(&self, id: QueryId) -> bool;
171
172    /// Record a dependency from the current query to another.
173    fn record_dependency(&self, from: QueryId, to: QueryId);
174}
175
176/// Storage for a specific query type.
177pub struct QueryStorage<Q: Query> {
178    /// The query implementation.
179    query: Q,
180    /// Memoized results keyed by query key.
181    results: DashMap<Q::Key, MemoizedResult<Q::Value>, BuildHasherDefault<FxHasher>>,
182}
183
184impl<Q: Query> QueryStorage<Q> {
185    /// Create new storage for a query.
186    pub fn new(query: Q) -> Self {
187        Self {
188            query,
189            results: DashMap::with_hasher(BuildHasherDefault::default()),
190        }
191    }
192
193    /// Execute the query, returning a memoized result if available.
194    pub fn get(&self, db: &dyn QueryDatabase, key: &Q::Key) -> QueryResult<Q::Value> {
195        // Check for cached result
196        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        // Compute the result
207        let query_id = self.make_query_id(key);
208
209        // Check for cycles
210        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        // Memoize the result
219        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    /// Invalidate all cached results for this query.
239    pub fn invalidate_all(&self) {
240        self.results.clear();
241    }
242
243    /// Invalidate a specific key.
244    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
259/// A simple in-memory query database implementation.
260pub 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    /// Create a new empty database.
268    #[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
313/// A runtime for executing queries with a database.
314pub struct QueryRuntime<DB> {
315    /// The underlying database.
316    pub db: DB,
317}
318
319impl<DB: QueryDatabase> QueryRuntime<DB> {
320    /// Create a new runtime with the given database.
321    pub fn new(db: DB) -> Self {
322        Self { db }
323    }
324
325    /// Execute a query and return its result.
326    pub fn query<Q: Query>(&self, storage: &QueryStorage<Q>, key: &Q::Key) -> QueryResult<Q::Value> {
327        storage.get(&self.db, key)
328    }
329}
330
331/// Marker trait for input queries that can be set directly.
332pub trait InputQuery: Query + Sized {
333    /// Set the value for an input query.
334    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        // Second call should be cached
366        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}