Merge pull request #33 from nikomatsakis/set-unchecked

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

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "salsa"
-version = "0.4.0"
+version = "0.4.1"
 authors = ["Niko Matsakis <niko@alum.mit.edu>"]
 edition = "2018"
 license = "Apache-2.0 OR MIT"

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(),
         }

tests/set_unchecked.rs

@@ -0,0 +1,95 @@
+use salsa::Database;
+
+salsa::query_group! {
+    trait HelloWorldDatabase: salsa::Database {
+        fn input(key: ()) -> String {
+            type Input;
+            storage input;
+        }
+
+        fn length(key: ()) -> usize {
+            type Length;
+        }
+
+        fn double_length(key: ()) -> usize {
+            type DoubleLength;
+        }
+    }
+}
+
+fn length(db: &impl HelloWorldDatabase, (): ()) -> usize {
+    let l = db.input(()).len();
+    assert!(l > 0); // not meant to be invoked with no input
+    l
+}
+
+fn double_length(db: &impl HelloWorldDatabase, (): ()) -> usize {
+    db.length(()) * 2
+}
+
+#[derive(Default)]
+struct DatabaseStruct {
+    runtime: salsa::runtime::Runtime<DatabaseStruct>,
+}
+
+impl salsa::Database for DatabaseStruct {
+    fn salsa_runtime(&self) -> &salsa::runtime::Runtime<DatabaseStruct> {
+        &self.runtime
+    }
+}
+
+salsa::database_storage! {
+    struct DatabaseStorage for DatabaseStruct {
+        impl HelloWorldDatabase {
+            fn input() for Input;
+            fn length() for Length;
+            fn double_length() for DoubleLength;
+        }
+    }
+}
+
+#[test]
+fn normal() {
+    let db = DatabaseStruct::default();
+    db.query(Input).set((), format!("Hello, world"));
+    assert_eq!(db.double_length(()), 24);
+    db.query(Input).set((), format!("Hello, world!"));
+    assert_eq!(db.double_length(()), 26);
+}
+
+#[test]
+#[should_panic]
+fn use_without_set() {
+    let db = DatabaseStruct::default();
+    db.double_length(());
+}
+
+#[test]
+fn using_set_unchecked_on_input() {
+    let db = DatabaseStruct::default();
+    db.query(Input).set_unchecked((), format!("Hello, world"));
+    assert_eq!(db.double_length(()), 24);
+}
+
+#[test]
+fn using_set_unchecked_on_input_after() {
+    let db = DatabaseStruct::default();
+    db.query(Input).set((), format!("Hello, world"));
+    assert_eq!(db.double_length(()), 24);
+
+    // If we use `set_unchecked`, we don't notice that `double_length`
+    // is out of date. Oh well, don't do that.
+    db.query(Input).set_unchecked((), format!("Hello, world!"));
+    assert_eq!(db.double_length(()), 24);
+}
+
+#[test]
+fn using_set_unchecked() {
+    let db = DatabaseStruct::default();
+
+    // Use `set_unchecked` to intentionally set the wrong value,
+    // demonstrating that the code never runs.
+    db.query(Length).set_unchecked((), 24);
+
+    assert_eq!(db.double_length(()), 48);
+}