starknet_committer,starknet_patricia: layout-dependent FilledTree serialization

ArielElp · ArielElp · commit e91d422be987 · 2026-01-14T13:29:51.000+02:00
diff --git a/crates/starknet_committer/src/db/external_test_utils.rs b/crates/starknet_committer/src/db/external_test_utils.rs
@@ -2,6 +2,7 @@ use std::collections::HashMap;
 
 use serde_json::json;
 use starknet_api::hash::HashOutput;
+use starknet_patricia::db_layout::NodeLayoutFor;
 use starknet_patricia::patricia_merkle_tree::filled_tree::tree::{FilledTree, FilledTreeImpl};
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{
     Leaf,
@@ -18,29 +19,28 @@ use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::{
 use starknet_patricia_storage::db_object::HasStaticPrefix;
 use starknet_patricia_storage::map_storage::MapStorage;
 
-use crate::db::facts_db::db::FactsNodeLayout;
 use crate::db::trie_traversal::create_original_skeleton_tree;
 
-// TODO(Ariel, 14/12/2025): make this generic over the layout.
-pub async fn tree_computation_flow<L, TH>(
+pub async fn tree_computation_flow<L, Layout, TH>(
     leaf_modifications: LeafModifications<L>,
     storage: &mut MapStorage,
     root_hash: HashOutput,
     config: impl OriginalSkeletonTreeConfig,
-    key_context: &<L as HasStaticPrefix>::KeyContext,
+    key_context: &<Layout::DbLeaf as HasStaticPrefix>::KeyContext,
 ) -> FilledTreeImpl<L>
 where
     TH: TreeHashFunction<L> + 'static,
     L: Leaf + 'static,
+    Layout: NodeLayoutFor<L> + 'static,
 {
     let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
     let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices);
-    let mut original_skeleton = create_original_skeleton_tree::<L, FactsNodeLayout>(
+    let mut original_skeleton = create_original_skeleton_tree::<Layout::DbLeaf, Layout>(
         storage,
         root_hash,
         sorted_leaf_indices,
         &config,
-        &leaf_modifications,
+        &leaf_modifications.iter().map(|(k, v)| (*k, v.clone().into())).collect(),
         None,
         key_context,
     )
@@ -64,27 +64,36 @@ where
     )
     .expect("Failed to create the updated skeleton tree");
 
-    FilledTreeImpl::<L>::create_with_existing_leaves::<TH>(updated_skeleton, leaf_modifications)
+    FilledTreeImpl::create_with_existing_leaves::<TH>(updated_skeleton, leaf_modifications)
         .await
         .expect("Failed to create the filled tree")
 }
 
-pub async fn single_tree_flow_test<L: Leaf + 'static, TH: TreeHashFunction<L> + 'static>(
+pub async fn single_tree_flow_test<
+    L: Leaf + 'static,
+    Layout: for<'a> NodeLayoutFor<L> + 'static,
+    TH: TreeHashFunction<L> + 'static,
+>(
     leaf_modifications: LeafModifications<L>,
     storage: &mut MapStorage,
     root_hash: HashOutput,
     config: impl OriginalSkeletonTreeConfig,
-    key_context: &<L as HasStaticPrefix>::KeyContext,
+    key_context: &<Layout::DbLeaf as HasStaticPrefix>::KeyContext,
 ) -> String {
     // Move from leaf number to actual index.
     let leaf_modifications = leaf_modifications
         .into_iter()
         .map(|(k, v)| (NodeIndex::FIRST_LEAF + k, v))
-        .collect::<LeafModifications<L>>();
+        .collect::<LeafModifications<_>>();
 
-    let filled_tree =
-        tree_computation_flow::<L, TH>(leaf_modifications, storage, root_hash, config, key_context)
-            .await;
+    let filled_tree: FilledTreeImpl<L> = tree_computation_flow::<L, Layout, TH>(
+        leaf_modifications,
+        storage,
+        root_hash,
+        config,
+        key_context,
+    )
+    .await;
 
     let hash_result = filled_tree.get_root_hash();
 
@@ -93,7 +102,7 @@ pub async fn single_tree_flow_test<L: Leaf + 'static, TH: TreeHashFunction<L> +
     let json_hash = &json!(hash_result.0.to_hex_string());
     result_map.insert("root_hash", json_hash);
     // Serlialize the storage modifications.
-    let json_storage = &json!(filled_tree.serialize(key_context).unwrap());
+    let json_storage = &json!(filled_tree.serialize::<Layout>(key_context).unwrap());
     result_map.insert("storage_changes", json_storage);
     serde_json::to_string(&result_map).expect("serialization failed")
 }
diff --git a/crates/starknet_committer/src/db/facts_db/db.rs b/crates/starknet_committer/src/db/facts_db/db.rs
@@ -6,17 +6,16 @@ use starknet_api::hash::HashOutput;
 use starknet_patricia::db_layout::{NodeLayout, NodeLayoutFor};
 use starknet_patricia::patricia_merkle_tree::filled_tree::node::{FactDbFilledNode, FilledNode};
 use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::FactNodeDeserializationContext;
-use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
-use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
+use starknet_patricia_storage::db_object::{DBObject, HasStaticPrefix};
 use starknet_patricia_storage::errors::SerializationResult;
 use starknet_patricia_storage::map_storage::MapStorage;
 use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, Storage};
 
 use crate::block_committer::input::{ReaderConfig, StarknetStorageValue};
 use crate::db::facts_db::types::{FactsDbInitialRead, FactsSubTree};
-use crate::db::forest_trait::{read_forest, ForestReader, ForestWriter};
+use crate::db::forest_trait::{read_forest, serialize_forest, ForestReader, ForestWriter};
 use crate::forest::filled_forest::FilledForest;
 use crate::forest::forest_errors::ForestResult;
 use crate::forest::original_skeleton_forest::{ForestSortedIndices, OriginalSkeletonForest};
@@ -111,18 +110,7 @@ impl<S: Storage> ForestReader<FactsDbInitialRead> for FactsDb<S> {
 #[async_trait]
 impl<S: Storage> ForestWriter for FactsDb<S> {
     fn serialize_forest(filled_forest: &FilledForest) -> SerializationResult<DbHashMap> {
-        let mut serialized_forest = DbHashMap::new();
-
-        // Storage tries.
-        for (contract_address, tree) in &filled_forest.storage_tries {
-            serialized_forest.extend(tree.serialize(contract_address)?);
-        }
-
-        // Contracts and classes tries.
-        serialized_forest.extend(filled_forest.contracts_trie.serialize(&EmptyKeyContext)?);
-        serialized_forest.extend(filled_forest.classes_trie.serialize(&EmptyKeyContext)?);
-
-        Ok(serialized_forest)
+        serialize_forest::<FactsNodeLayout>(filled_forest)
     }
 
     async fn write_updates(&mut self, updates: DbHashMap) -> usize {
diff --git a/crates/starknet_committer/src/db/forest_trait.rs b/crates/starknet_committer/src/db/forest_trait.rs
@@ -4,6 +4,7 @@ use async_trait::async_trait;
 use serde::{Deserialize, Serialize};
 use starknet_api::core::ContractAddress;
 use starknet_patricia::db_layout::NodeLayoutFor;
+use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::LeafModifications;
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
 use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
@@ -118,6 +119,36 @@ where
     ))
 }
 
+/// Helper function containing layout-common write logic.
+pub(crate) fn serialize_forest<Layout>(
+    filled_forest: &FilledForest,
+) -> SerializationResult<DbHashMap>
+where
+    Layout: NodeLayoutFor<StarknetStorageValue>
+        + NodeLayoutFor<ContractState>
+        + NodeLayoutFor<CompiledClassHash>,
+    <Layout as NodeLayoutFor<StarknetStorageValue>>::DbLeaf:
+        HasStaticPrefix<KeyContext = ContractAddress>,
+    <Layout as NodeLayoutFor<ContractState>>::DbLeaf: HasStaticPrefix<KeyContext = EmptyKeyContext>,
+    <Layout as NodeLayoutFor<CompiledClassHash>>::DbLeaf:
+        HasStaticPrefix<KeyContext = EmptyKeyContext>,
+{
+    let mut serialized_forest = DbHashMap::new();
+
+    // Storage tries.
+    for (contract_address, tree) in &filled_forest.storage_tries {
+        serialized_forest.extend(tree.serialize::<Layout>(contract_address)?);
+    }
+
+    // Contracts trie.
+    serialized_forest.extend(filled_forest.contracts_trie.serialize::<Layout>(&EmptyKeyContext)?);
+
+    // Classes trie.
+    serialized_forest.extend(filled_forest.classes_trie.serialize::<Layout>(&EmptyKeyContext)?);
+
+    Ok(serialized_forest)
+}
+
 #[async_trait]
 pub trait ForestWriter: Send {
     /// Serializes a filled forest into a hash map.
diff --git a/crates/starknet_committer/src/db/index_db/db.rs b/crates/starknet_committer/src/db/index_db/db.rs
@@ -5,17 +5,16 @@ use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
 use starknet_patricia::db_layout::{NodeLayout, NodeLayoutFor};
 use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
-use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
 use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunction;
-use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
+use starknet_patricia_storage::db_object::{DBObject, HasStaticPrefix};
 use starknet_patricia_storage::errors::SerializationResult;
 use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, Storage};
 
 use crate::block_committer::input::{ReaderConfig, StarknetStorageValue};
 use crate::db::facts_db::types::FactsDbInitialRead;
-use crate::db::forest_trait::{read_forest, ForestReader, ForestWriter};
+use crate::db::forest_trait::{read_forest, serialize_forest, ForestReader, ForestWriter};
 use crate::db::index_db::leaves::{
     IndexLayoutCompiledClassHash,
     IndexLayoutContractState,
@@ -114,17 +113,7 @@ impl<S: Storage> ForestReader<FactsDbInitialRead> for IndexDb<S> {
 #[async_trait]
 impl<S: Storage> ForestWriter for IndexDb<S> {
     fn serialize_forest(filled_forest: &FilledForest) -> SerializationResult<DbHashMap> {
-        let mut serialized_forest = DbHashMap::new();
-
-        for (contract_address, tree) in &filled_forest.storage_tries {
-            serialized_forest.extend(tree.serialize(contract_address)?);
-        }
-
-        // Contracts and classes tries.
-        serialized_forest.extend(filled_forest.contracts_trie.serialize(&EmptyKeyContext)?);
-        serialized_forest.extend(filled_forest.classes_trie.serialize(&EmptyKeyContext)?);
-
-        Ok(serialized_forest)
+        serialize_forest::<IndexNodeLayout>(filled_forest)
     }
 
     async fn write_updates(&mut self, updates: DbHashMap) -> usize {
diff --git a/crates/starknet_committer_and_os_cli/benches/main.rs b/crates/starknet_committer_and_os_cli/benches/main.rs
@@ -13,6 +13,7 @@ use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use starknet_api::core::ContractAddress;
 use starknet_committer::block_committer::input::StarknetStorageValue;
 use starknet_committer::db::external_test_utils::tree_computation_flow;
+use starknet_committer::db::facts_db::db::FactsNodeLayout;
 use starknet_committer::hash_function::hash::TreeHashFunctionImpl;
 use starknet_committer::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
 use starknet_committer_and_os_cli::committer_cli::commands::commit;
@@ -46,15 +47,17 @@ pub fn single_tree_flow_benchmark(criterion: &mut Criterion) {
         benchmark.iter_batched(
             || leaf_modifications.clone(),
             |leaf_modifications_input| {
-                runtime.block_on(
-                    tree_computation_flow::<StarknetStorageValue, TreeHashFunctionImpl>(
-                        leaf_modifications_input,
-                        &mut storage,
-                        root_hash,
-                        OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(false),
-                        &dummy_contract_address,
-                    ),
-                );
+                runtime.block_on(tree_computation_flow::<
+                    StarknetStorageValue,
+                    FactsNodeLayout,
+                    TreeHashFunctionImpl,
+                >(
+                    leaf_modifications_input,
+                    &mut storage,
+                    root_hash,
+                    OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(false),
+                    &dummy_contract_address,
+                ));
             },
             BatchSize::LargeInput,
         )
diff --git a/crates/starknet_committer_and_os_cli/src/committer_cli/tests/python_tests.rs b/crates/starknet_committer_and_os_cli/src/committer_cli/tests/python_tests.rs
@@ -12,6 +12,7 @@ use starknet_committer::block_committer::input::{
 };
 use starknet_committer::block_committer::random_structs::DummyRandomValue;
 use starknet_committer::db::external_test_utils::single_tree_flow_test;
+use starknet_committer::db::facts_db::db::FactsNodeLayout;
 use starknet_committer::forest::filled_forest::FilledForest;
 use starknet_committer::hash_function::hash::{
     TreeHashFunctionImpl,
@@ -164,7 +165,11 @@ impl PythonTestRunner for CommitterPythonTestRunner {
                     serde_json::from_str(Self::non_optional_input(input)?)?;
                 // 2. Run the test.
                 let dummy_contract_address = ContractAddress::from(0_u128);
-                let output = single_tree_flow_test::<StarknetStorageValue, TreeHashFunctionImpl>(
+                let output = single_tree_flow_test::<
+                    StarknetStorageValue,
+                    FactsNodeLayout,
+                    TreeHashFunctionImpl,
+                >(
                     leaf_modifications,
                     &mut storage,
                     root_hash,
diff --git a/crates/starknet_committer_and_os_cli/src/committer_cli/tests/regression_tests.rs b/crates/starknet_committer_and_os_cli/src/committer_cli/tests/regression_tests.rs
@@ -7,6 +7,7 @@ use serde_json::{Map, Value};
 use starknet_api::core::ContractAddress;
 use starknet_committer::block_committer::input::StarknetStorageValue;
 use starknet_committer::db::external_test_utils::single_tree_flow_test;
+use starknet_committer::db::facts_db::db::FactsNodeLayout;
 use starknet_committer::hash_function::hash::TreeHashFunctionImpl;
 use starknet_committer::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
 use tempfile::NamedTempFile;
@@ -108,14 +109,15 @@ pub async fn test_regression_single_tree() {
     let start = std::time::Instant::now();
     // Benchmark the single tree flow test.
     let dummy_contract_address = ContractAddress::from(0_u128);
-    let output = single_tree_flow_test::<StarknetStorageValue, TreeHashFunctionImpl>(
-        leaf_modifications,
-        &mut storage,
-        root_hash,
-        OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(false),
-        &dummy_contract_address,
-    )
-    .await;
+    let output =
+        single_tree_flow_test::<StarknetStorageValue, FactsNodeLayout, TreeHashFunctionImpl>(
+            leaf_modifications,
+            &mut storage,
+            root_hash,
+            OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(false),
+            &dummy_contract_address,
+        )
+        .await;
     let execution_time = std::time::Instant::now() - start;
 
     // Assert correctness of the output of the single tree flow test.
diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/node.rs b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/node.rs
@@ -14,3 +14,9 @@ pub struct FilledNode<L: Leaf, ChildData> {
 // starknet_committer. This can happen after serialization of FilledTree is made generic in the
 // layout.
 pub type FactDbFilledNode<L> = FilledNode<L, HashOutput>;
+
+/// A node in a filled tree, where all the hashes were computed. Used in the `FilledTree` trait.
+///
+/// While the same underlying type as `FactDbFilledNode`, this alias is not used in the context of
+/// the DB-representation of the node.
+pub type HashFilledNode<L> = FilledNode<L, HashOutput>;
diff --git a/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree.rs b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree.rs
