starknet_committer: abstract original tree creation

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;
 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,12 +34,13 @@ 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,
         &config,
         &leaf_modifications,
+        None,
         &EmptyKeyContext,
     )
     .await

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

@@ -5,61 +5,19 @@ use starknet_patricia::patricia_merkle_tree::node_data::leaf::{
     LeafModifications,
     LeafWithEmptyKeyContext,
 };
-use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonTreeConfig;
-use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::tree::{
-    OriginalSkeletonTreeImpl,
-    OriginalSkeletonTreeResult,
-};
-use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait;
+use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonTreeResult;
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
 use starknet_patricia_storage::db_object::EmptyKeyContext;
 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;
+use crate::db::trie_traversal::create_original_skeleton_tree;
 use crate::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
 
 #[cfg(test)]
 #[path = "create_facts_tree_test.rs"]
 pub mod create_facts_tree_test;
 
-pub async fn create_original_skeleton_tree_and_get_previous_leaves<
-    'a,
-    L: LeafWithEmptyKeyContext,
->(
-    storage: &mut impl Storage,
-    root_hash: HashOutput,
-    sorted_leaf_indices: SortedLeafIndices<'a>,
-    leaf_modifications: &LeafModifications<L>,
-    config: &impl OriginalSkeletonTreeConfig,
-) -> OriginalSkeletonTreeResult<(OriginalSkeletonTreeImpl<'a>, HashMap<NodeIndex, L>)> {
-    if sorted_leaf_indices.is_empty() {
-        let unmodified = OriginalSkeletonTreeImpl::create_unmodified(root_hash);
-        return Ok((unmodified, HashMap::new()));
-    }
-    if root_hash == HashOutput::ROOT_OF_EMPTY_TREE {
-        return Ok((
-            OriginalSkeletonTreeImpl::create_empty(sorted_leaf_indices),
-            sorted_leaf_indices.get_indices().iter().map(|idx| (*idx, L::default())).collect(),
-        ));
-    }
-    let main_subtree = FactsSubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash);
-    let mut skeleton_tree = OriginalSkeletonTreeImpl { nodes: HashMap::new(), sorted_leaf_indices };
-    let mut leaves = HashMap::new();
-    fetch_nodes::<L, FactsNodeLayout>(
-        &mut skeleton_tree,
-        vec![main_subtree],
-        storage,
-        leaf_modifications,
-        config,
-        Some(&mut leaves),
-        &EmptyKeyContext,
-    )
-    .await?;
-    Ok((skeleton_tree, leaves))
-}
-
 /// 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.
@@ -70,12 +28,15 @@ pub async fn get_leaves<'a, L: LeafWithEmptyKeyContext>(
 ) -> OriginalSkeletonTreeResult<HashMap<NodeIndex, L>> {
     let config = OriginalSkeletonTrieConfig::default();
     let leaf_modifications = LeafModifications::new();
-    let (_, previous_leaves) = create_original_skeleton_tree_and_get_previous_leaves(
+    let mut previous_leaves = HashMap::new();
+    create_original_skeleton_tree::<L, FactsNodeLayout>(
         storage,
         root_hash,
         sorted_leaf_indices,
-        &leaf_modifications,
         &config,
+        &leaf_modifications,
+        Some(&mut previous_leaves),
+        &EmptyKeyContext,
     )
     .await?;
     Ok(previous_leaves)

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,12 +212,13 @@ 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,
         &config,
         &leaf_modifications,
+        None,
         &EmptyKeyContext,
     )
     .await

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

@@ -80,21 +80,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

@@ -10,11 +10,7 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     EdgeData,
     NodeData,
 };
-use starknet_patricia::patricia_merkle_tree::node_data::leaf::{
-    Leaf,
-    LeafModifications,
-    LeafWithEmptyKeyContext,
-};
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
 use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonTreeConfig;
 use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::node::OriginalSkeletonNode;
 use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::tree::{
@@ -27,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;
@@ -38,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;
@@ -61,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,
@@ -73,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();
@@ -251,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()
@@ -266,7 +256,8 @@ 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::NodeDbObject::deserialize(val, &subtree.get_root_context())?.into();
         subtrees_roots.push(filled_node);
     }
     Ok(subtrees_roots)
@@ -292,12 +283,27 @@ pub(crate) fn log_warning_for_empty_leaves<L: Leaf, T: Borrow<NodeIndex> + Debug
     Ok(())
 }
 
-pub async fn create_original_skeleton_tree<'a, L: LeafWithEmptyKeyContext>(
+/// Creates an original skeleton tree by fetching Patricia nodes from storage.
+///
+/// Traverses the trie from the root towards the modified leaves, collecting all nodes needed
+/// to construct the skeleton. If `previous_leaves` is provided, it will be populated with the
+/// previous values of modified leaves.
+///
+/// # Layout-Dependent Arguments
+///
+/// The following arguments depend on the `Layout` type parameter:
+/// - `storage`: storage to fetch nodes from, expected to match the serialization of
+///   `Layout::NodeDbObject`.
+/// - `key_context`: Generated via `Layout::generate_key_context`, determines the DB key prefix used
+///   when fetching nodes from storage.
+/// - `leaf_modifications` and `previous_leaves`: Their leaf type `L` is constrained by `Layout`.
+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>,
     config: &impl OriginalSkeletonTreeConfig,
     leaf_modifications: &LeafModifications<L>,
+    previous_leaves: Option<&mut HashMap<NodeIndex, L>>,
     key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> OriginalSkeletonTreeResult<OriginalSkeletonTreeImpl<'a>> {
     if sorted_leaf_indices.is_empty() {
@@ -309,24 +315,38 @@ pub async fn create_original_skeleton_tree<'a, L: LeafWithEmptyKeyContext>(
             leaf_modifications,
             config,
         )?;
+        if let Some(previous_leaves) = previous_leaves {
+            previous_leaves.extend(
+                sorted_leaf_indices
+                    .get_indices()
+                    .iter()
+                    .map(|idx| (*idx, L::default()))
+                    .collect::<HashMap<NodeIndex, L>>(),
+            );
+        }
         return Ok(OriginalSkeletonTreeImpl::create_empty(sorted_leaf_indices));
     }
-    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,
         leaf_modifications,
         config,
-        None,
+        previous_leaves,
         key_context,
     )
     .await?;
     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>,
@@ -345,13 +365,16 @@ 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,
+            // TODO(Ariel): Change `LeafModifications` in `actual_storage_updates` to be an
+            // iterator over borrowed data so that the conversion below is costless.
+            &updates.iter().map(|(idx, value)| (*idx, L::from(*value))).collect(),
+            None,
+            &Layout::generate_key_context(TrieType::StorageTrie(*address)),
         )
         .await?;
         storage_tries.insert(*address, original_skeleton);
@@ -361,22 +384,36 @@ 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(
+    let config = OriginalSkeletonTrieConfig::new_for_contracts_trie();
+
+    let mut leaves = HashMap::new();
+    let skeleton_tree = create_original_skeleton_tree::<L, Layout>(
         storage,
         contracts_trie_root_hash,
         contracts_trie_sorted_indices,
+        &config,
         &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,
@@ -387,13 +424,16 @@ 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,
+        // TODO(Ariel): Change `actual_classes_updates` to be an iterator over borrowed data so
+        // that the conversion below is costless.
+        &actual_classes_updates.iter().map(|(idx, value)| (*idx, L::from(*value))).collect(),
+        None,
+        &Layout::generate_key_context(TrieType::ClassesTrie),
     )
     .await?)
 }