starknet_committer: abstract original tree creation

Author: ArielElp

crates/starknet_committer/src/db/db_layout.rs

@@ -1,7 +1,11 @@
+use starknet_api::hash::HashOutput;
 use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
 use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait;
-use starknet_patricia_storage::db_object::DBObject;
+use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
+use starknet_patricia_storage::db_object::{DBObject, HasStaticPrefix};
+
+use crate::db::index_db::leaves::TrieType;
 
 /// Specifies the trie db layout.
 pub trait NodeLayout<'a, L: Leaf>
@@ -21,4 +25,12 @@ where
             NodeData = Self::NodeData,
             NodeDeserializeContext = Self::DeserializationContext,
         >;
+
+    fn create_subtree(
+        sorted_leaf_indices: SortedLeafIndices<'a>,
+        root_index: NodeIndex,
+        root_hash: HashOutput,
+    ) -> Self::SubTree;
+
+    fn generate_key_context(trie_type: TrieType) -> <L as HasStaticPrefix>::KeyContext;
 }

crates/starknet_committer/src/db/external_test_utils.rs

@@ -18,7 +18,8 @@ use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::{
 use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::map_storage::MapStorage;
 
-use crate::db::facts_db::create_facts_tree::create_original_skeleton_tree;
+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>(
@@ -33,7 +34,7 @@ where
 {
     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(
+    let mut original_skeleton = create_original_skeleton_tree::<L, FactsNodeLayout>(
         storage,
         root_hash,
         sorted_leaf_indices,

crates/starknet_committer/src/db/facts_db/create_facts_tree.rs

@@ -9,59 +9,27 @@ use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::tree::{
 };
 use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait;
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
-use starknet_patricia_storage::db_object::HasStaticPrefix;
+use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::storage_trait::Storage;
 
 use crate::db::facts_db::db::FactsNodeLayout;
 use crate::db::facts_db::types::FactsSubTree;
-use crate::db::trie_traversal::{fetch_nodes, log_warning_for_empty_leaves};
+use crate::db::trie_traversal::fetch_nodes;
 use crate::patricia_merkle_tree::tree::OriginalSkeletonTrieDontCompareConfig;
 
 #[cfg(test)]
 #[path = "create_facts_tree_test.rs"]
 pub mod create_facts_tree_test;
 
-pub async fn create_original_skeleton_tree<'a, L: Leaf>(
-    storage: &mut impl Storage,
-    root_hash: HashOutput,
-    sorted_leaf_indices: SortedLeafIndices<'a>,
-    config: &impl OriginalSkeletonTreeConfig,
-    leaf_modifications: &LeafModifications<L>,
-    key_context: &<L as HasStaticPrefix>::KeyContext,
-) -> OriginalSkeletonTreeResult<OriginalSkeletonTreeImpl<'a>> {
-    if sorted_leaf_indices.is_empty() {
-        return Ok(OriginalSkeletonTreeImpl::create_unmodified(root_hash));
-    }
-    if root_hash == HashOutput::ROOT_OF_EMPTY_TREE {
-        log_warning_for_empty_leaves(
-            sorted_leaf_indices.get_indices(),
-            leaf_modifications,
-            config,
-        )?;
-        return Ok(OriginalSkeletonTreeImpl::create_empty(sorted_leaf_indices));
-    }
-    let main_subtree = FactsSubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash);
-    let mut skeleton_tree = OriginalSkeletonTreeImpl { nodes: HashMap::new(), sorted_leaf_indices };
-    fetch_nodes::<L, FactsNodeLayout>(
-        &mut skeleton_tree,
-        vec![main_subtree],
-        storage,
-        leaf_modifications,
-        config,
-        None,
-        key_context,
-    )
-    .await?;
-    Ok(skeleton_tree)
-}
-
-pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
+pub async fn create_original_skeleton_tree_and_get_previous_leaves<
+    'a,
+    L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext>,
+>(
     storage: &mut impl Storage,
     root_hash: HashOutput,
     sorted_leaf_indices: SortedLeafIndices<'a>,
     leaf_modifications: &LeafModifications<L>,
     config: &impl OriginalSkeletonTreeConfig,
-    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> OriginalSkeletonTreeResult<(OriginalSkeletonTreeImpl<'a>, HashMap<NodeIndex, L>)> {
     if sorted_leaf_indices.is_empty() {
         let unmodified = OriginalSkeletonTreeImpl::create_unmodified(root_hash);
@@ -83,7 +51,7 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
         leaf_modifications,
         config,
         Some(&mut leaves),
-        key_context,
+        &EmptyKeyContext,
     )
     .await?;
     Ok((skeleton_tree, leaves))
@@ -92,11 +60,10 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
 /// Prepares the OS inputs by fetching paths to the given leaves (i.e. their induced Skeleton tree).
 /// Note that ATM, the Rust committer does not manage history and is not used for storage proofs;
 /// Thus, this function assumes facts layout.
-pub async fn get_leaves<'a, L: Leaf>(
+pub async fn get_leaves<'a, L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext>>(
     storage: &mut impl Storage,
     root_hash: HashOutput,
     sorted_leaf_indices: SortedLeafIndices<'a>,
-    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> OriginalSkeletonTreeResult<HashMap<NodeIndex, L>> {
     let config = OriginalSkeletonTrieDontCompareConfig;
     let leaf_modifications = LeafModifications::new();
@@ -106,7 +73,6 @@ pub async fn get_leaves<'a, L: Leaf>(
         sorted_leaf_indices,
         &leaf_modifications,
         &config,
-        key_context,
     )
     .await?;
     Ok(previous_leaves)

crates/starknet_committer/src/db/facts_db/create_facts_tree_test.rs

@@ -22,7 +22,8 @@ use starknet_patricia_storage::map_storage::MapStorage;
 use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue};
 use starknet_types_core::felt::Felt;
 
-use crate::db::facts_db::create_facts_tree::create_original_skeleton_tree;
+use crate::db::facts_db::db::FactsNodeLayout;
+use crate::db::trie_traversal::create_original_skeleton_tree;
 use crate::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
 
 #[tokio::test]
@@ -211,7 +212,7 @@ async fn test_create_tree(
     let config = OriginalSkeletonTrieConfig::new(compare_modified_leaves);
     let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
     let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices);
-    let skeleton_tree = create_original_skeleton_tree(
+    let skeleton_tree = create_original_skeleton_tree::<MockLeaf, FactsNodeLayout>(
         &mut storage,
         root_hash,
         sorted_leaf_indices,

crates/starknet_committer/src/db/facts_db/db.rs

@@ -6,9 +6,9 @@ use starknet_api::hash::HashOutput;
 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::original_skeleton_tree::tree::OriginalSkeletonTreeImpl;
+use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait;
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
-use starknet_patricia_storage::db_object::EmptyKeyContext;
+use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::errors::SerializationResult;
 use starknet_patricia_storage::map_storage::MapStorage;
 use starknet_patricia_storage::storage_trait::{
@@ -19,25 +19,16 @@ use starknet_patricia_storage::storage_trait::{
     Storage,
 };
 
-use crate::block_committer::input::{
-    contract_address_into_node_index,
-    Config,
-    FactsDbInitialRead,
-    ReaderConfig,
-    StarknetStorageValue,
-};
+use crate::block_committer::input::{FactsDbInitialRead, ReaderConfig, StarknetStorageValue};
 use crate::db::db_layout::NodeLayout;
-use crate::db::facts_db::create_facts_tree::{
-    create_original_skeleton_tree,
-    create_original_skeleton_tree_and_get_previous_leaves,
-};
 use crate::db::facts_db::types::FactsSubTree;
 use crate::db::forest_trait::{ForestMetadata, ForestMetadataType, ForestReader, ForestWriter};
+use crate::db::index_db::leaves::TrieType;
+use crate::db::trie_traversal::{create_classes_trie, create_contracts_trie, create_storage_tries};
 use crate::forest::filled_forest::FilledForest;
-use crate::forest::forest_errors::{ForestError, ForestResult};
+use crate::forest::forest_errors::ForestResult;
 use crate::forest::original_skeleton_forest::{ForestSortedIndices, OriginalSkeletonForest};
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
-use crate::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
 use crate::patricia_merkle_tree::types::CompiledClassHash;
 
 /// Facts DB node layout.
@@ -47,12 +38,27 @@ use crate::patricia_merkle_tree::types::CompiledClassHash;
 /// have the db keys of its children.
 pub struct FactsNodeLayout {}
 
-impl<'a, L: Leaf> NodeLayout<'a, L> for FactsNodeLayout {
+impl<'a, L: Leaf> NodeLayout<'a, L> for FactsNodeLayout
+where
+    L: HasStaticPrefix<KeyContext = EmptyKeyContext>,
+{
     type NodeData = HashOutput;
 
     type DeserializationContext = FactNodeDeserializationContext;
 
     type SubTree = FactsSubTree<'a>;
+
+    fn create_subtree(
+        sorted_leaf_indices: SortedLeafIndices<'a>,
+        root_index: NodeIndex,
+        root_hash: HashOutput,
+    ) -> Self::SubTree {
+        FactsSubTree::create(sorted_leaf_indices, root_index, root_hash)
+    }
+
+    fn generate_key_context(_trie_type: TrieType) -> <L as HasStaticPrefix>::KeyContext {
+        EmptyKeyContext
+    }
 }
 
 pub struct FactsDb<S: Storage> {
@@ -68,75 +74,6 @@ impl<S: Storage> FactsDb<S> {
     pub fn new(storage: S) -> Self {
         Self { storage }
     }
-
-    /// Creates the contracts trie original skeleton.
-    /// Also returns the previous contracts state of the modified contracts.
-    async fn create_contracts_trie<'a>(
-        &mut self,
-        contracts_trie_root_hash: HashOutput,
-        contracts_trie_sorted_indices: SortedLeafIndices<'a>,
-    ) -> ForestResult<(OriginalSkeletonTreeImpl<'a>, HashMap<NodeIndex, ContractState>)> {
-        Ok(create_original_skeleton_tree_and_get_previous_leaves(
-            &mut self.storage,
-            contracts_trie_root_hash,
-            contracts_trie_sorted_indices,
-            &HashMap::new(),
-            &OriginalSkeletonTrieConfig::new(true),
-            &EmptyKeyContext,
-        )
-        .await?)
-    }
-
-    async fn create_storage_tries<'a>(
-        &mut self,
-        actual_storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
-        original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
-        config: &impl Config,
-        storage_tries_sorted_indices: &HashMap<ContractAddress, SortedLeafIndices<'a>>,
-    ) -> ForestResult<HashMap<ContractAddress, OriginalSkeletonTreeImpl<'a>>> {
-        let mut storage_tries = HashMap::new();
-        for (address, updates) in actual_storage_updates {
-            let sorted_leaf_indices = storage_tries_sorted_indices
-                .get(address)
-                .ok_or(ForestError::MissingSortedLeafIndices(*address))?;
-            let contract_state = original_contracts_trie_leaves
-                .get(&contract_address_into_node_index(address))
-                .ok_or(ForestError::MissingContractCurrentState(*address))?;
-            let config = OriginalSkeletonTrieConfig::new(config.warn_on_trivial_modifications());
-
-            let original_skeleton = create_original_skeleton_tree(
-                &mut self.storage,
-                contract_state.storage_root_hash,
-                *sorted_leaf_indices,
-                &config,
-                updates,
-                &EmptyKeyContext,
-            )
-            .await?;
-            storage_tries.insert(*address, original_skeleton);
-        }
-        Ok(storage_tries)
-    }
-
-    async fn create_classes_trie<'a>(
-        &mut self,
-        actual_classes_updates: &LeafModifications<CompiledClassHash>,
-        classes_trie_root_hash: HashOutput,
-        config: &impl Config,
-        contracts_trie_sorted_indices: SortedLeafIndices<'a>,
-    ) -> ForestResult<OriginalSkeletonTreeImpl<'a>> {
-        let config = OriginalSkeletonTrieConfig::new(config.warn_on_trivial_modifications());
-
-        Ok(create_original_skeleton_tree(
-            &mut self.storage,
-            classes_trie_root_hash,
-            contracts_trie_sorted_indices,
-            &config,
-            actual_classes_updates,
-            &EmptyKeyContext,
-        )
-        .await?)
-    }
 }
 
 impl FactsDb<MapStorage> {
@@ -158,28 +95,29 @@ impl<S: Storage> ForestReader<FactsDbInitialRead> for FactsDb<S> {
         forest_sorted_indices: &'a ForestSortedIndices<'a>,
         config: ReaderConfig,
     ) -> ForestResult<(OriginalSkeletonForest<'a>, HashMap<NodeIndex, ContractState>)> {
-        let (contracts_trie, original_contracts_trie_leaves) = self
-            .create_contracts_trie(
+        let (contracts_trie, original_contracts_trie_leaves) =
+            create_contracts_trie::<ContractState, FactsNodeLayout>(
+                &mut self.storage,
                 context.0.contracts_trie_root_hash,
                 forest_sorted_indices.contracts_trie_sorted_indices,
             )
             .await?;
-        let storage_tries = self
-            .create_storage_tries(
-                storage_updates,
-                &original_contracts_trie_leaves,
-                &config,
-                &forest_sorted_indices.storage_tries_sorted_indices,
-            )
-            .await?;
-        let classes_trie = self
-            .create_classes_trie(
-                classes_updates,
-                context.0.classes_trie_root_hash,
-                &config,
-                forest_sorted_indices.classes_trie_sorted_indices,
-            )
-            .await?;
+        let storage_tries = create_storage_tries::<StarknetStorageValue, FactsNodeLayout>(
+            &mut self.storage,
+            storage_updates,
+            &original_contracts_trie_leaves,
+            &config,
+            &forest_sorted_indices.storage_tries_sorted_indices,
+        )
+        .await?;
+        let classes_trie = create_classes_trie::<CompiledClassHash, FactsNodeLayout>(
+            &mut self.storage,
+            classes_updates,
+            context.0.classes_trie_root_hash,
+            &config,
+            forest_sorted_indices.classes_trie_sorted_indices,
+        )
+        .await?;
 
         Ok((
             OriginalSkeletonForest { classes_trie, contracts_trie, storage_tries },