Track cycles in volatile queries

src/lib.rs

@@ -368,7 +368,7 @@ macro_rules! query_definition {
     (
         @storage_ty[$QC:ident, $Self:ident, volatile]
     ) => {
-        $crate::volatile::VolatileStorage
+        $crate::volatile::VolatileStorage<$QC, $Self>
     };
 
     // Accept a "field-like" query definition (input)

src/volatile.rs

@@ -5,8 +5,8 @@ use crate::QueryContext;
 use crate::QueryStorageOps;
 use crate::QueryTable;
 use log::debug;
-use parking_lot::{RwLock, RwLockUpgradableReadGuard};
-use rustc_hash::FxHashMap;
+use parking_lot::Mutex;
+use rustc_hash::FxHashSet;
 use std::any::Any;
 use std::cell::RefCell;
 use std::collections::hash_map::Entry;
@@ -17,10 +17,29 @@ use std::hash::Hash;
 
 /// Volatile Storage is just **always** considered dirty. Any time you
 /// ask for the result of such a query, it is recomputed.
-#[derive(Default)]
-pub struct VolatileStorage;
+pub struct VolatileStorage<QC, Q>
+where
+    Q: Query<QC>,
+    QC: QueryContext,
+{
+    /// We don't store the results of volatile queries,
+    /// but we track in-progress set to detect cycles.
+    in_progress: Mutex<FxHashSet<Q::Key>>,
+}
 
-impl<QC, Q> QueryStorageOps<QC, Q> for VolatileStorage
+impl<QC, Q> Default for VolatileStorage<QC, Q>
+where
+    Q: Query<QC>,
+    QC: QueryContext,
+{
+    fn default() -> Self {
+        VolatileStorage {
+            in_progress: Mutex::new(FxHashSet::default()),
+        }
+    }
+}
+
+impl<QC, Q> QueryStorageOps<QC, Q> for VolatileStorage<QC, Q>
 where
     Q: Query<QC>,
     QC: QueryContext,
@@ -31,12 +50,18 @@ where
         key: &Q::Key,
         descriptor: &QC::QueryDescriptor,
     ) -> Result<Q::Value, CycleDetected> {
+        if !self.in_progress.lock().insert(key.clone()) {
+            return Err(CycleDetected);
+        }
+
         // FIXME: Should we even call `execute_query_implementation`
         // here? Or should we just call `Q::execute`, and maybe
         // separate out the `push`/`pop` operations.
         let (value, _inputs) = query
             .salsa_runtime()
             .execute_query_implementation::<Q>(query, descriptor, key);
+        let was_in_progress = self.in_progress.lock().remove(key);
+        assert!(was_in_progress);
 
         query.salsa_runtime().report_query_read(descriptor);