add a set_unchecked method that can be used to do mocking in tests

src/input.rs

@@ -8,6 +8,7 @@ use crate::Query;
 use crate::QueryDescriptor;
 use crate::QueryStorageOps;
 use crate::QueryTable;
+use crate::UncheckedMutQueryStorageOps;
 use log::debug;
 use parking_lot::{RwLock, RwLockUpgradableReadGuard};
 use rustc_hash::FxHashMap;
@@ -135,3 +136,22 @@ where
         map_write.insert(key, StampedValue { value, changed_at });
     }
 }
+
+impl<DB, Q> UncheckedMutQueryStorageOps<DB, Q> for InputStorage<DB, Q>
+where
+    Q: Query<DB>,
+    DB: Database,
+    Q::Value: Default,
+{
+    fn set_unchecked(&self, db: &DB, key: &Q::Key, value: Q::Value) {
+        let key = key.clone();
+
+        let mut map_write = self.map.write();
+
+        // Unlike with `set`, here we use the **current revision** and
+        // do not create a new one.
+        let changed_at = db.salsa_runtime().current_revision();
+
+        map_write.insert(key, StampedValue { value, changed_at });
+    }
+}

src/lib.rs

@@ -136,6 +136,17 @@ where
     fn set(&self, db: &DB, key: &Q::Key, new_value: Q::Value);
 }
 
+/// An optional trait that is implemented for "user mutable" storage:
+/// that is, storage whose value is not derived from other storage but
+/// is set independently.
+pub trait UncheckedMutQueryStorageOps<DB, Q>: Default
+where
+    DB: Database,
+    Q: Query<DB>,
+{
+    fn set_unchecked(&self, db: &DB, key: &Q::Key, new_value: Q::Value);
+}
+
 #[derive(new)]
 pub struct QueryTable<'me, DB, Q>
 where
@@ -180,6 +191,22 @@ where
         self.storage.set(self.db, &key, value);
     }
 
+    /// Assigns a value to the query **bypassing the normal
+    /// incremental checking** -- this value becomes the value for the
+    /// query in the current revision. This can even be used on
+    /// "derived" queries (so long as their results are memoized).
+    ///
+    /// **This is only meant to be used for "mocking" purposes in
+    /// tests** -- when testing a given query, you can use
+    /// `set_unchecked` to assign the values for its various inputs
+    /// and thus control what it sees when it executes.
+    pub fn set_unchecked(&self, key: Q::Key, value: Q::Value)
+    where
+        Q::Storage: UncheckedMutQueryStorageOps<DB, Q>,
+    {
+        self.storage.set_unchecked(self.db, &key, value);
+    }
+
     fn descriptor(&self, key: &Q::Key) -> DB::QueryDescriptor {
         (self.descriptor_fn)(self.db, key)
     }

src/memoized.rs

@@ -7,6 +7,7 @@ use crate::QueryDescriptor;
 use crate::QueryFunction;
 use crate::QueryStorageOps;
 use crate::QueryTable;
+use crate::UncheckedMutQueryStorageOps;
 use log::debug;
 use parking_lot::{RwLock, RwLockUpgradableReadGuard};
 use rustc_hash::FxHashMap;
@@ -260,3 +261,26 @@ where
         }
     }
 }
+
+impl<DB, Q> UncheckedMutQueryStorageOps<DB, Q> for MemoizedStorage<DB, Q>
+where
+    Q: QueryFunction<DB>,
+    DB: Database,
+{
+    fn set_unchecked(&self, db: &DB, key: &Q::Key, value: Q::Value) {
+        let key = key.clone();
+
+        let mut map_write = self.map.write();
+
+        let changed_at = db.salsa_runtime().current_revision();
+
+        map_write.insert(
+            key,
+            QueryState::Memoized(Memo {
+                stamped_value: StampedValue { value, changed_at },
+                inputs: QueryDescriptorSet::new(),
+                verified_at: changed_at,
+            }),
+        );
+    }
+}

src/runtime.rs

@@ -227,7 +227,7 @@ impl<DB: Database> std::fmt::Debug for QueryDescriptorSet<DB> {
 }
 
 impl<DB: Database> QueryDescriptorSet<DB> {
-    fn new() -> Self {
+    crate fn new() -> Self {
         QueryDescriptorSet {
             set: FxIndexSet::default(),
         }