starknet_committer,starknet_patricia: layout-dependent FilledTree serialization

ArielElp · ArielElp · commit 1eebe4f19a4c · 2025-12-31T16:31:21.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::{NodeLayout, TrieType};
 use starknet_patricia::patricia_merkle_tree::filled_tree::tree::{FilledTree, FilledTreeImpl};
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{
     Leaf,
@@ -15,32 +16,32 @@ use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::{
     UpdatedSkeletonTree,
     UpdatedSkeletonTreeImpl,
 };
-use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
+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,
 ) -> FilledTreeImpl<L>
 where
     TH: TreeHashFunction<L> + 'static,
-    L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext> + 'static,
+    L: Leaf + 'static,
+    Layout: for<'a> NodeLayout<'a, 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::<L, Layout>(
         storage,
         root_hash,
         sorted_leaf_indices,
         &config,
         &leaf_modifications,
-        &EmptyKeyContext,
+        key_context,
     )
     .await
     .expect("Failed to create the original skeleton tree");
@@ -62,28 +63,38 @@ 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 + HasStaticPrefix<KeyContext = EmptyKeyContext> + 'static,
+    L: Leaf + 'static,
+    Layout: for<'a> NodeLayout<'a, L> + 'static,
     TH: TreeHashFunction<L> + 'static,
 >(
     leaf_modifications: LeafModifications<L>,
     storage: &mut MapStorage,
     root_hash: HashOutput,
     config: impl OriginalSkeletonTreeConfig,
+    // TODO(Ariel): Move trie_type to the config trait or to the leaf type.
+    trie_type: TrieType,
 ) -> 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>>();
 
-    let filled_tree =
-        tree_computation_flow::<L, TH>(leaf_modifications, storage, root_hash, config).await;
+    let key_context = Layout::generate_key_context(trie_type);
+    let filled_tree = tree_computation_flow::<L, Layout, TH>(
+        leaf_modifications,
+        storage,
+        root_hash,
+        config,
+        &key_context,
+    )
+    .await;
 
     let hash_result = filled_tree.get_root_hash();
 
@@ -92,7 +103,7 @@ pub async fn single_tree_flow_test<
     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(&EmptyKeyContext).unwrap());
+    let json_storage = &json!(filled_tree.serialize::<L, 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
@@ -142,12 +142,21 @@ impl<S: Storage> ForestWriter for FactsDb<S> {
 
         // Storage tries.
         for tree in filled_forest.storage_tries.values() {
-            serialized_forest.extend(tree.serialize(&EmptyKeyContext)?);
+            serialized_forest
+                .extend(tree.serialize::<StarknetStorageValue, FactsNodeLayout>(&EmptyKeyContext)?);
         }
 
         // Contracts and classes tries.
-        serialized_forest.extend(filled_forest.contracts_trie.serialize(&EmptyKeyContext)?);
-        serialized_forest.extend(filled_forest.classes_trie.serialize(&EmptyKeyContext)?);
+        serialized_forest.extend(
+            filled_forest
+                .contracts_trie
+                .serialize::<ContractState, FactsNodeLayout>(&EmptyKeyContext)?,
+        );
+        serialized_forest.extend(
+            filled_forest
+                .classes_trie
+                .serialize::<CompiledClassHash, FactsNodeLayout>(&EmptyKeyContext)?,
+        );
 
         Ok(serialized_forest)
     }
diff --git a/crates/starknet_committer/src/db/index_db/db.rs b/crates/starknet_committer/src/db/index_db/db.rs
@@ -14,7 +14,7 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
 use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::hash_function::TreeHashFunction;
-use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
+use starknet_patricia_storage::db_object::HasStaticPrefix;
 use starknet_patricia_storage::errors::SerializationResult;
 use starknet_patricia_storage::storage_trait::{DbHashMap, Storage};
 
@@ -138,13 +138,25 @@ impl<S: Storage> ForestWriter for IndexDb<S> {
         let mut serialized_forest = DbHashMap::new();
 
         // TODO(Ariel): use a different key context when FilledForest is generic over leaf types.
-        for tree in filled_forest.storage_tries.values() {
-            serialized_forest.extend(tree.serialize(&EmptyKeyContext)?);
+        for (contract_address, tree) in &filled_forest.storage_tries {
+            serialized_forest.extend(
+                tree.serialize::<IndexLayoutStarknetStorageValue, IndexNodeLayout>(
+                    &TrieType::StorageTrie(*contract_address),
+                )?,
+            );
         }
 
         // Contracts and classes tries.
-        serialized_forest.extend(filled_forest.contracts_trie.serialize(&EmptyKeyContext)?);
-        serialized_forest.extend(filled_forest.classes_trie.serialize(&EmptyKeyContext)?);
+        serialized_forest.extend(
+            filled_forest
+                .contracts_trie
+                .serialize::<IndexLayoutContractState, IndexNodeLayout>(&TrieType::ContractsTrie)?,
+        );
+        serialized_forest.extend(
+            filled_forest.classes_trie.serialize::<IndexLayoutCompiledClassHash, IndexNodeLayout>(
+                &TrieType::ClassesTrie,
+            )?,
+        );
 
         Ok(serialized_forest)
     }
diff --git a/crates/starknet_committer_and_os_cli/benches/main.rs b/crates/starknet_committer_and_os_cli/benches/main.rs
@@ -12,6 +12,7 @@ use std::collections::HashMap;
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 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;
@@ -43,14 +44,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),
-                    ),
-                );
+                runtime.block_on(tree_computation_flow::<
+                    StarknetStorageValue,
+                    FactsNodeLayout,
+                    TreeHashFunctionImpl,
+                >(
+                    leaf_modifications_input,
+                    &mut storage,
+                    root_hash,
+                    OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(false),
+                    &EmptyKeyContext,
+                ));
             },
             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,
@@ -21,6 +22,7 @@ use starknet_committer::hash_function::hash::{
 use starknet_committer::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use starknet_committer::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
 use starknet_committer::patricia_merkle_tree::types::CompiledClassHash;
+use starknet_patricia::db_layout::TrieType;
 use starknet_patricia::patricia_merkle_tree::filled_tree::node::FactDbFilledNode;
 use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     BinaryData,
@@ -163,11 +165,16 @@ impl PythonTestRunner for CommitterPythonTestRunner {
                 let TreeFlowInput { leaf_modifications, mut storage, root_hash } =
                     serde_json::from_str(Self::non_optional_input(input)?)?;
                 // 2. Run the test.
-                let output = single_tree_flow_test::<StarknetStorageValue, TreeHashFunctionImpl>(
+                let output = single_tree_flow_test::<
+                    StarknetStorageValue,
+                    FactsNodeLayout,
+                    TreeHashFunctionImpl,
+                >(
                     leaf_modifications,
                     &mut storage,
                     root_hash,
                     OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(false),
+                    &TrieType::StorageTrie(ContractAddress::from(1_u128)),
                 )
                 .await;
                 // 3. Serialize and return output.
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
@@ -4,10 +4,13 @@ use std::fs;
 use clap::Error;
 use serde::{Deserialize, Deserializer};
 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 starknet_patricia::db_layout::TrieType;
 use tempfile::NamedTempFile;
 
 use crate::committer_cli::commands::commit;
@@ -106,13 +109,15 @@ pub async fn test_regression_single_tree() {
 
     let start = std::time::Instant::now();
     // Benchmark the single tree flow test.
-    let output = single_tree_flow_test::<StarknetStorageValue, TreeHashFunctionImpl>(
-        leaf_modifications,
-        &mut storage,
-        root_hash,
-        OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(false),
-    )
-    .await;
+    let output =
+        single_tree_flow_test::<StarknetStorageValue, FactsNodeLayout, TreeHashFunctionImpl>(
+            leaf_modifications,
+            &mut storage,
+            root_hash,
+            OriginalSkeletonTrieConfig::new_for_classes_or_storage_trie(false),
+            &TrieType::StorageTrie(ContractAddress::from(1_u128)),
+        )
+        .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/tree.rs b/crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/tree.rs
