Fix verification of tracked struct from high-durability query

Author: carljm

src/tracked_struct/struct_map.rs

@@ -6,7 +6,7 @@ use std::{
 use crossbeam::queue::SegQueue;
 use dashmap::mapref::one::RefMut;
 
-use crate::{alloc::Alloc, hash::FxDashMap, Id, Revision};
+use crate::{alloc::Alloc, hash::FxDashMap, zalsa::Zalsa, Id, Revision};
 
 use super::{Configuration, KeyStruct, Value};
 
@@ -100,14 +100,22 @@ where
     }
 
     pub fn validate(&self, current_revision: Revision, id: Id) {
-        let mut data = self.map.get_mut(&id).unwrap();
+        Self::validate_in_map(&self.map, current_revision, id)
+    }
+
+    pub fn validate_in_map(
+        map: &FxDashMap<Id, Alloc<Value<C>>>,
+        current_revision: Revision,
+        id: Id,
+    ) {
+        let mut data = map.get_mut(&id).unwrap();
 
         // UNSAFE: We never permit `&`-access in the current revision until data.created_at
         // has been updated to the current revision (which we check below).
         let data = unsafe { data.as_mut() };
 
         // Never update a struct twice in the same revision.
-        assert!(data.created_at < current_revision);
+        assert!(data.created_at <= current_revision);
         data.created_at = current_revision;
     }
 
@@ -185,7 +193,7 @@ where
         let data_ref: &Value<C> = unsafe { data.as_ref() };
 
         // Before we drop the lock, check that the value has
-        // been updated in this revision. This is what allows us to return a ``
+        // been updated in this revision. This is what allows us to return a Struct
         let created_at = data_ref.created_at;
         assert!(
             created_at == current_revision,
@@ -200,6 +208,56 @@ where
         unsafe { C::struct_from_raw(data.as_raw()) }
     }
 
+    /// Lookup an existing tracked struct from the map, maybe validating it to current revision.
+    ///
+    /// Validates to current revision if the struct was last created/validated in a revision that
+    /// is still current for the struct's durability. That is, if the struct is HIGH durability
+    /// (created by a HIGH durability query) and was created in R2, and we are now at R3 but no
+    /// HIGH durability input has changed since R2, the struct is still valid and we can validate
+    /// it to R3.
+    ///
+    /// # Panics
+    ///
+    /// * If the value is not present in the map.
+    /// * If the value has not been updated in the last-changed revision for its durability.
+    fn get_and_validate_last_changed<'db>(
+        map: &'db FxDashMap<Id, Alloc<Value<C>>>,
+        zalsa: &Zalsa,
+        id: Id,
+    ) -> C::Struct<'db> {
+        let data = map.get(&id).unwrap();
+
+        // UNSAFE: We permit `&`-access in the current revision once data.created_at
+        // has been updated to the current revision (which we ensure below).
+        let data_ref: &Value<C> = unsafe { data.as_ref() };
+
+        // Before we drop the lock, check that the value has been updated in the most recent
+        // version in which the query that created it could have changed (based on durability).
+        let created_at = data_ref.created_at;
+        let last_changed = zalsa.last_changed_revision(data_ref.durability);
+        assert!(
+            created_at >= last_changed,
+            "access to tracked struct from obsolete revision"
+        );
+
+        // Unsafety clause:
+        //
+        // * Value will not be updated again in this revision,
+        //   and revision will not change so long as runtime is shared
+        // * We only remove values from the map when we have `&mut self`
+        let ret = unsafe { C::struct_from_raw(data.as_raw()) };
+
+        drop(data);
+
+        // Validate in current revision, if necessary.
+        let current_revision = zalsa.current_revision();
+        if last_changed < current_revision {
+            Self::validate_in_map(map, current_revision, id);
+        }
+
+        ret
+    }
+
     /// Remove the entry for `id` from the map.
     ///
     /// NB. the data won't actually be freed until `drop_deleted_entries` is called.
@@ -233,6 +291,10 @@ where
     pub fn get(&self, current_revision: Revision, id: Id) -> C::Struct<'_> {
         StructMap::get_from_map(&self.map, current_revision, id)
     }
+
+    pub fn get_and_validate_last_changed<'db>(&'db self, zalsa: &Zalsa, id: Id) -> C::Struct<'db> {
+        StructMap::get_and_validate_last_changed(&self.map, zalsa, id)
+    }
 }
 
 /// A mutable reference to the data for a single struct.

src/tracked_struct/tracked_field.rs

@@ -83,9 +83,10 @@ where
         input: Option<Id>,
         revision: crate::Revision,
     ) -> bool {
-        let current_revision = db.zalsa().current_revision();
         let id = input.unwrap();
-        let data = self.struct_map.get(current_revision, id);
+        let data = self
+            .struct_map
+            .get_and_validate_last_changed(db.zalsa(), id);
         let data = C::deref_struct(data);
         let field_changed_at = data.revisions[self.field_index];
         field_changed_at > revision

tests/tracked_struct_not_validated_due_to_durability.rs

@@ -0,0 +1,131 @@
+/// Test that high durabilities can't cause "access tracked struct from previous revision" panic.
+///
+/// The test models a situation where we have two File inputs (0, 1), where `File(0)` has LOW
+/// durability and `File(1)` has HIGH durability. We can query an `index` for each file, and a
+/// `definitions` from that index (just a sub-part of the index), and we can `infer` each file. The
+/// `index` and `definitions` queries depend only on the `File` they operate on, but the `infer`
+/// query has some other dependencies: `infer(0)` depends on `infer(1)`, and `infer(1)` also
+/// depends directly on `File(0)`.
+///
+/// The panic occurs (in versions of Salsa without a fix) because `definitions(1)` is high
+/// durability, and depends on `index(1)` which is also high durability. `index(1)` creates the
+/// tracked struct `Definition(1)`, and `infer(1)` (which is low durability) depends on
+/// `Definition.file(1)`.
+///
+/// After a change to `File(0)` (low durability), we only shallowly verify `definitions(1)` -- it
+/// passes shallow verification due to durability. We take care to mark-validated the outputs of
+/// `definitions(1)`, but we never verify `index(1)` at all (deeply or shallowly), which means we
+/// never mark `Definition(1)` validated. So when we deep-verify `infer(1)`, we try to access its
+/// dependency `Definition.file(1)`, and hit the panic because we are accessing a tracked struct
+/// that has never been re-validated or re-recreated in R2.
+use salsa::{Durability, Setter};
+
+#[salsa::db]
+trait Db: salsa::Database {
+    fn file(&self, idx: usize) -> File;
+}
+
+#[salsa::input]
+struct File {
+    field: usize,
+}
+
+#[salsa::tracked]
+struct Definition<'db> {
+    file: File,
+}
+
+#[salsa::tracked]
+struct Index<'db> {
+    definitions: Definitions<'db>,
+}
+
+#[salsa::tracked]
+struct Definitions<'db> {
+    definition: Definition<'db>,
+}
+
+#[salsa::tracked]
+struct Inference<'db> {
+    definition: Definition<'db>,
+}
+
+#[salsa::tracked]
+fn index<'db>(db: &'db dyn Db, file: File) -> Index<'db> {
+    let _ = file.field(db);
+    Index::new(db, Definitions::new(db, Definition::new(db, file)))
+}
+
+#[salsa::tracked]
+fn definitions<'db>(db: &'db dyn Db, file: File) -> Definitions<'db> {
+    index(db, file).definitions(db)
+}
+
+#[salsa::tracked]
+fn infer<'db>(db: &'db dyn Db, definition: Definition<'db>) -> Inference<'db> {
+    let file = definition.file(db);
+    if file.field(db) < 1 {
+        let dependent_file = db.file(1);
+        infer(db, definitions(db, dependent_file).definition(db))
+    } else {
+        db.file(0).field(db);
+        index(db, file);
+        Inference::new(db, definition)
+    }
+}
+
+#[salsa::tracked]
+fn check<'db>(db: &'db dyn Db, file: File) -> Inference<'db> {
+    let defs = definitions(db, file);
+    infer(db, defs.definition(db))
+}
+
+#[test]
+fn execute() {
+    #[salsa::db]
+    #[derive(Default)]
+    struct Database {
+        storage: salsa::Storage<Self>,
+        files: Vec<File>,
+    }
+
+    #[salsa::db]
+    impl salsa::Database for Database {
+        fn salsa_event(&self, _event: &dyn Fn() -> salsa::Event) {}
+    }
+
+    #[salsa::db]
+    impl Db for Database {
+        fn file(&self, idx: usize) -> File {
+            self.files[idx]
+        }
+    }
+
+    let mut db = Database::default();
+    // Create a file 0 with low durability, and a file 1 with high durability.
+
+    let file0 = File::new(&db, 0);
+    db.files.push(file0);
+
+    let file1 = File::new(&db, 1);
+    file1
+        .set_field(&mut db)
+        .with_durability(Durability::HIGH)
+        .to(1);
+    db.files.push(file1);
+
+    // check(0) -> infer(0) -> definitions(0) -> index(0)
+    //                     \-> infer(1) -> definitions(1) -> index(1)
+
+    assert_eq!(check(&db, file0).definition(&db).file(&db).field(&db), 1);
+
+    // update the low durability file 0
+    file0.set_field(&mut db).to(0);
+
+    // Re-query check(0). definitions(1) is high durability so it short-circuits in shallow-verify,
+    // meaning we never verify index(1) at all, but index(1) created the tracked struct
+    // Definition(1), so we never validate Definition(1) in R2, so when we try to verify
+    // Definition.file(1) (as an input of infer(1) ) we hit a panic for trying to use a struct that
+    // isn't validated in R2.
+    check(&db, file0);
+}