track "dependent functions"

components/salsa-2022-macros/src/input.rs

@@ -205,6 +205,10 @@ impl InputStruct {
             where
                 DB: ?Sized + salsa::DbWithJar<#jar_ty>,
             {
+                fn register_dependent_fn(_db: &DB, _index: salsa::routes::IngredientIndex) {
+                    // Do nothing here, at least for now.
+                    // If/when we add ability to delete inputs, this would become relevant.
+                }
             }
         }
     }

components/salsa-2022-macros/src/interned.rs

@@ -152,6 +152,10 @@ impl InternedStruct {
             where
                 DB: ?Sized + salsa::DbWithJar<#jar_ty>,
             {
+                fn register_dependent_fn(_db: &DB, _index: salsa::routes::IngredientIndex) {
+                    // Do nothing here, at least for now.
+                    // If/when we add ability to delete inputs, this would become relevant.
+                }
             }
         }
     }

components/salsa-2022-macros/src/tracked_struct.rs

@@ -211,11 +211,17 @@ impl TrackedStruct {
     fn salsa_struct_in_db_impl(&self) -> syn::ItemImpl {
         let ident = self.id_ident();
         let jar_ty = self.jar_ty();
+        let tracked_struct_index: Literal = self.tracked_struct_index();
         parse_quote! {
             impl<DB> salsa::salsa_struct::SalsaStructInDb<DB> for #ident
             where
                 DB: ?Sized + salsa::DbWithJar<#jar_ty>,
             {
+                fn register_dependent_fn(db: &DB, index: salsa::routes::IngredientIndex) {
+                    let (jar, _) = <_ as salsa::storage::HasJar<#jar_ty>>::jar(db);
+                    let ingredients = <#jar_ty as salsa::storage::HasIngredientsFor<#ident>>::ingredient(jar);
+                    ingredients.#tracked_struct_index.register_dependent_fn(index)
+                }
             }
         }
     }

components/salsa-2022/src/function.rs

@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use arc_swap::ArcSwap;
-use crossbeam::queue::SegQueue;
+use crossbeam::{atomic::AtomicCell, epoch::Atomic, queue::SegQueue};
 
 use crate::{
     cycle::CycleRecoveryStrategy,
@@ -67,6 +67,10 @@ pub struct FunctionIngredient<C: Configuration> {
     /// we don't know that we can trust the database to give us the same runtime
     /// everytime and so forth.
     deleted_entries: SegQueue<ArcSwap<memo::Memo<C::Value>>>,
+
+    /// Set to true once we invoke `register_dependent_fn` for `C::SalsaStruct`.
+    /// Prevents us from registering more than once.
+    registered: AtomicCell<bool>,
 }
 
 pub trait Configuration {
@@ -140,6 +144,7 @@ where
             lru: Default::default(),
             sync_map: Default::default(),
             deleted_entries: Default::default(),
+            registered: Default::default(),
         }
     }
 
@@ -169,7 +174,13 @@ where
         std::mem::transmute(memo_value)
     }
 
-    fn insert_memo(&self, key: C::Key, memo: memo::Memo<C::Value>) -> Option<&C::Value> {
+    fn insert_memo(
+        &self,
+        db: &DynDb<'_, C>,
+        key: C::Key,
+        memo: memo::Memo<C::Value>,
+    ) -> Option<&C::Value> {
+        self.register(db);
         let memo = Arc::new(memo);
         let value = unsafe {
             // Unsafety conditions: memo must be in the map (it's not yet, but it will be by the time this
@@ -183,6 +194,15 @@ where
         }
         value
     }
+
+    /// Register this function as a dependent fn of the given salsa struct.
+    /// When instances of that salsa struct are deleted, we'll get a callback
+    /// so we can remove any data keyed by them.
+    fn register(&self, db: &DynDb<'_, C>) {
+        if !self.registered.fetch_or(true) {
+            <C::SalsaStruct as SalsaStructInDb<_>>::register_dependent_fn(db, self.index)
+        }
+    }
 }
 
 impl<DB, C> Ingredient<DB> for FunctionIngredient<C>

components/salsa-2022/src/function/execute.rs

@@ -91,7 +91,11 @@ where
         }
 
         let value = self
-            .insert_memo(key, Memo::new(Some(value), revision_now, revisions.clone()))
+            .insert_memo(
+                db,
+                key,
+                Memo::new(Some(value), revision_now, revisions.clone()),
+            )
             .unwrap();
 
         let stamped_value = revisions.stamped_value(value);

components/salsa-2022/src/function/specify.rs

@@ -75,7 +75,7 @@ where
         };
 
         log::debug!("specify: about to add memo {:#?} for key {:?}", memo, key);
-        self.insert_memo(key, memo);
+        self.insert_memo(db, key, memo);
     }
 
     /// Specify the value for `key` but do not record it is an output.

components/salsa-2022/src/ingredient_list.rs

@@ -0,0 +1,83 @@
+use std::sync::Arc;
+
+use arc_swap::{ArcSwapOption, AsRaw};
+
+use crate::IngredientIndex;
+
+/// A list of ingredients that can be added to in parallel.
+pub(crate) struct IngredientList {
+    /// A list of each tracked functions.
+    /// tracked struct.
+    ///
+    /// Whenever an instance `i` of this struct is deleted,
+    /// each of these functions will be notified
+    /// so they can remove any data tied to that instance.
+    list: ArcSwapOption<Vec<IngredientIndex>>,
+}
+
+impl IngredientList {
+    pub fn new() -> Self {
+        Self {
+            list: ArcSwapOption::new(None),
+        }
+    }
+
+    /// Returns an iterator over the items in the list.
+    /// This is a snapshot of the list as it was when this function is called.
+    /// Items could still be added in parallel via `add_ingredient`
+    /// that will not be returned by this iterator.
+    pub(crate) fn iter(&self) -> impl Iterator<Item = IngredientIndex> {
+        let guard = self.list.load();
+        let mut index = 0;
+        std::iter::from_fn(move || match &*guard {
+            Some(list) if index < list.len() => {
+                let r = list[index];
+                index += 1;
+                Some(r)
+            }
+            _ => None,
+        })
+    }
+
+    /// Adds an ingredient to the list (if not already present).
+    pub(crate) fn push(&self, index: IngredientIndex) {
+        // This function is called whenever a value is stored,
+        // so other tracked functions and things may be executing,
+        // and there could even be two calls to this function in parallel.
+        //
+        // We use a "compare-and-swap" strategy of reading the old vector, creating a new vector,
+        // and then installing it, hoping that nobody has conflicted with us.
+        // If that fails, we start over.
+
+        loop {
+            let guard = self.list.load();
+            let empty_vec = vec![];
+            let old_vec = match &*guard {
+                Some(v) => v,
+                None => &empty_vec,
+            };
+
+            // First check whether the index is already present.
+            if old_vec.contains(&index) {
+                return;
+            }
+
+            // If not, construct a new vector that has all the old values, followed by `index`.
+            let vec: Arc<Vec<IngredientIndex>> = Arc::new(
+                old_vec
+                    .iter()
+                    .copied()
+                    .chain(std::iter::once(index))
+                    .collect(),
+            );
+
+            // Try to replace the old vector with the new one. If we fail, loop around again.
+            assert_eq!(vec.len(), vec.capacity());
+            let previous = self.list.compare_and_swap(&guard, Some(vec));
+            if guard.as_raw() == previous.as_raw() {
+                // swap was successful
+                break;
+            }
+        }
+    }
+}

components/salsa-2022/src/lib.rs

@@ -9,6 +9,7 @@ pub mod function;
 pub mod hash;
 pub mod id;
 pub mod ingredient;
+pub mod ingredient_list;
 pub mod input;
 pub mod input_field;
 pub mod interned;

components/salsa-2022/src/salsa_struct.rs

@@ -1,3 +1,5 @@
-use crate::Database;
+use crate::{Database, IngredientIndex};
 
-pub trait SalsaStructInDb<DB: ?Sized + Database> {}
+pub trait SalsaStructInDb<DB: ?Sized + Database> {
+    fn register_dependent_fn(db: &DB, index: IngredientIndex);
+}

components/salsa-2022/src/tracked_struct.rs

@@ -1,6 +1,12 @@
+use std::sync::Arc;
+
+use arc_swap::{ArcSwap, ArcSwapOption, AsRaw};
+
 use crate::{
     cycle::CycleRecoveryStrategy,
     ingredient::{Ingredient, MutIngredient},
+    ingredient_list::IngredientList,
+    input,
     interned::{InternedData, InternedId, InternedIngredient},
     key::{DatabaseKeyIndex, DependencyIndex},
     runtime::{local_state::QueryOrigin, Runtime},
@@ -33,6 +39,14 @@ where
     Data: TrackedStructData,
 {
     interned: InternedIngredient<Id, TrackedStructKey<Data>>,
+
+    /// A list of each tracked function `f` whose key is this
+    /// tracked struct.
+    ///
+    /// Whenever an instance `i` of this struct is deleted,
+    /// each of these functions will be notified
+    /// so they can remove any data tied to that instance.
+    dependent_fns: IngredientList,
 }
 
 #[derive(Debug, PartialEq, Eq, Hash, PartialOrd, Ord, Copy, Clone)]
@@ -53,6 +67,7 @@ where
     pub fn new(index: IngredientIndex) -> Self {
         Self {
             interned: InternedIngredient::new(index),
+            dependent_fns: IngredientList::new(),
         }
     }
 
@@ -99,6 +114,13 @@ where
             self.interned.delete_index(id);
         }
     }
+
+    /// Adds a dependent function (one keyed by this tracked struct) to our list.
+    /// When instances of this struct are deleted, these dependent functions
+    /// will be notified.
+    pub fn register_dependent_fn(&self, index: IngredientIndex) {
+        self.dependent_fns.push(index);
+    }
 }
 
 impl<DB: ?Sized, Id, Data> Ingredient<DB> for TrackedStructIngredient<Id, Data>