add plumbing to reuse interned slots

Author: ibraheemdev

src/active_query.rs

@@ -142,7 +142,7 @@ impl ActiveQuery {
     /// Used during cycle recovery, see [`Runtime::unblock_cycle_and_maybe_throw`].
     pub(super) fn remove_cycle_participants(&mut self, cycle: &Cycle) {
         for p in cycle.participant_keys() {
-            let p: DatabaseKeyIndex = p.into();
+            let p: DatabaseKeyIndex = p;
             self.input_outputs.shift_remove(&(EdgeKind::Input, p));
         }
     }

src/function/fetch.rs

@@ -21,7 +21,7 @@ where
             self.evict_value_from_memo_for(zalsa, evicted);
         }
 
-        zalsa_local.report_tracked_read(self.database_key_index(id).into(), durability, changed_at);
+        zalsa_local.report_tracked_read(self.database_key_index(id), durability, changed_at);
 
         value
     }

src/function/specify.rs

@@ -39,8 +39,7 @@ where
         //
         // Now, if We invoke Q3 first, We get one result for Q2, but if We invoke Q4 first, We get a different value. That's no good.
         let database_key_index = <C::Input<'db>>::database_key_index(db.as_dyn_database(), key);
-        let dependency_index = database_key_index.into();
-        if !zalsa_local.is_output_of_active_query(dependency_index) {
+        if !zalsa_local.is_output_of_active_query(database_key_index) {
             panic!("can only use `specify` on salsa structs created during the current tracked fn");
         }
 
@@ -92,7 +91,7 @@ where
 
         // Record that the current query *specified* a value for this cell.
         let database_key_index = self.database_key_index(key);
-        zalsa_local.add_output(database_key_index.into());
+        zalsa_local.add_output(database_key_index);
     }
 
     /// Invoked when the query `executor` has been validated as having green inputs

src/interned.rs

@@ -2,6 +2,7 @@ use crate::durability::Durability;
 use crate::id::AsId;
 use crate::ingredient::fmt_index;
 use crate::plumbing::{Jar, JarAux};
+use crate::revision::AtomicRevision;
 use crate::table::memo::MemoTable;
 use crate::table::sync::SyncTable;
 use crate::table::Slot;
@@ -67,6 +68,10 @@ where
     data: C::Data<'static>,
     memos: MemoTable,
     syncs: SyncTable,
+    /// The revision the value was first interned in.
+    first_interned_at: Revision,
+    /// The most recent interned revision.
+    last_interned_at: AtomicRevision,
 }
 
 impl<C: Configuration> Default for JarImpl<C> {
@@ -120,7 +125,8 @@ where
         db: &'db dyn crate::Database,
         data: impl Lookup<C::Data<'db>>,
     ) -> C::Struct<'db> {
-        let zalsa_local = db.zalsa_local();
+        let (zalsa, zalsa_local) = db.zalsas();
+        let current_revision = zalsa.current_revision();
 
         // Optimisation to only get read lock on the map if the data has already
         // been interned.
@@ -142,9 +148,13 @@ where
                 // SAFETY: Read lock on map is held during this block
                 let id = unsafe { *bucket.as_ref().1.get() };
 
+                // Sync the value's revision.
+                let value = zalsa.table().get::<Value<C>>(id);
+                value.last_interned_at.store(current_revision);
+
                 // Record a dependency on this value.
                 let index = self.database_key_index(id);
-                zalsa_local.report_tracked_read(index, Durability::MAX, Revision::start());
+                zalsa_local.report_tracked_read(index, Durability::MAX, current_revision);
 
                 return C::struct_from_id(id);
             }
@@ -160,9 +170,13 @@ where
                 let id = *entry.get();
                 drop(entry);
 
+                // Sync the value's revision.
+                let value = zalsa.table().get::<Value<C>>(id);
+                value.last_interned_at.store(current_revision);
+
                 // Record a dependency on this value.
                 let index = self.database_key_index(id);
-                zalsa_local.report_tracked_read(index, Durability::MAX, Revision::start());
+                zalsa_local.report_tracked_read(index, Durability::MAX, current_revision);
 
                 C::struct_from_id(id)
             }
@@ -171,43 +185,56 @@ where
             dashmap::mapref::entry::Entry::Vacant(entry) => {
                 let zalsa = db.zalsa();
                 let table = zalsa.table();
+
                 let next_id = zalsa_local.allocate(table, self.ingredient_index, || Value::<C> {
                     data: internal_data,
                     memos: Default::default(),
                     syncs: Default::default(),
+                    first_interned_at: current_revision,
+                    last_interned_at: AtomicRevision::from(current_revision),
                 });
                 entry.insert(next_id);
 
                 // Record a dependency on this value.
                 let index = self.database_key_index(next_id);
-                zalsa_local.report_tracked_read(index, Durability::MAX, Revision::start());
+                zalsa_local.report_tracked_read(index, Durability::MAX, current_revision);
 
                 C::struct_from_id(next_id)
             }
         }
     }
 
     /// Returns the database key index for an interned value with the given id.
-    pub fn database_key_index(&self, id: Id) -> DatabaseKeyIndex {
+    pub fn database_key_index(&self, key_index: Id) -> DatabaseKeyIndex {
         DatabaseKeyIndex {
+            key_index,
             ingredient_index: self.ingredient_index,
-            key_index: id,
         }
     }
 
     /// Lookup the data for an interned value based on its id.
     /// Rarely used since end-users generally carry a struct with a pointer directly
     /// to the interned item.
     pub fn data<'db>(&'db self, db: &'db dyn Database, id: Id) -> &'db C::Data<'db> {
-        let internal_data = db.zalsa().table().get::<Value<C>>(id);
-        unsafe { Self::from_internal_data(&internal_data.data) }
+        let value = db.zalsa().table().get::<Value<C>>(id);
+        assert!(
+            value.last_interned_at.load() >= db.zalsa().current_revision(),
+            "Data was not interned in the current revision."
+        );
+        unsafe { Self::from_internal_data(&value.data) }
     }
 
     /// Lookup the fields from an interned struct.
     /// Note that this is not "leaking" since no dependency edge is required.
     pub fn fields<'db>(&'db self, db: &'db dyn Database, s: C::Struct<'db>) -> &'db C::Data<'db> {
         self.data(db, C::deref_struct(s))
     }
+
+    pub fn reset(&mut self, db: &mut dyn Database) {
+        // Trigger a new revision.
+        let _zalsa_mut = db.zalsa_mut();
+        self.key_map.clear();
+    }
 }
 
 impl<C> Ingredient for IngredientImpl<C>
@@ -218,8 +245,16 @@ where
         self.ingredient_index
     }
 
-    fn maybe_changed_after(&self, _db: &dyn Database, _input: Id, _revision: Revision) -> bool {
-        // Interned data currently never changes.
+    fn maybe_changed_after(&self, db: &dyn Database, input: Id, revision: Revision) -> bool {
+        let value = db.zalsa().table().get::<Value<C>>(input);
+        if value.first_interned_at > revision {
+            // The slot was reused.
+            return true;
+        }
+
+        // The slot is valid in this revision but we have to sync the value's revision.
+        value.last_interned_at.store(db.zalsa().current_revision());
+
         false
     }
 

src/revision.rs

@@ -62,3 +62,11 @@ impl AtomicRevision {
         self.data.store(r.as_usize(), Ordering::SeqCst);
     }
 }
+
+impl From<Revision> for AtomicRevision {
+    fn from(value: Revision) -> Self {
+        Self {
+            data: AtomicUsize::from(value.as_usize()),
+        }
+    }
+}

src/tracked_struct.rs

@@ -291,7 +291,7 @@ where
         match zalsa_local.tracked_struct_id(&identity) {
             Some(id) => {
                 // The struct already exists in the intern map.
-                zalsa_local.add_output(self.database_key_index(id).into());
+                zalsa_local.add_output(self.database_key_index(id));
                 self.update(zalsa, current_revision, id, &current_deps, fields);
                 C::struct_from_id(id)
             }
@@ -300,7 +300,7 @@ where
                 // This is a new tracked struct, so create an entry in the struct map.
                 let id = self.allocate(zalsa, zalsa_local, current_revision, &current_deps, fields);
                 let key = self.database_key_index(id);
-                zalsa_local.add_output(key.into());
+                zalsa_local.add_output(key);
                 zalsa_local.store_tracked_struct_id(identity, id);
                 C::struct_from_id(id)
             }

tests/interned-revisions.rs

@@ -0,0 +1,72 @@
+//! Test that a `tracked` fn on a `salsa::input`
+//! compiles and executes successfully.
+
+mod common;
+use salsa::{Database, Setter};
+use test_log::test;
+
+#[salsa::input]
+struct Input {
+    field1: usize,
+}
+
+#[salsa::interned]
+struct Interned<'db> {
+    field1: usize,
+}
+
+#[test]
+fn test_shallow_memo() {
+    #[salsa::tracked]
+    fn function<'db>(db: &'db dyn Database, _input: Input) -> Interned<'db> {
+        Interned::new(db, 0)
+    }
+
+    let mut db = common::EventLoggerDatabase::default();
+    let input = Input::new(&db, 0);
+
+    let result_in_rev_1 = function(&db, input);
+    assert_eq!(result_in_rev_1.field1(&db), 0);
+
+    input.set_field1(&mut db).to(1);
+
+    let result_in_rev_2 = function(&db, input);
+    assert_eq!(result_in_rev_2.field1(&db), 0);
+}
+
+// #[test]
+// fn test_no_reintern_input() {
+//     #[salsa::tracked]
+//     fn function<'db>(db: &'db dyn Database, input: Input) -> Interned<'db> {
+//         function2(db, input.field1(db))
+//     }
+//
+//     fn function2<'db>(db: &'db dyn Database, value: usize) -> Interned<'db> {
+//         Interned::new(db, value)
+//     }
+//
+//     let mut db = common::EventLoggerDatabase::default();
+//
+//     let input = Input::new(&db, 0);
+//     let result_in_rev_1 = function(&db, input);
+//     db.assert_logs(expect![[r#"
+//         [
+//             "Event { thread_id: ThreadId(3), kind: WillCheckCancellation }",
+//             "Event { thread_id: ThreadId(3), kind: WillExecute { database_key: function(Id(0)) } }",
+//         ]"#]]);
+//
+//     assert_eq!(result_in_rev_1.field1(&db), 0);
+//
+//     // Modify the input to force the value to be re-interned.
+//     input.set_field1(&mut db).to(1);
+//
+//     let result_in_rev_2 = function(&db, input);
+//     db.assert_logs(expect![[r#"
+//         [
+//             "Event { thread_id: ThreadId(3), kind: DidSetCancellationFlag }",
+//             "Event { thread_id: ThreadId(3), kind: WillCheckCancellation }",
+//             "Event { thread_id: ThreadId(3), kind: WillExecute { database_key: function(Id(0)) } }",
+//         ]"#]]);
+//
+//     assert_eq!(result_in_rev_2.field1(&db), 1);
+// }

tests/preverify-struct-with-leaked-data-2.rs

@@ -95,6 +95,6 @@ fn test_leaked_inputs_ignored() {
     // value of 100 since the struct has already been read during
     // this revision.
     //
-    // Contrast with preverify-struct-with-leaked-data-2.rs.
+    // Contrast with preverify-struct-with-leaked-data.rs.
     assert_eq!(result_in_rev_2, (0, 0));
 }