starknet_committer: abstract original tree creation

Author: ArielElp

crates/starknet_committer/src/db/db_layout.rs

@@ -7,10 +7,7 @@ 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>
-where
-    FilledNode<L, Self::NodeData>: DBObject<DeserializeContext = Self::DeserializationContext>,
-{
+pub trait NodeLayout<'a, L: Leaf> {
     /// Additional data that a node stores about its children.
     type NodeData: Clone + From<HashOutput>;
 

crates/starknet_committer/src/db/external_test_utils.rs

@@ -18,6 +18,7 @@ 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::db::FactsNodeLayout;
 use crate::db::trie_traversal::create_original_skeleton_tree;
 
 // TODO(Ariel, 14/12/2025): make this generic over the layout.
@@ -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_test.rs

@@ -22,6 +22,7 @@ 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::db::FactsNodeLayout;
 use crate::db::trie_traversal::create_original_skeleton_tree;
 use crate::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
 
@@ -211,7 +212,7 @@ async fn test_create_tree(
     let config = OriginalSkeletonTrieConfig::new_for_tests(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

@@ -84,21 +84,22 @@ 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) = create_contracts_trie(
-            &mut self.storage,
-            context.0.contracts_trie_root_hash,
-            forest_sorted_indices.contracts_trie_sorted_indices,
-        )
-        .await?;
-        let storage_tries = create_storage_tries(
+        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 = 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(
+        let classes_trie = create_classes_trie::<CompiledClassHash, FactsNodeLayout>(
             &mut self.storage,
             classes_updates,
             context.0.classes_trie_root_hash,

crates/starknet_committer/src/db/trie_traversal.rs

@@ -23,7 +23,7 @@ use starknet_patricia::patricia_merkle_tree::traversal::{
     UnmodifiedChildTraversal,
 };
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
-use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
+use starknet_patricia_storage::db_object::{DBObject, HasStaticPrefix};
 use starknet_patricia_storage::errors::StorageError;
 use starknet_patricia_storage::storage_trait::{create_db_key, DbKey, Storage};
 use tracing::warn;
@@ -34,9 +34,7 @@ use crate::block_committer::input::{
     StarknetStorageValue,
 };
 use crate::db::db_layout::NodeLayout;
-use crate::db::facts_db::create_facts_tree::create_original_skeleton_tree_and_get_previous_leaves;
-use crate::db::facts_db::db::FactsNodeLayout;
-use crate::db::facts_db::types::FactsSubTree;
+use crate::db::index_db::leaves::TrieType;
 use crate::forest::forest_errors::{ForestError, ForestResult};
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use crate::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
@@ -57,7 +55,7 @@ macro_rules! log_trivial_modification {
 ///
 /// The function is generic over the DB layout (`Layout`), which controls the concrete node data
 /// (`Layout::NodeData`) and traversal strategy (via `Layout::SubTree`).
-pub async fn fetch_nodes<'a, L, Layout>(
+pub(crate) async fn fetch_nodes<'a, L, Layout>(
     skeleton_tree: &mut OriginalSkeletonTreeImpl<'a>,
     subtrees: Vec<Layout::SubTree>,
     storage: &mut impl Storage,
@@ -69,7 +67,6 @@ pub async fn fetch_nodes<'a, L, Layout>(
 where
     L: Leaf,
     Layout: NodeLayout<'a, L>,
-    FilledNode<L, Layout::NodeData>: DBObject<DeserializeContext = Layout::DeserializationContext>,
 {
     let mut current_subtrees = subtrees;
     let mut next_subtrees = Vec::new();
@@ -247,10 +244,7 @@ pub async fn get_roots_from_storage<'a, L: Leaf, Layout: NodeLayout<'a, L>>(
     subtrees: &[Layout::SubTree],
     storage: &mut impl Storage,
     key_context: &<L as HasStaticPrefix>::KeyContext,
-) -> TraversalResult<Vec<FilledNode<L, Layout::NodeData>>>
-where
-    FilledNode<L, Layout::NodeData>: DBObject<DeserializeContext = Layout::DeserializationContext>,
-{
+) -> TraversalResult<Vec<FilledNode<L, Layout::NodeData>>> {
     let mut subtrees_roots = vec![];
     let db_keys: Vec<DbKey> = subtrees
         .iter()
@@ -262,7 +256,10 @@ where
     let db_vals = storage.mget(&db_keys.iter().collect::<Vec<&DbKey>>()).await?;
     for ((subtree, optional_val), db_key) in subtrees.iter().zip(db_vals.iter()).zip(db_keys) {
         let Some(val) = optional_val else { Err(StorageError::MissingKey(db_key))? };
-        let filled_node = FilledNode::deserialize(val, &subtree.get_root_context())?;
+        let filled_node = Layout::get_filled_node(Layout::NodeDbObject::deserialize(
+            val,
+            &subtree.get_root_context(),
+        )?);
         subtrees_roots.push(filled_node);
     }
     Ok(subtrees_roots)
@@ -288,10 +285,7 @@ pub(crate) fn log_warning_for_empty_leaves<L: Leaf, T: Borrow<NodeIndex> + Debug
     Ok(())
 }
 
-pub async fn create_original_skeleton_tree<
-    'a,
-    L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext>,
->(
+pub async fn create_original_skeleton_tree<'a, L: Leaf, Layout: NodeLayout<'a, L>>(
     storage: &mut impl Storage,
     root_hash: HashOutput,
     sorted_leaf_indices: SortedLeafIndices<'a>,
@@ -308,11 +302,12 @@ pub async fn create_original_skeleton_tree<
             leaf_modifications,
             config,
         )?;
-        return Ok(OriginalSkeletonTreeImpl::create_empty(sorted_leaf_indices));
+        return Ok(handle_empty_subtree::<L>(sorted_leaf_indices).0);
     }
-    let main_subtree = FactsSubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash);
+    let main_subtree =
+        Layout::SubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash.into());
     let mut skeleton_tree = OriginalSkeletonTreeImpl { nodes: HashMap::new(), sorted_leaf_indices };
-    fetch_nodes::<L, FactsNodeLayout>(
+    fetch_nodes::<L, Layout>(
         &mut skeleton_tree,
         vec![main_subtree],
         storage,
@@ -325,7 +320,11 @@ pub async fn create_original_skeleton_tree<
     Ok(skeleton_tree)
 }
 
-pub async fn create_storage_tries<'a>(
+pub async fn create_storage_tries<
+    'a,
+    L: Leaf + From<StarknetStorageValue>,
+    Layout: NodeLayout<'a, L>,
+>(
     storage: &mut impl Storage,
     actual_storage_updates: &HashMap<ContractAddress, LeafModifications<StarknetStorageValue>>,
     original_contracts_trie_leaves: &HashMap<NodeIndex, ContractState>,
@@ -344,13 +343,13 @@ pub async fn create_storage_tries<'a>(
             config.warn_on_trivial_modifications(),
         );
 
-        let original_skeleton = create_original_skeleton_tree(
+        let original_skeleton = create_original_skeleton_tree::<L, Layout>(
             storage,
             contract_state.storage_root_hash,
             *sorted_leaf_indices,
             &config,
-            updates,
-            &EmptyKeyContext,
+            &updates.iter().map(|(idx, value)| (*idx, L::from(*value))).collect(),
+            &Layout::generate_key_context(TrieType::StorageTrie(*address)),
         )
         .await?;
         storage_tries.insert(*address, original_skeleton);
@@ -360,22 +359,50 @@ pub async fn create_storage_tries<'a>(
 
 /// Creates the contracts trie original skeleton.
 /// Also returns the previous contracts state of the modified contracts.
-pub async fn create_contracts_trie<'a>(
+pub async fn create_contracts_trie<'a, L: Leaf + Into<ContractState>, Layout: NodeLayout<'a, L>>(
     storage: &mut impl Storage,
     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(
-        storage,
-        contracts_trie_root_hash,
+    if contracts_trie_sorted_indices.is_empty() {
+        let unmodified = OriginalSkeletonTreeImpl::create_unmodified(contracts_trie_root_hash);
+        return Ok((unmodified, HashMap::new()));
+    }
+    if contracts_trie_root_hash == HashOutput::ROOT_OF_EMPTY_TREE {
+        return Ok(handle_empty_subtree(contracts_trie_sorted_indices));
+    }
+    let main_subtree = Layout::SubTree::create(
         contracts_trie_sorted_indices,
+        NodeIndex::ROOT,
+        contracts_trie_root_hash.into(),
+    );
+    let mut skeleton_tree = OriginalSkeletonTreeImpl {
+        nodes: HashMap::new(),
+        sorted_leaf_indices: contracts_trie_sorted_indices,
+    };
+    let mut leaves = HashMap::new();
+    fetch_nodes::<L, Layout>(
+        &mut skeleton_tree,
+        vec![main_subtree],
+        storage,
         &HashMap::new(),
         &OriginalSkeletonTrieConfig::new_for_contracts_trie(),
+        Some(&mut leaves),
+        &Layout::generate_key_context(TrieType::ContractsTrie),
     )
-    .await?)
+    .await?;
+
+    let leaves: HashMap<NodeIndex, ContractState> =
+        leaves.into_iter().map(|(idx, leaf)| (idx, leaf.into())).collect();
+
+    Ok((skeleton_tree, leaves))
 }
 
-pub async fn create_classes_trie<'a>(
+pub async fn create_classes_trie<
+    'a,
+    L: Leaf + From<CompiledClassHash>,
+    Layout: NodeLayout<'a, L>,
+>(
     storage: &mut impl Storage,
     actual_classes_updates: &LeafModifications<CompiledClassHash>,
     classes_trie_root_hash: HashOutput,
@@ -386,13 +413,22 @@ pub async fn create_classes_trie<'a>(
         config.warn_on_trivial_modifications(),
     );
 
-    Ok(create_original_skeleton_tree(
+    Ok(create_original_skeleton_tree::<L, Layout>(
         storage,
         classes_trie_root_hash,
         contracts_trie_sorted_indices,
         &config,
-        actual_classes_updates,
-        &EmptyKeyContext,
+        &actual_classes_updates.iter().map(|(idx, value)| (*idx, L::from(*value))).collect(),
+        &Layout::generate_key_context(TrieType::ClassesTrie),
     )
     .await?)
 }
+
+pub(crate) fn handle_empty_subtree<'a, L: Leaf>(
+    sorted_leaf_indices: SortedLeafIndices<'a>,
+) -> (OriginalSkeletonTreeImpl<'a>, HashMap<NodeIndex, L>) {
+    (
+        OriginalSkeletonTreeImpl::create_empty(sorted_leaf_indices),
+        sorted_leaf_indices.get_indices().iter().map(|idx| (*idx, L::default())).collect(),
+    )
+}