2 files changed, 76 insertions, 12 deletions

diff --git a/crates/salsa/src/lib.rs b/crates/salsa/src/lib.rs
index a54f2a06f..5de3c7774 100644
--- a/crates/salsa/src/lib.rs
+++ b/crates/salsa/src/lib.rs
@@ -3,7 +3,7 @@ extern crate parking_lot;
 
 use std::{
     sync::Arc,
-    collections::HashMap,
+    collections::{HashSet, HashMap},
     cell::RefCell,
 };
 use parking_lot::Mutex;
@@ -138,7 +138,16 @@ where
             return (record.output.clone(), record.output_fingerprint)
         }
         if self.query_config.ground_fn.contains_key(&query_id.0) {
-            return self.force(query_id, params);
+            let (invalidated, record) = {
+                let guard = self.db.graph.lock();
+                let (gen, ref record) = guard[&query_id];
+                (gen == INVALIDATED, record.clone())
+            };
+            if invalidated {
+                return self.force(query_id, params);
+            } else {
+                return (record.output.clone(), record.output_fingerprint);
+            }
         }
         for (dep_query_id, prev_fingerprint) in record.deps.iter().cloned() {
             let dep_params: D = {
@@ -198,6 +207,28 @@ where
     }
 }
 
+pub struct Invalidations {
+    types: HashSet<QueryTypeId>,
+    ids: Vec<QueryId>,
+}
+
+impl Invalidations {
+    pub fn new() -> Invalidations {
+        Invalidations {
+            types: HashSet::new(),
+            ids: Vec::new(),
+        }
+    }
+    pub fn invalidate(
+        &mut self,
+        query_type: QueryTypeId,
+        params: impl Iterator<Item=InputFingerprint>,
+    ) {
+        self.types.insert(query_type);
+        self.ids.extend(params.map(|it| QueryId(query_type, it)))
+    }
+}
+
 impl<T, D> Db<T, D>
 where
     D: Clone
@@ -209,9 +240,25 @@ where
         }
     }
 
-    pub fn with_ground_data(&self, ground_data: T) -> Db<T, D> {
+    pub fn with_ground_data(
+        &self,
+        ground_data: T,
+        invalidations: Invalidations,
+    ) -> Db<T, D> {
+        for id in self.query_config.ground_fn.keys() {
+            assert!(
+                invalidations.types.contains(id),
+                "all ground queries must be invalidated"
+            );
+        }
+
         let gen = Gen(self.db.gen.0 + 1);
-        let graph = self.db.graph.lock().clone();
+        let mut graph = self.db.graph.lock().clone();
+        for id in invalidations.ids {
+            if let Some((gen, _)) = graph.get_mut(&id) {
+                *gen = INVALIDATED;
+            }
+        }
         let graph = Mutex::new(graph);
         Db {
             db: Arc::new(DbState { ground_data, gen, graph }),
@@ -232,6 +279,7 @@ where
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 struct Gen(u64);
+const INVALIDATED: Gen = Gen(!0);
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
 pub struct InputFingerprint(pub u64);
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
diff --git a/crates/salsa/tests/integration.rs b/crates/salsa/tests/integration.rs
index 2872d3913..3cec330e6 100644
--- a/crates/salsa/tests/integration.rs
+++ b/crates/salsa/tests/integration.rs
@@ -1,5 +1,6 @@
 extern crate salsa;
 use std::{
+    iter::once,
     sync::Arc,
     collections::hash_map::{HashMap, DefaultHasher},
     any::Any,
@@ -113,30 +114,45 @@ fn test_number_of_lines() {
     assert_eq!(trace.len(), 0);
 
     state.insert(1, "hello\nworld".to_string());
-    let db = db.with_ground_data(state.clone());
+    let mut inv = salsa::Invalidations::new();
+    inv.invalidate(GET_TEXT, once(i_print(&1u32)));
+    inv.invalidate(GET_FILES, once(i_print(&())));
+    let db = db.with_ground_data(state.clone(), inv);
     let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
     assert_eq!(*newlines, 2);
     assert_eq!(trace.len(), 4);
 
     state.insert(2, "spam\neggs".to_string());
-    let db = db.with_ground_data(state.clone());
+    let mut inv = salsa::Invalidations::new();
+    inv.invalidate(GET_TEXT, once(i_print(&2u32)));
+    inv.invalidate(GET_FILES, once(i_print(&())));
+    let db = db.with_ground_data(state.clone(), inv);
     let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
     assert_eq!(*newlines, 4);
-    assert_eq!(trace.len(), 5);
+    assert_eq!(trace.len(), 4);
 
+    let mut invs = vec![];
     for i in 0..10 {
-        state.insert(i + 10, "spam".to_string());
+        let id = i + 10;
+        invs.push(i_print(&id));
+        state.insert(id, "spam".to_string());
     }
-    let db = db.with_ground_data(state.clone());
+    let mut inv = salsa::Invalidations::new();
+    inv.invalidate(GET_TEXT, invs.into_iter());
+    inv.invalidate(GET_FILES, once(i_print(&())));
+    let db = db.with_ground_data(state.clone(), inv);
     let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
     assert_eq!(*newlines, 14);
-    assert_eq!(trace.len(), 24);
+    assert_eq!(trace.len(), 22);
 
     state.insert(15, String::new());
-    let db = db.with_ground_data(state.clone());
+    let mut inv = salsa::Invalidations::new();
+    inv.invalidate(GET_TEXT, once(i_print(&15u32)));
+    inv.invalidate(GET_FILES, once(i_print(&())));
+    let db = db.with_ground_data(state.clone(), inv);
     let (newlines, trace) = get::<(), usize>(&db, TOTAL_NEWLINES, ());
     assert_eq!(*newlines, 13);
-    assert_eq!(trace.len(), 15);
+    assert_eq!(trace.len(), 4);
 }
 
 fn o_print<T: Hash>(x: &T) -> salsa::OutputFingerprint {