Flatten unserializable query dependencies (#975)

* flatten unserializable query dependencies

* avoid allocating intermediary serialized dependency edges

Author: ibraheemdev

components/salsa-macro-rules/src/setup_tracked_fn.rs

@@ -172,7 +172,7 @@ macro_rules! setup_tracked_fn {
                             line: line!(),
                         };
                         const DEBUG_NAME: &'static str = concat!($(stringify!($self_ty), "::",)? stringify!($fn_name), "::interned_arguments");
-                        const PERSIST: bool = true;
+                        const PERSIST: bool = $persist;
 
                         type Fields<$db_lt> = ($($interned_input_ty),*);
 

src/accumulator.rs

@@ -8,11 +8,13 @@ use std::panic::UnwindSafe;
 use accumulated::{Accumulated, AnyAccumulated};
 
 use crate::function::{VerifyCycleHeads, VerifyResult};
+use crate::hash::FxIndexSet;
 use crate::ingredient::{Ingredient, Jar};
 use crate::plumbing::ZalsaLocal;
 use crate::sync::Arc;
 use crate::table::memo::MemoTableTypes;
 use crate::zalsa::{IngredientIndex, JarKind, Zalsa};
+use crate::zalsa_local::QueryEdge;
 use crate::{Database, Id, Revision};
 
 mod accumulated;
@@ -110,6 +112,15 @@ impl<A: Accumulator> Ingredient for IngredientImpl<A> {
         panic!("nothing should ever depend on an accumulator directly")
     }
 
+    fn collect_minimum_serialized_edges(
+        &self,
+        _zalsa: &Zalsa,
+        _edge: QueryEdge,
+        _serialized_edges: &mut FxIndexSet<QueryEdge>,
+    ) {
+        panic!("nothing should ever depend on an accumulator directly")
+    }
+
     fn debug_name(&self) -> &'static str {
         A::DEBUG_NAME
     }

src/function.rs

@@ -13,6 +13,7 @@ use crate::cycle::{
 use crate::database::RawDatabase;
 use crate::function::delete::DeletedEntries;
 use crate::function::sync::{ClaimResult, SyncTable};
+use crate::hash::FxIndexSet;
 use crate::ingredient::{Ingredient, WaitForResult};
 use crate::key::DatabaseKeyIndex;
 use crate::plumbing::{self, MemoIngredientMap};
@@ -22,7 +23,7 @@ use crate::table::memo::MemoTableTypes;
 use crate::table::Table;
 use crate::views::DatabaseDownCaster;
 use crate::zalsa::{IngredientIndex, JarKind, MemoIngredientIndex, Zalsa};
-use crate::zalsa_local::QueryOriginRef;
+use crate::zalsa_local::{QueryEdge, QueryOriginRef};
 use crate::{Id, Revision};
 
 #[cfg(feature = "accumulator")]
@@ -277,7 +278,7 @@ where
 
     unsafe fn maybe_changed_after(
         &self,
-        _zalsa: &crate::zalsa::Zalsa,
+        _zalsa: &Zalsa,
         db: RawDatabase<'_>,
         input: Id,
         revision: Revision,
@@ -288,6 +289,29 @@ where
         self.maybe_changed_after(db, input, revision, cycle_heads)
     }
 
+    fn collect_minimum_serialized_edges(
+        &self,
+        zalsa: &Zalsa,
+        edge: QueryEdge,
+        serialized_edges: &mut FxIndexSet<QueryEdge>,
+    ) {
+        let input = edge.key().key_index();
+
+        let Some(memo) =
+            self.get_memo_from_table_for(zalsa, input, self.memo_ingredient_index(zalsa, input))
+        else {
+            return;
+        };
+
+        let origin = memo.revisions.origin.as_ref();
+
+        // Collect the minimum dependency tree.
+        for edge in origin.edges() {
+            let dependency = zalsa.lookup_ingredient(edge.key().ingredient_index());
+            dependency.collect_minimum_serialized_edges(zalsa, *edge, serialized_edges)
+        }
+    }
+
     /// Returns `final` only if the memo has the `verified_final` flag set and the cycle recovery strategy is not `FallbackImmediate`.
     ///
     /// Otherwise, the value is still provisional. For both final and provisional, it also
@@ -470,8 +494,10 @@ where
 #[cfg(feature = "persistence")]
 mod persistence {
     use super::{Configuration, IngredientImpl, Memo};
-    use crate::plumbing::{Ingredient, MemoIngredientMap, SalsaStructInDb};
+    use crate::hash::FxIndexSet;
+    use crate::plumbing::{MemoIngredientMap, SalsaStructInDb};
     use crate::zalsa::Zalsa;
+    use crate::zalsa_local::{QueryEdge, QueryOrigin, QueryOriginRef};
     use crate::{Id, IngredientIndex};
 
     use serde::de;
@@ -499,16 +525,6 @@ mod persistence {
 
             let mut map = serializer.serialize_map(None)?;
 
-            for struct_index in
-                <C::SalsaStruct<'_> as SalsaStructInDb>::lookup_ingredient_index(zalsa).iter()
-            {
-                let struct_ingredient = zalsa.lookup_ingredient(struct_index);
-                assert!(
-                    struct_ingredient.is_persistable(),
-                    "the input of a serialized tracked function must be serialized"
-                );
-            }
-
             for entry in <C::SalsaStruct<'_> as SalsaStructInDb>::entries(zalsa) {
                 let memo_ingredient_index = ingredient
                     .memo_ingredient_indices
@@ -521,18 +537,30 @@ mod persistence {
                 );
 
                 if let Some(memo) = memo.filter(|memo| memo.should_serialize()) {
-                    for edge in memo.revisions.origin.as_ref().edges() {
-                        let dependency = zalsa.lookup_ingredient(edge.key().ingredient_index());
-
-                        // TODO: This is not strictly necessary, we only need the transitive input
-                        // dependencies of this query to serialize a valid memo.
-                        assert!(
-                            dependency.is_persistable(),
-                            "attempted to serialize query `{}`, but dependency `{}` is not persistable",
-                            ingredient.debug_name(),
-                            dependency.debug_name()
-                        );
-                    }
+                    // Flatten the dependencies of this query down to the base inputs.
+                    let flattened_origin = match memo.revisions.origin.as_ref() {
+                        QueryOriginRef::Derived(edges) => {
+                            QueryOrigin::derived(flatten_edges(zalsa, edges))
+                        }
+                        QueryOriginRef::DerivedUntracked(edges) => {
+                            QueryOrigin::derived_untracked(flatten_edges(zalsa, edges))
+                        }
+                        QueryOriginRef::Assigned(key) => {
+                            let dependency = zalsa.lookup_ingredient(key.ingredient_index());
+                            assert!(
+                                dependency.is_persistable(),
+                                "specified query `{}` must be persistable",
+                                dependency.debug_name()
+                            );
+
+                            QueryOrigin::assigned(key)
+                        }
+                        QueryOriginRef::FixpointInitial => unreachable!(
+                            "`should_serialize` returns `false` for provisional queries"
+                        ),
+                    };
+
+                    let memo = memo.with_origin(flattened_origin);
 
                     // TODO: Group structs by ingredient index into a nested map.
                     let key = format!(
@@ -541,14 +569,34 @@ mod persistence {
                         entry.key_index().as_bits()
                     );
 
-                    map.serialize_entry(&key, memo)?;
+                    map.serialize_entry(&key, &memo)?;
                 }
             }
 
             map.end()
         }
     }
 
+    // Flatten the dependency edges before serialization.
+    fn flatten_edges(zalsa: &Zalsa, edges: &[QueryEdge]) -> FxIndexSet<QueryEdge> {
+        let mut flattened_edges =
+            FxIndexSet::with_capacity_and_hasher(edges.len(), Default::default());
+
+        for &edge in edges {
+            let dependency = zalsa.lookup_ingredient(edge.key().ingredient_index());
+
+            if dependency.is_persistable() {
+                // If the dependency will be serialized, we can serialize the edge directly.
+                flattened_edges.insert(edge);
+            } else {
+                // Otherwise, serialize the minimum edges necessary to cover the dependency.
+                dependency.collect_minimum_serialized_edges(zalsa, edge, &mut flattened_edges);
+            }
+        }
+
+        flattened_edges
+    }
+
     pub struct DeserializeIngredient<'db, C>
     where
         C: Configuration,

src/function/memo.rs

@@ -359,12 +359,36 @@ mod persistence {
     use crate::function::memo::Memo;
     use crate::function::Configuration;
     use crate::revision::AtomicRevision;
-    use crate::zalsa_local::QueryRevisions;
+    use crate::zalsa_local::persistence::MappedQueryRevisions;
+    use crate::zalsa_local::{QueryOrigin, QueryRevisions};
 
     use serde::ser::SerializeStruct;
     use serde::Deserialize;
 
-    impl<C> serde::Serialize for Memo<'_, C>
+    /// A reference to the fields of a [`Memo`], with its [`QueryRevisions`] transformed.
+    pub(crate) struct MappedMemo<'memo, 'db, C: Configuration> {
+        value: Option<&'memo C::Output<'db>>,
+        verified_at: AtomicRevision,
+        revisions: MappedQueryRevisions<'memo>,
+    }
+
+    impl<'db, C: Configuration> Memo<'db, C> {
+        pub(crate) fn with_origin(&self, origin: QueryOrigin) -> MappedMemo<'_, 'db, C> {
+            let Memo {
+                ref verified_at,
+                ref value,
+                ref revisions,
+            } = *self;
+
+            MappedMemo {
+                value: value.as_ref(),
+                verified_at: AtomicRevision::from(verified_at.load()),
+                revisions: revisions.with_origin(origin),
+            }
+        }
+    }
+
+    impl<C> serde::Serialize for MappedMemo<'_, '_, C>
     where
         C: Configuration,
     {
@@ -386,13 +410,15 @@ mod persistence {
                 }
             }
 
-            let Memo {
+            let MappedMemo {
                 value,
                 verified_at,
                 revisions,
             } = self;
 
-            let value = value.as_ref().expect("attempted to serialize empty memo");
+            let value = value.expect(
+                "attempted to serialize memo where `Memo::should_serialize` returned `false`",
+            );
 
             let mut s = serializer.serialize_struct("Memo", 3)?;
             s.serialize_field("value", &SerializeValue::<C>(value))?;

src/id.rs

@@ -18,7 +18,6 @@ use crate::zalsa::Zalsa;
 /// As an end-user of `Salsa` you will generally not use `Id` directly,
 /// it is wrapped in new types.
 #[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
-#[cfg_attr(feature = "persistence", derive(serde::Serialize, serde::Deserialize))]
 pub struct Id {
     index: NonZeroU32,
     generation: u32,
@@ -133,6 +132,26 @@ impl Hash for Id {
     }
 }
 
+#[cfg(feature = "persistence")]
+impl serde::Serialize for Id {
+    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
+    where
+        S: serde::Serializer,
+    {
+        serde::Serialize::serialize(&self.as_bits(), serializer)
+    }
+}
+
+#[cfg(feature = "persistence")]
+impl<'de> serde::Deserialize<'de> for Id {
+    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
+    where
+        D: serde::Deserializer<'de>,
+    {
+        serde::Deserialize::deserialize(deserializer).map(Self::from_bits)
+    }
+}
+
 impl Debug for Id {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         if self.generation() == 0 {

src/ingredient.rs

@@ -6,12 +6,13 @@ use crate::cycle::{
 };
 use crate::database::RawDatabase;
 use crate::function::{VerifyCycleHeads, VerifyResult};
+use crate::hash::FxIndexSet;
 use crate::runtime::Running;
 use crate::sync::Arc;
 use crate::table::memo::MemoTableTypes;
 use crate::table::Table;
 use crate::zalsa::{transmute_data_mut_ptr, transmute_data_ptr, IngredientIndex, JarKind, Zalsa};
-use crate::zalsa_local::QueryOriginRef;
+use crate::zalsa_local::{QueryEdge, QueryOriginRef};
 use crate::{DatabaseKeyIndex, Id, Revision};
 
 /// A "jar" is a group of ingredients that are added atomically.
@@ -55,6 +56,20 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
         cycle_heads: &mut VerifyCycleHeads,
     ) -> VerifyResult;
 
+    /// Collects the minimum edges necessary to serialize a given dependency edge on this ingredient,
+    /// without necessarily serializing the dependency edge itself.
+    ///
+    /// This generally only returns any transitive input dependencies, i.e. the leaves of the dependency
+    /// tree, as most other fine-grained dependencies are covered by the inputs.
+    ///
+    /// Note that any ingredients returned by this function must be persistable.
+    fn collect_minimum_serialized_edges(
+        &self,
+        zalsa: &Zalsa,
+        edge: QueryEdge,
+        serialized_edges: &mut FxIndexSet<QueryEdge>,
+    );
+
     /// Returns information about the current provisional status of `input`.
     ///
     /// Is it a provisional value or has it been finalized and in which iteration.
@@ -217,7 +232,7 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
         _zalsa: &'db Zalsa,
         _f: &mut dyn FnMut(&dyn erased_serde::Serialize),
     ) {
-        unimplemented!("called `serialize` on ingredient where `is_persistable` returns `false`")
+        unimplemented!("called `serialize` on ingredient where `should_serialize` returns `false`")
     }
 
     /// Deserialize the ingredient.
@@ -227,7 +242,9 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
         _zalsa: &mut Zalsa,
         _deserializer: &mut dyn erased_serde::Deserializer,
     ) -> Result<(), erased_serde::Error> {
-        unimplemented!("called `deserialize` on ingredient where `is_persistable` returns `false`")
+        unimplemented!(
+            "called `deserialize` on ingredient where `should_serialize` returns `false`"
+        )
     }
 }
 

src/input.rs

@@ -9,6 +9,7 @@ pub mod singleton;
 use input_field::FieldIngredientImpl;
 
 use crate::function::{VerifyCycleHeads, VerifyResult};
+use crate::hash::FxIndexSet;
 use crate::id::{AsId, FromId, FromIdWithDb};
 use crate::ingredient::Ingredient;
 use crate::input::singleton::{Singleton, SingletonChoice};
@@ -18,6 +19,7 @@ use crate::sync::Arc;
 use crate::table::memo::{MemoTable, MemoTableTypes};
 use crate::table::{Slot, Table};
 use crate::zalsa::{IngredientIndex, JarKind, Zalsa};
+use crate::zalsa_local::QueryEdge;
 use crate::{Durability, Id, Revision, Runtime};
 
 pub trait Configuration: Any {
@@ -274,7 +276,16 @@ impl<C: Configuration> Ingredient for IngredientImpl<C> {
     ) -> VerifyResult {
         // Input ingredients are just a counter, they store no data, they are immortal.
         // Their *fields* are stored in function ingredients elsewhere.
-        VerifyResult::unchanged()
+        panic!("nothing should ever depend on an input struct directly")
+    }
+
+    fn collect_minimum_serialized_edges(
+        &self,
+        _zalsa: &Zalsa,
+        _edge: QueryEdge,
+        _serialized_edges: &mut FxIndexSet<QueryEdge>,
+    ) {
+        panic!("nothing should ever depend on an input struct directly")
     }
 
     fn debug_name(&self) -> &'static str {