Merge pull request #194 from matklad/on-demand

On demand input queries

book/src/common_patterns.md

@@ -1 +1,51 @@
 # Common patterns
+
+## On Demand (Lazy) Inputs
+
+Salsa input quries work best if you can easily provide all of the inputs upfront.
+However sometimes the set of inputs is not known beforehand.
+
+A typical example is reading files from disk.
+While it is possible to eagerly scan a particular directory and create an in-memory file tree in a salsa input query, a more straight-forward approach is to read the files lazily.
+That is, when someone requests the text of a file for the first time:
+
+1. Read the file from disk and cache it.
+2. Setup a file-system watcher for this path.
+3. Invalidate the cached file once the watcher sends a change notification.
+
+This is possible to achieve in salsa, using a derived query and `report_synthetic_read` and `invalidate` queries.
+The setup looks roughtly like this:
+
+```rust,ignore
+#[salsa::query_group(VfsDatabaseStorage)]
+trait VfsDatabase: salsa::Database + FileWathcer {
+    fn read(&self, path: PathBuf) -> String;
+}
+
+trait FileWatcher {
+    fn watch(&self, path: &Path);
+    fn did_change_file(&self, path: &Path);
+}
+
+fn read(db: &impl salsa::Database, path: PathBuf) -> String {
+    db.salsa_runtime()
+        .report_synthetic_read(salsa::Durability::LOW);
+    db.watch(&path);
+    std::fs::read_to_string(&path).unwrap_or_default()
+}
+
+#[salsa::database(VfsDatabaseStorage)]
+struct MyDatabase { ... }
+
+impl FileWatcher for MyDatabase {
+    fn watch(&self, path: &Path) { ... }
+    fn did_change_file(&self, path: &Path) {
+        self.query_mut(ReadQuery).invalidate(path);
+    }
+}
+```
+
+* We declare the query as a derived query (which is the default).
+* In the query implementation, we don't call any other query and just directly read file from disk.
+* Because the query doesn't read any inputs, it will be assigned a `HIGH` durability by default, which we override with `report_synthetic_read`.
+* The result of the query is cached, and we must call `invalidate` to clear this cache.

src/derived.rs

@@ -1,6 +1,7 @@
 use crate::debug::TableEntry;
 use crate::durability::Durability;
 use crate::lru::Lru;
+use crate::plumbing::DerivedQueryStorageOps;
 use crate::plumbing::HasQueryGroup;
 use crate::plumbing::LruQueryStorageOps;
 use crate::plumbing::QueryFunction;
@@ -189,3 +190,22 @@ where
         self.lru_list.set_lru_capacity(new_capacity);
     }
 }
+
+impl<DB, Q, MP> DerivedQueryStorageOps<DB, Q> for DerivedStorage<DB, Q, MP>
+where
+    Q: QueryFunction<DB>,
+    DB: Database + HasQueryGroup<Q::Group>,
+    MP: MemoizationPolicy<DB, Q>,
+{
+    fn invalidate(&self, db: &DB, key: &Q::Key) {
+        db.salsa_runtime().with_incremented_revision(|guard| {
+            let map_read = self.slot_map.read();
+
+            if let Some(slot) = map_read.get(key) {
+                if let Some(durability) = slot.invalidate() {
+                    guard.mark_durability_as_changed(durability);
+                }
+            }
+        })
+    }
+}

src/derived/slot.rs

@@ -534,6 +534,15 @@ where
         }
     }
 
+    pub(super) fn invalidate(&self) -> Option<Durability> {
+        if let QueryState::Memoized(memo) = &mut *self.state.write() {
+            memo.inputs = MemoInputs::Untracked;
+            Some(memo.durability)
+        } else {
+            None
+        }
+    }
+
     /// Helper:
     ///
     /// When we encounter an `InProgress` indicator, we need to either

src/lib.rs

@@ -24,6 +24,7 @@ pub mod debug;
 #[doc(hidden)]
 pub mod plumbing;
 
+use crate::plumbing::DerivedQueryStorageOps;
 use crate::plumbing::InputQueryStorageOps;
 use crate::plumbing::LruQueryStorageOps;
 use crate::plumbing::QueryStorageMassOps;
@@ -552,6 +553,19 @@ where
     {
         self.storage.set_lru_capacity(cap);
     }
+
+    /// Marks the computed value as outdated.
+    ///
+    /// This causes salsa to re-execute the query function on the next access to
+    /// the query, even if all dependencies are up to date.
+    ///
+    /// This is most commonly used for invaliding on-demand inputs; see the [Salsa Book](https://salsa-rs.github.io/salsa/) for more information.
+    pub fn invalidate(&self, key: &Q::Key)
+    where
+        Q::Storage: plumbing::DerivedQueryStorageOps<DB, Q>,
+    {
+        self.storage.invalidate(self.db, key)
+    }
 }
 
 /// The error returned when a query could not be resolved due to a cycle

src/plumbing.rs

@@ -191,3 +191,11 @@ where
 pub trait LruQueryStorageOps: Default {
     fn set_lru_capacity(&self, new_capacity: usize);
 }
+
+pub trait DerivedQueryStorageOps<DB, Q>: Default
+where
+    DB: Database,
+    Q: Query<DB>,
+{
+    fn invalidate(&self, db: &DB, key: &Q::Key);
+}

src/runtime.rs

@@ -396,6 +396,14 @@ where
             .report_untracked_read(self.current_revision());
     }
 
+    /// Acts as though the current query had read an input with the given durability; this will force the current query's durability to be at most `durability`.
+    ///
+    /// This is mostly useful to control the durability level for on-demand inputs, as described in [the salsa book](https://salsa-rs.github.io/salsa/).
+    pub fn report_synthetic_read(&self, durability: Durability) {
+        self.local_state
+            .report_synthetic_read(durability);
+    }
+
     /// An "anonymous" read is a read that doesn't come from executing
     /// a query, but from some other internal operation. It just
     /// modifies the "changed at" to be at least the given revision.
@@ -693,6 +701,10 @@ impl<DB: Database> ActiveQuery<DB> {
         self.changed_at = changed_at;
     }
 
+    fn add_synthetic_read(&mut self, durability: Durability) {
+        self.durability = self.durability.min(durability);
+    }
+
     fn add_anon_read(&mut self, changed_at: Revision) {
         self.changed_at = self.changed_at.max(changed_at);
     }

src/runtime/local_state.rs

@@ -82,6 +82,12 @@ impl<DB: Database> LocalState<DB> {
         }
     }
 
+    pub(super) fn report_synthetic_read(&self, durability: Durability) {
+        if let Some(top_query) = self.query_stack.borrow_mut().last_mut() {
+            top_query.add_synthetic_read(durability);
+        }
+    }
+
     pub(super) fn report_anon_read(&self, revision: Revision) {
         if let Some(top_query) = self.query_stack.borrow_mut().last_mut() {
             top_query.add_anon_read(revision);

tests/on_demand_inputs.rs

@@ -0,0 +1,110 @@
+//! Test that "on-demand" input pattern works.
+//!
+//! On-demand inputs are inputs computed lazily on the fly. They are simulated
+//! via a b query with zero inputs, which uses `add_synthetic_read` to
+//! tweak durability and `invalidate` to clear the input.
+
+use std::{cell::Cell, collections::HashMap, rc::Rc};
+
+use salsa::{Database as _, Durability};
+
+#[salsa::query_group(QueryGroupStorage)]
+trait QueryGroup: salsa::Database + AsRef<HashMap<u32, u32>> {
+    fn a(&self, x: u32) -> u32;
+    fn b(&self, x: u32) -> u32;
+    fn c(&self, x: u32) -> u32;
+}
+
+fn a(db: &impl QueryGroup, x: u32) -> u32 {
+    let durability = if x % 2 == 0 {
+        Durability::LOW
+    } else {
+        Durability::HIGH
+    };
+    db.salsa_runtime().report_synthetic_read(durability);
+    let external_state: &HashMap<u32, u32> = db.as_ref();
+    external_state[&x]
+}
+
+fn b(db: &impl QueryGroup, x: u32) -> u32 {
+    db.a(x)
+}
+
+fn c(db: &impl QueryGroup, x: u32) -> u32 {
+    db.b(x)
+}
+
+#[salsa::database(QueryGroupStorage)]
+#[derive(Default)]
+struct Database {
+    runtime: salsa::Runtime<Database>,
+    external_state: HashMap<u32, u32>,
+    on_event: Option<Box<dyn Fn(salsa::Event<Database>)>>,
+}
+
+impl salsa::Database for Database {
+    fn salsa_runtime(&self) -> &salsa::Runtime<Database> {
+        &self.runtime
+    }
+
+    fn salsa_event(&self, event_fn: impl Fn() -> salsa::Event<Self>) {
+        if let Some(cb) = &self.on_event {
+            cb(event_fn())
+        }
+    }
+}
+
+impl AsRef<HashMap<u32, u32>> for Database {
+    fn as_ref(&self) -> &HashMap<u32, u32> {
+        &self.external_state
+    }
+}
+
+#[test]
+fn on_demand_input_works() {
+    let mut db = Database::default();
+
+    db.external_state.insert(1, 10);
+    assert_eq!(db.b(1), 10);
+    assert_eq!(db.a(1), 10);
+
+    // We changed external state, but haven't signaled about this yet,
+    // so we expect to see the old answer
+    db.external_state.insert(1, 92);
+    assert_eq!(db.b(1), 10);
+    assert_eq!(db.a(1), 10);
+
+    db.query_mut(AQuery).invalidate(&1);
+    assert_eq!(db.b(1), 92);
+    assert_eq!(db.a(1), 92);
+}
+
+#[test]
+fn on_demand_input_durability() {
+    let mut db = Database::default();
+    db.external_state.insert(1, 10);
+    db.external_state.insert(2, 20);
+    assert_eq!(db.b(1), 10);
+    assert_eq!(db.b(2), 20);
+
+    let validated = Rc::new(Cell::new(0));
+    db.on_event = Some(Box::new({
+        let validated = Rc::clone(&validated);
+        move |event| match event.kind {
+            salsa::EventKind::DidValidateMemoizedValue { .. } => validated.set(validated.get() + 1),
+            _ => (),
+        }
+    }));
+
+    db.salsa_runtime().synthetic_write(Durability::LOW);
+    validated.set(0);
+    assert_eq!(db.c(1), 10);
+    assert_eq!(db.c(2), 20);
+    assert_eq!(validated.get(), 2);
+
+    db.salsa_runtime().synthetic_write(Durability::HIGH);
+    validated.set(0);
+    assert_eq!(db.c(1), 10);
+    assert_eq!(db.c(2), 20);
+    assert_eq!(validated.get(), 4);
+}