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 rustc_hash::FxHasher;
36use std::any::TypeId;
37use std::fmt::Debug;
38use std::hash::{BuildHasherDefault, Hash};
39use std::sync::atomic::{AtomicU64, Ordering};
40
41/// A revision number tracking database state changes.
42#[derive(Clone, Copy, Debug, Default, PartialEq, Eq, PartialOrd, Ord, Hash)]
43pub struct Revision(u64);
44
45impl Revision {
46    /// The initial revision.
47    pub const INITIAL: Self = Self(0);
48
49    /// Create a new revision with the given value.
50    #[must_use]
51    pub const fn new(value: u64) -> Self {
52        Self(value)
53    }
54
55    /// Get the next revision.
56    #[must_use]
57    pub const fn next(self) -> Self {
58        Self(self.0 + 1)
59    }
60
61    /// Get the raw revision value.
62    #[must_use]
63    pub const fn as_u64(self) -> u64 {
64        self.0
65    }
66}
67
68/// A unique identifier for a query invocation.
69#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
70pub struct QueryId {
71    /// The type of query.
72    pub query_type: TypeId,
73    /// A hash of the query key.
74    pub key_hash: u64,
75}
76
77/// The state of a query computation.
78#[derive(Clone, Copy, Debug, PartialEq, Eq)]
79pub enum QueryState {
80    /// Query has not been computed yet.
81    NotComputed,
82    /// Query is currently being computed (cycle detection).
83    InProgress,
84    /// Query has been computed and memoized.
85    Memoized(Revision),
86    /// Query result is stale and needs recomputation.
87    Stale,
88}
89
90/// Error types for the query system.
91#[derive(Debug, Clone)]
92pub enum QueryError {
93    /// A cycle was detected in query dependencies.
94    Cycle(Vec<QueryId>),
95    /// The query panicked during computation.
96    Panic(String),
97    /// The query was cancelled.
98    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
113/// Result type for query computations.
114pub type QueryResult<T> = Result<T, QueryError>;
115
116/// A memoized query result with revision tracking.
117#[derive(Debug)]
118pub struct MemoizedResult<V> {
119    /// The cached value.
120    pub value: V,
121    /// The revision when this was computed.
122    pub computed_at: Revision,
123    /// The revision when this was last verified.
124    pub verified_at: Revision,
125    /// Dependencies of this query (for invalidation).
126    pub dependencies: Vec<QueryId>,
127}
128
129/// Trait for types that can be used as query keys.
130pub trait QueryKey: Clone + Eq + Hash + Debug + Send + Sync + 'static {}
131
132impl<T: Clone + Eq + Hash + Debug + Send + Sync + 'static> QueryKey for T {}
133
134/// Trait for types that can be used as query values.
135pub trait QueryValue: Clone + Debug + Send + Sync + 'static {}
136
137impl<T: Clone + Debug + Send + Sync + 'static> QueryValue for T {}
138
139/// A query definition with its computation function.
140pub trait Query: Send + Sync + 'static {
141    /// The key type for this query.
142    type Key: QueryKey;
143    /// The value type for this query.
144    type Value: QueryValue;
145
146    /// Compute the query result for the given key.
147    fn compute(&self, db: &dyn QueryDatabase, key: &Self::Key) -> Self::Value;
148
149    /// Get a human-readable name for this query.
150    fn name(&self) -> &'static str;
151}
152
153/// The central database trait that all query databases must implement.
154pub trait QueryDatabase: Send + Sync {
155    /// Get the current revision of the database.
156    fn current_revision(&self) -> Revision;
157
158    /// Increment the revision (called when inputs change).
159    fn increment_revision(&self) -> Revision;
160
161    /// Mark a query as in-progress (for cycle detection).
162    fn mark_in_progress(&self, id: QueryId);
163
164    /// Unmark a query as in-progress.
165    fn unmark_in_progress(&self, id: QueryId);
166
167    /// Check if a query is in progress (cycle detection).
168    fn is_in_progress(&self, id: QueryId) -> bool;
169
170    /// Record a dependency from the current query to another.
171    fn record_dependency(&self, from: QueryId, to: QueryId);
172}
173
174/// Storage for a specific query type.
175pub struct QueryStorage<Q: Query> {
176    /// The query implementation.
177    query: Q,
178    /// Memoized results keyed by query key.
179    results: DashMap<Q::Key, MemoizedResult<Q::Value>, BuildHasherDefault<FxHasher>>,
180}
181
182impl<Q: Query> QueryStorage<Q> {
183    /// Create new storage for a query.
184    pub fn new(query: Q) -> Self {
185        Self {
186            query,
187            results: DashMap::with_hasher(BuildHasherDefault::default()),
188        }
189    }
190
191    /// Execute the query, returning a memoized result if available.
192    pub fn get(&self, db: &dyn QueryDatabase, key: &Q::Key) -> QueryResult<Q::Value> {
193        // Check for cached result
194        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        // Compute the result
202        let query_id = self.make_query_id(key);
203
204        // Check for cycles
205        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        // Memoize the result
214        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    /// Invalidate all cached results for this query.
231    pub fn invalidate_all(&self) {
232        self.results.clear();
233    }
234
235    /// Invalidate a specific key.
236    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
251/// A simple in-memory query database implementation.
252pub 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    /// Create a new empty database.
260    #[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
302/// A runtime for executing queries with a database.
303pub struct QueryRuntime<DB> {
304    /// The underlying database.
305    pub db: DB,
306}
307
308impl<DB: QueryDatabase> QueryRuntime<DB> {
309    /// Create a new runtime with the given database.
310    pub fn new(db: DB) -> Self {
311        Self { db }
312    }
313
314    /// Execute a query and return its result.
315    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
324/// Marker trait for input queries that can be set directly.
325pub trait InputQuery: Query + Sized {
326    /// Set the value for an input query.
327    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        // Second call should be cached
359        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}