From e13187f747cabef31a72de45c2773c978a795464 Mon Sep 17 00:00:00 2001
From: Aleksey Kladov <aleksey.kladov@gmail.com>
Date: Mon, 1 Oct 2018 13:58:18 +0300
Subject: [PATCH] Track cycles in volatile queries

---
 src/lib.rs      |  2 +-
 src/volatile.rs | 35 ++++++++++++++++++++++++++++++-----
 2 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index 911be173..b3d226cc 100644
--- a/src/lib.rs
+++ b/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)
diff --git a/src/volatile.rs b/src/volatile.rs
index 911da65a..5312ffd4 100644
--- a/src/volatile.rs
+++ b/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);