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!(query = Q::name(&self.query), "cache hit");
199                return Ok(entry.value.clone());
200            }
201        }
202
203        // Compute the result
204        let query_id = self.make_query_id(key);
205
206        // Check for cycles
207        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        // Memoize the result
216        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    /// Invalidate all cached results for this query.
233    pub fn invalidate_all(&self) {
234        self.results.clear();
235    }
236
237    /// Invalidate a specific key.
238    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
253/// A simple in-memory query database implementation.
254pub 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    /// Create a new empty database.
262    #[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
307/// A runtime for executing queries with a database.
308pub struct QueryRuntime<DB> {
309    /// The underlying database.
310    pub db: DB,
311}
312
313impl<DB: QueryDatabase> QueryRuntime<DB> {
314    /// Create a new runtime with the given database.
315    pub fn new(db: DB) -> Self {
316        Self { db }
317    }
318
319    /// Execute a query and return its result.
320    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
329/// Marker trait for input queries that can be set directly.
330pub trait InputQuery: Query + Sized {
331    /// Set the value for an input query.
332    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        // Second call should be cached
364        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}