Fix verification of tracked struct from high-durability query

Author: carljm

src/tracked_struct/struct_map.rs

@@ -100,7 +100,15 @@ where
     }
 
     pub fn validate<'db>(&'db self, runtime: &'db Runtime, id: Id) {
-        let mut data = self.map.get_mut(&id).unwrap();
+        Self::validate_in_map(&self.map, runtime, id)
+    }
+
+    pub fn validate_in_map<'db>(
+        map: &'db FxDashMap<Id, Alloc<Value<C>>>,
+        runtime: &'db Runtime,
+        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).
@@ -186,7 +194,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,
@@ -201,6 +209,50 @@ where
         unsafe { C::struct_from_raw(data.as_raw()) }
     }
 
+    fn get_and_validate_last_changed<'db>(
+        map: &'db FxDashMap<Id, Alloc<Value<C>>>,
+        runtime: &'db Runtime,
+        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 check 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), and validate to the current revision if so.
+        let created_at = data_ref.created_at;
+        let last_changed = runtime.last_changed_revision(data_ref.durability);
+        dbg!(
+            runtime.current_revision(),
+            created_at,
+            last_changed,
+            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.
+        if last_changed < runtime.current_revision() {
+            Self::validate_in_map(map, runtime, id);
+        }
+
+        ret
+    }
+
     /// Remove the entry for `id` from the map.
     ///
     /// NB. the data won't actually be freed until `drop_deleted_entries` is called.
@@ -234,6 +286,14 @@ where
     pub fn get<'db>(&'db self, current_revision: Revision, id: Id) -> C::Struct<'db> {
         StructMap::get_from_map(&self.map, current_revision, id)
     }
+
+    pub fn get_and_validate_last_changed<'db>(
+        &'db self,
+        runtime: &'db Runtime,
+        id: Id,
+    ) -> C::Struct<'db> {
+        StructMap::get_and_validate_last_changed(&self.map, runtime, id)
+    }
 }
 
 /// A mutable reference to the data for a single struct.

src/tracked_struct/tracked_field.rs

@@ -85,9 +85,10 @@ where
         input: Option<Id>,
         revision: crate::Revision,
     ) -> bool {
-        let current_revision = db.runtime().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.runtime(), 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,108 @@
+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);
+        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 {}
+
+    #[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);
+}