starknet_committer: make fetch nodes depend on db layout

ArielElp · ArielElp · commit 780f9597c651 · 2025-12-11T19:08:13.000+02:00
diff --git a/crates/starknet_committer/src/db.rs b/crates/starknet_committer/src/db.rs
@@ -1,3 +1,4 @@
+mod db_layout;
 #[cfg(any(feature = "testing", test))]
 pub mod external_test_utils;
 pub mod facts_db;
diff --git a/crates/starknet_committer/src/db/db_layout.rs b/crates/starknet_committer/src/db/db_layout.rs
@@ -0,0 +1,13 @@
+use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
+use starknet_patricia_storage::errors::DeserializationError;
+use starknet_patricia_storage::storage_trait::DbValue;
+
+pub trait NodeLayout<L: Leaf> {
+    type ChildData: Copy;
+    type DeserializationContext;
+    fn deserialize_node(
+        value: &DbValue,
+        deserialize_context: &Self::DeserializationContext,
+    ) -> Result<FilledNode<L, Self::ChildData>, DeserializationError>;
+}
diff --git a/crates/starknet_committer/src/db/facts_db/create_facts_tree.rs b/crates/starknet_committer/src/db/facts_db/create_facts_tree.rs
@@ -24,7 +24,9 @@ use starknet_patricia_storage::db_object::HasStaticPrefix;
 use starknet_patricia_storage::storage_trait::Storage;
 use tracing::warn;
 
-use crate::db::facts_db::traversal::calculate_subtrees_roots;
+use crate::db::db_layout::NodeLayout;
+use crate::db::facts_db::db::FactsNodeLayout;
+use crate::db::facts_db::traversal::get_roots_from_storage;
 use crate::db::facts_db::types::FactsSubTree;
 
 #[cfg(test)]
@@ -42,65 +44,85 @@ macro_rules! log_trivial_modification {
 /// Given a list of subtrees, traverses towards their leaves and fetches all non-empty,
 /// unmodified nodes. If `compare_modified_leaves` is set, function logs out a warning when
 /// encountering a trivial modification. Fills the previous leaf values if it is not none.
-async fn fetch_nodes<'a, L: Leaf>(
+async fn fetch_nodes<'a, L, Layout, SubTree>(
     skeleton_tree: &mut OriginalSkeletonTreeImpl<'a>,
-    subtrees: Vec<FactsSubTree<'a>>,
+    subtrees: Vec<SubTree>,
     storage: &mut impl Storage,
     leaf_modifications: &LeafModifications<L>,
     config: &impl OriginalSkeletonTreeConfig<L>,
     mut previous_leaves: Option<&mut HashMap<NodeIndex, L>>,
     key_context: &<L as HasStaticPrefix>::KeyContext,
-) -> OriginalSkeletonTreeResult<()> {
+) -> OriginalSkeletonTreeResult<()>
+where
+    L: Leaf,
+    Layout: NodeLayout<L>,
+    SubTree: SubTreeTrait<'a, NodeData = Layout::ChildData, NodeContext = Layout::DeserializationContext>,
+{
     let mut current_subtrees = subtrees;
     let mut next_subtrees = Vec::new();
     while !current_subtrees.is_empty() {
         let should_fetch_modified_leaves =
             config.compare_modified_leaves() || previous_leaves.is_some();
         let filled_roots =
-            calculate_subtrees_roots::<L>(&current_subtrees, storage, key_context).await?;
+            get_roots_from_storage::<L, Layout>(&current_subtrees, storage, key_context).await?;
         for (filled_root, subtree) in filled_roots.into_iter().zip(current_subtrees.iter()) {
             match filled_root.data {
                 // Binary node.
-                NodeData::<L, HashOutput>::Binary(BinaryData { left_data, right_data }) => {
+                NodeData::Binary(BinaryData { left_data, right_data }) => {
                     if subtree.is_unmodified() {
                         skeleton_tree.nodes.insert(
-                            subtree.root_index,
+                            subtree.get_root_index(),
                             OriginalSkeletonNode::UnmodifiedSubTree(filled_root.hash),
                         );
                         continue;
                     }
-                    skeleton_tree.nodes.insert(subtree.root_index, OriginalSkeletonNode::Binary);
+                    skeleton_tree
+                        .nodes
+                        .insert(subtree.get_root_index(), OriginalSkeletonNode::Binary);
                     let (left_subtree, right_subtree) =
                         subtree.get_children_subtrees(left_data, right_data);
 
                     handle_subtree(
                         skeleton_tree,
                         &mut next_subtrees,
                         left_subtree,
+                        left_data,
                         should_fetch_modified_leaves,
                     );
                     handle_subtree(
                         skeleton_tree,
                         &mut next_subtrees,
                         right_subtree,
+                        right_data,
                         should_fetch_modified_leaves,
                     )
                 }
                 // Edge node.
-                NodeData::<L, HashOutput>::Edge(EdgeData { bottom_data, path_to_bottom }) => {
-                    skeleton_tree
-                        .nodes
-                        .insert(subtree.root_index, OriginalSkeletonNode::Edge(path_to_bottom));
-                    if subtree.is_unmodified() {
-                        skeleton_tree.nodes.insert(
-                            path_to_bottom.bottom_index(subtree.root_index),
-                            OriginalSkeletonNode::UnmodifiedSubTree(bottom_data),
-                        );
-                        continue;
-                    }
+                NodeData::Edge(EdgeData { bottom_data, path_to_bottom }) => {
+                    skeleton_tree.nodes.insert(
+                        subtree.get_root_index(),
+                        OriginalSkeletonNode::Edge(path_to_bottom),
+                    );
+
                     // Parse bottom.
                     let (bottom_subtree, previously_empty_leaves_indices) =
                         subtree.get_bottom_subtree(&path_to_bottom, bottom_data);
+
+                    if subtree.is_unmodified() {
+                        if !SubTree::should_traverse_unmodified_children() {
+                            skeleton_tree.nodes.insert(
+                                path_to_bottom.bottom_index(subtree.get_root_index()),
+                                OriginalSkeletonNode::UnmodifiedSubTree(
+                                    SubTree::unmodified_child_hash(bottom_data).unwrap(),
+                                ),
+                            );
+                        } else {
+                            // With index layout we need to traverse an unmodified bottom node.
+                            next_subtrees.push(bottom_subtree)
+                        }
+                        continue;
+                    }
+
                     if let Some(ref mut leaves) = previous_leaves {
                         leaves.extend(
                             previously_empty_leaves_indices
@@ -119,23 +141,28 @@ async fn fetch_nodes<'a, L: Leaf>(
                         skeleton_tree,
                         &mut next_subtrees,
                         bottom_subtree,
+                        bottom_data,
                         should_fetch_modified_leaves,
                     );
                 }
                 // Leaf node.
                 NodeData::Leaf(previous_leaf) => {
                     if subtree.is_unmodified() {
-                        warn!("Unexpectedly deserialized leaf sibling.")
+                        skeleton_tree.nodes.insert(
+                            subtree.get_root_index(),
+                            OriginalSkeletonNode::UnmodifiedSubTree(filled_root.hash),
+                        );
                     } else {
+                        let root_index = subtree.get_root_index();
                         // Modified leaf.
                         if config.compare_modified_leaves()
-                            && L::compare(leaf_modifications, &subtree.root_index, &previous_leaf)?
+                            && L::compare(leaf_modifications, &root_index, &previous_leaf)?
                         {
-                            log_trivial_modification!(subtree.root_index, previous_leaf);
+                            log_trivial_modification!(root_index, previous_leaf);
                         }
                         // If previous values of modified leaves are requested, add this leaf.
                         if let Some(ref mut leaves) = previous_leaves {
-                            leaves.insert(subtree.root_index, previous_leaf);
+                            leaves.insert(root_index, previous_leaf);
                         }
                     }
                 }
@@ -166,9 +193,9 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf>(
         )?;
         return Ok(OriginalSkeletonTreeImpl::create_empty(sorted_leaf_indices));
     }
-    let main_subtree = FactsSubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash);
+    let main_subtree = FactsSubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
     let mut skeleton_tree = OriginalSkeletonTreeImpl { nodes: HashMap::new(), sorted_leaf_indices };
-    fetch_nodes::<L>(
+    fetch_nodes::<L, FactsNodeLayout, FactsSubTree<'a>>(
         &mut skeleton_tree,
         vec![main_subtree],
         storage,
@@ -202,7 +229,7 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
     let main_subtree = FactsSubTree { sorted_leaf_indices, root_index: NodeIndex::ROOT, root_hash };
     let mut skeleton_tree = OriginalSkeletonTreeImpl { nodes: HashMap::new(), sorted_leaf_indices };
     let mut leaves = HashMap::new();
-    fetch_nodes::<L>(
+    fetch_nodes::<L, FactsNodeLayout, FactsSubTree<'a>>(
         &mut skeleton_tree,
         vec![main_subtree],
         storage,
@@ -237,19 +264,40 @@ pub async fn get_leaves<'a, L: Leaf>(
 
 /// Handles a subtree referred by an edge or a binary node. Decides whether we deserialize the
 /// referred subtree or not.
-fn handle_subtree<'a>(
+fn handle_subtree<'a, SubTree: SubTreeTrait<'a>>(
     skeleton_tree: &mut OriginalSkeletonTreeImpl<'a>,
-    next_subtrees: &mut Vec<FactsSubTree<'a>>,
-    subtree: FactsSubTree<'a>,
+    next_subtrees: &mut Vec<SubTree>,
+    subtree: SubTree,
+    subtree_data: SubTree::NodeData,
     should_fetch_modified_leaves: bool,
 ) {
-    if !subtree.is_leaf() || (should_fetch_modified_leaves && !subtree.is_unmodified()) {
+    let is_leaf = subtree.is_leaf();
+    let is_unmodified = subtree.is_unmodified();
+
+    // 1. Internal node → always traverse.
+    if !is_leaf {
+        next_subtrees.push(subtree);
+        return;
+    }
+
+    // 2. Modified leaf.
+    if !is_unmodified {
+        if should_fetch_modified_leaves {
+            next_subtrees.push(subtree);
+        }
+        return;
+    }
+
+    // 3. Unmodified leaf sibling.
+    if !SubTree::should_traverse_unmodified_children() {
+        skeleton_tree.nodes.insert(
+            subtree.get_root_index(),
+            OriginalSkeletonNode::UnmodifiedSubTree(
+                SubTree::unmodified_child_hash(subtree_data).unwrap(),
+            ),
+        );
+    } else {
         next_subtrees.push(subtree);
-    } else if subtree.is_unmodified() {
-        // Leaf sibling.
-        skeleton_tree
-            .nodes
-            .insert(subtree.root_index, OriginalSkeletonNode::UnmodifiedSubTree(subtree.root_hash));
     }
 }
 
diff --git a/crates/starknet_committer/src/db/facts_db/db.rs b/crates/starknet_committer/src/db/facts_db/db.rs
@@ -2,23 +2,27 @@ use std::collections::HashMap;
 
 use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
-use starknet_patricia::patricia_merkle_tree::node_data::leaf::LeafModifications;
+use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::FactNodeDeserializationContext;
+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::types::{NodeIndex, SortedLeafIndices};
-use starknet_patricia_storage::db_object::EmptyKeyContext;
+use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext};
+use starknet_patricia_storage::errors::DeserializationError;
 use starknet_patricia_storage::map_storage::MapStorage;
-use starknet_patricia_storage::storage_trait::Storage;
+use starknet_patricia_storage::storage_trait::{DbValue, Storage};
 
 use crate::block_committer::input::{
     contract_address_into_node_index,
     Config,
     ConfigImpl,
     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::FactDbFilledNode;
 use crate::db::forest_trait::{ForestReader, ForestWriter};
 use crate::forest::filled_forest::FilledForest;
 use crate::forest::forest_errors::{ForestError, ForestResult};
@@ -31,6 +35,18 @@ use crate::patricia_merkle_tree::tree::{
 };
 use crate::patricia_merkle_tree::types::CompiledClassHash;
 
+pub struct FactsNodeLayout {}
+
+impl<L: Leaf> NodeLayout<L> for FactsNodeLayout {
+    type ChildData = HashOutput;
+    type DeserializationContext = FactNodeDeserializationContext;
+    fn deserialize_node(
+        value: &DbValue,
+        deserialize_context: &Self::DeserializationContext,
+    ) -> Result<FactDbFilledNode<L>, DeserializationError> {
+        DBObject::deserialize(value, deserialize_context)
+    }
+}
 pub struct FactsDb<S: Storage> {
     // TODO(Yoav): Define StorageStats trait and impl it here. Then, make the storage field
     // private.
diff --git a/crates/starknet_committer/src/db/facts_db/traversal.rs b/crates/starknet_committer/src/db/facts_db/traversal.rs
@@ -2,7 +2,6 @@ use std::collections::HashMap;
 
 use starknet_api::hash::HashOutput;
 use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
-use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::FactNodeDeserializationContext;
 use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     NodeData,
     Preimage,
@@ -11,43 +10,41 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
 use starknet_patricia::patricia_merkle_tree::traversal::{SubTreeTrait, TraversalResult};
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
-use starknet_patricia_storage::db_object::{DBObject, HasStaticPrefix};
+use starknet_patricia_storage::db_object::HasStaticPrefix;
 use starknet_patricia_storage::errors::StorageError;
 use starknet_patricia_storage::storage_trait::{create_db_key, DbKey, Storage};
 
+use crate::db::db_layout::NodeLayout;
+use crate::db::facts_db::db::FactsNodeLayout;
 use crate::db::facts_db::types::FactsSubTree;
 
 #[cfg(test)]
 #[path = "traversal_test.rs"]
 pub mod traversal_test;
 
 // TODO(Aviv, 17/07/2024): Split between storage prefix implementation and function logic.
-pub async fn calculate_subtrees_roots<'a, L: Leaf>(
-    subtrees: &[FactsSubTree<'a>],
+pub async fn get_roots_from_storage<'a, L: Leaf, Layout: NodeLayout<L>>(
+    subtrees: &[impl SubTreeTrait<
+        'a,
+        NodeData = Layout::ChildData,
+        NodeContext = Layout::DeserializationContext,
+    >],
     storage: &mut impl Storage,
     key_context: &<L as HasStaticPrefix>::KeyContext,
-) -> TraversalResult<Vec<FilledNode<L>>> {
+) -> TraversalResult<Vec<FilledNode<L, Layout::ChildData>>> {
     let mut subtrees_roots = vec![];
     let db_keys: Vec<DbKey> = subtrees
         .iter()
         .map(|subtree| {
-            create_db_key(
-                subtree.get_root_prefix::<L>(key_context),
-                &subtree.root_hash.0.to_bytes_be(),
-            )
+            create_db_key(subtree.get_root_prefix::<L>(key_context), &subtree.get_root_suffix())
         })
         .collect();
 
     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))? };
-        subtrees_roots.push(FilledNode::deserialize(
-            val,
-            &FactNodeDeserializationContext {
-                is_leaf: subtree.is_leaf(),
-                node_hash: subtree.root_hash,
-            },
-        )?)
+        let filled_node = Layout::deserialize_node(val, &subtree.get_root_context())?;
+        subtrees_roots.push(filled_node);
     }
     Ok(subtrees_roots)
 }
@@ -100,7 +97,8 @@ pub(crate) async fn fetch_patricia_paths_inner<'a, L: Leaf>(
     let mut next_subtrees = Vec::new();
     while !current_subtrees.is_empty() {
         let filled_roots =
-            calculate_subtrees_roots::<L>(&current_subtrees, storage, key_context).await?;
+            get_roots_from_storage::<L, FactsNodeLayout>(&current_subtrees, storage, key_context)
+                .await?;
         for (filled_root, subtree) in filled_roots.into_iter().zip(current_subtrees.iter()) {
             // Always insert root.
             // No need to insert an unmodified node (which is not the root), because its parent is

Original file line number	Diff line number	Diff line change
`@@ -1,3 +1,4 @@`
	`1`	`+mod db_layout;`
`1`	`2`	`#[cfg(any(feature = "testing", test))]`
`2`	`3`	`pub mod external_test_utils;`
`3`	`4`	`pub mod facts_db;`