avoid allocating intermediary serialized dependency edges

Author: ibraheemdev

src/accumulator.rs

@@ -8,6 +8,7 @@ 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;
@@ -111,7 +112,12 @@ impl<A: Accumulator> Ingredient for IngredientImpl<A> {
         panic!("nothing should ever depend on an accumulator directly")
     }
 
-    fn minimum_serialized_edges(&self, _zalsa: &Zalsa, _edge: QueryEdge) -> Vec<QueryEdge> {
+    fn collect_minimum_serialized_edges(
+        &self,
+        _zalsa: &Zalsa,
+        _edge: QueryEdge,
+        _serialized_edges: &mut FxIndexSet<QueryEdge>,
+    ) {
         panic!("nothing should ever depend on an accumulator directly")
     }
 

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};
@@ -288,26 +289,27 @@ where
         self.maybe_changed_after(db, input, revision, cycle_heads)
     }
 
-    fn minimum_serialized_edges(&self, zalsa: &Zalsa, edge: QueryEdge) -> Vec<QueryEdge> {
+    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 Vec::new();
+            return;
         };
 
         let origin = memo.revisions.origin.as_ref();
 
         // Collect the minimum dependency tree.
-        origin
-            .edges()
-            .iter()
-            .flat_map(|&edge| {
-                let dependency = zalsa.lookup_ingredient(edge.key().ingredient_index());
-                dependency.minimum_serialized_edges(zalsa, edge)
-            })
-            .collect()
+        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`.
@@ -588,7 +590,7 @@ mod persistence {
                 flattened_edges.insert(edge);
             } else {
                 // Otherwise, serialize the minimum edges necessary to cover the dependency.
-                flattened_edges.extend(dependency.minimum_serialized_edges(zalsa, edge));
+                dependency.collect_minimum_serialized_edges(zalsa, edge, &mut flattened_edges);
             }
         }
 

src/ingredient.rs

@@ -6,6 +6,7 @@ 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;
@@ -55,14 +56,19 @@ pub trait Ingredient: Any + std::fmt::Debug + Send + Sync {
         cycle_heads: &mut VerifyCycleHeads,
     ) -> VerifyResult;
 
-    /// Returns the minimum edges necessary to serialize a given dependency edge on this ingredient,
+    /// 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 minimum_serialized_edges(&self, zalsa: &Zalsa, edge: QueryEdge) -> Vec<QueryEdge>;
+    fn collect_minimum_serialized_edges(
+        &self,
+        zalsa: &Zalsa,
+        edge: QueryEdge,
+        serialized_edges: &mut FxIndexSet<QueryEdge>,
+    );
 
     /// Returns information about the current provisional status of `input`.
     ///

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};
@@ -278,7 +279,12 @@ impl<C: Configuration> Ingredient for IngredientImpl<C> {
         panic!("nothing should ever depend on an input struct directly")
     }
 
-    fn minimum_serialized_edges(&self, _zalsa: &Zalsa, _edge: QueryEdge) -> Vec<QueryEdge> {
+    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")
     }
 

src/input/input_field.rs

@@ -2,6 +2,7 @@ use std::fmt;
 use std::marker::PhantomData;
 
 use crate::function::{VerifyCycleHeads, VerifyResult};
+use crate::hash::FxIndexSet;
 use crate::ingredient::Ingredient;
 use crate::input::{Configuration, IngredientImpl, Value};
 use crate::sync::Arc;
@@ -62,14 +63,19 @@ where
         VerifyResult::changed_if(value.revisions[self.field_index] > revision)
     }
 
-    fn minimum_serialized_edges(&self, _zalsa: &Zalsa, edge: QueryEdge) -> Vec<QueryEdge> {
+    fn collect_minimum_serialized_edges(
+        &self,
+        _zalsa: &Zalsa,
+        edge: QueryEdge,
+        serialized_edges: &mut FxIndexSet<QueryEdge>,
+    ) {
         assert!(
             C::PERSIST,
             "the inputs of a persistable tracked function must be persistable"
         );
 
         // Input dependencies are the leaves of the minimum dependency tree.
-        Vec::from([edge])
+        serialized_edges.insert(edge);
     }
 
     fn fmt_index(&self, index: crate::Id, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {

src/interned.rs

@@ -12,6 +12,7 @@ use rustc_hash::FxBuildHasher;
 
 use crate::durability::Durability;
 use crate::function::{VerifyCycleHeads, VerifyResult};
+use crate::hash::FxIndexSet;
 use crate::id::{AsId, FromId};
 use crate::ingredient::Ingredient;
 use crate::plumbing::{self, Jar, ZalsaLocal};
@@ -897,18 +898,22 @@ where
         VerifyResult::unchanged()
     }
 
-    fn minimum_serialized_edges(&self, _zalsa: &Zalsa, edge: QueryEdge) -> Vec<QueryEdge> {
+    fn collect_minimum_serialized_edges(
+        &self,
+        _zalsa: &Zalsa,
+        edge: QueryEdge,
+        serialized_edges: &mut FxIndexSet<QueryEdge>,
+    ) {
         if C::PERSIST {
             // If the interned struct is being persisted, it may be reachable through transitive queries.
             // Additionally, interned struct dependencies are impure in that garbage collection can
             // invalidate a dependency without a base input necessarily being updated. Thus, we must
             // preserve the transitive dependency on the interned struct.
-            Vec::from([edge])
-        } else {
-            // The dependency is covered by the base inputs, as the interned struct itself is
-            // not being persisted.
-            Vec::new()
+            serialized_edges.insert(edge);
         }
+
+        // Otherwise, the dependency is covered by the base inputs, as the interned struct itself is
+        // not being persisted.
     }
 
     fn debug_name(&self) -> &'static str {

src/tracked_struct.rs

@@ -12,6 +12,7 @@ use thin_vec::ThinVec;
 use tracked_field::FieldIngredientImpl;
 
 use crate::function::{VerifyCycleHeads, VerifyResult};
+use crate::hash::FxIndexSet;
 use crate::id::{AsId, FromId};
 use crate::ingredient::{Ingredient, Jar};
 use crate::key::DatabaseKeyIndex;
@@ -943,7 +944,19 @@ where
         panic!("nothing should ever depend on a tracked struct directly")
     }
 
-    fn minimum_serialized_edges(&self, _zalsa: &Zalsa, _edge: QueryEdge) -> Vec<QueryEdge> {
+    fn collect_minimum_serialized_edges(
+        &self,
+        _zalsa: &Zalsa,
+        _edge: QueryEdge,
+        _serialized_edges: &mut FxIndexSet<QueryEdge>,
+    ) {
+        // Note that tracked structs are referenced by the identity map, but that
+        // only matters if we are serializing the creating query, in which case
+        // the dependency edge will be serialized directly.
+        //
+        // TODO: We could flatten the identity map here if the tracked struct is being
+        // persisted, in order to more aggressively preserve the tracked struct IDs if
+        // the transitive query is re-executed.
         panic!("nothing should ever depend on a tracked struct directly")
     }
 

src/tracked_struct/tracked_field.rs

@@ -1,6 +1,7 @@
 use std::marker::PhantomData;
 
 use crate::function::{VerifyCycleHeads, VerifyResult};
+use crate::hash::FxIndexSet;
 use crate::ingredient::Ingredient;
 use crate::sync::Arc;
 use crate::table::memo::MemoTableTypes;
@@ -68,10 +69,14 @@ where
         VerifyResult::changed_if(field_changed_at > revision)
     }
 
-    fn minimum_serialized_edges(&self, _zalsa: &Zalsa, _edge: QueryEdge) -> Vec<QueryEdge> {
+    fn collect_minimum_serialized_edges(
+        &self,
+        _zalsa: &Zalsa,
+        _edge: QueryEdge,
+        _serialized_edges: &mut FxIndexSet<QueryEdge>,
+    ) {
         // Tracked structs do not have transitive dependencies, and their dependencies are covered by
         // the base inputs.
-        Vec::new()
     }
 
     fn fmt_index(&self, index: crate::Id, fmt: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {