starknet_committer: make fetch nodes depend on db layout

ArielElp · ArielElp · commit d711681f0783 · 2025-12-16T14:21:55.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,23 @@
+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;
+
+/// A trait that specifies the trie db layout.
+///
+/// The layout determines:
+/// [NodeData]: additional data that a node stores about its children.
+/// [DeserializationContext]: the context needed to deserialize the node from a raw [DbValue].
+/// [SubTree]: the type of the subtree that is used to traverse the trie.
+pub trait NodeLayout<'a, L: Leaf>
+where
+    FilledNode<L, Self::NodeData>: DBObject<DeserializeContext = Self::DeserializationContext>,
+{
+    type NodeData: Copy;
+    type DeserializationContext;
+    type SubTree: SubTreeTrait<
+            'a,
+            NodeData = Self::NodeData,
+            NodeDeserializeContext = Self::DeserializationContext,
+        >;
+}
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
@@ -3,6 +3,7 @@ use std::collections::HashMap;
 use std::fmt::Debug;
 
 use starknet_api::hash::HashOutput;
+use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
 use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     BinaryData,
     EdgeData,
@@ -18,13 +19,15 @@ 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::traversal::{SubTreeTrait, UnmodifiedChildTraversal};
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
-use starknet_patricia_storage::db_object::HasStaticPrefix;
+use starknet_patricia_storage::db_object::{DBObject, 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)]
@@ -39,68 +42,90 @@ macro_rules! log_trivial_modification {
 }
 
 /// Fetches the Patricia witnesses, required to build the original skeleton tree from storage.
+///
 /// 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>(
+///
+/// The function is generic over the DB layout (`Layout`), which controls the concrete node data
+/// (`Layout::NodeData`) and traversal strategy (via `Layout::SubTree`).
+async fn fetch_nodes<'a, L, Layout>(
     skeleton_tree: &mut OriginalSkeletonTreeImpl<'a>,
-    subtrees: Vec<FactsSubTree<'a>>,
+    subtrees: Vec<Layout::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<'a, L>,
+    FilledNode<L, Layout::NodeData>: DBObject<DeserializeContext = 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()) {
+            let root_index = subtree.get_root_index();
             match filled_root.data {
                 // Binary node.
                 NodeData::Binary(BinaryData { left_data, right_data }) => {
                     if subtree.is_unmodified() {
                         skeleton_tree.nodes.insert(
-                            subtree.root_index,
+                            root_index,
                             OriginalSkeletonNode::UnmodifiedSubTree(filled_root.hash),
                         );
                         continue;
                     }
-                    skeleton_tree.nodes.insert(subtree.root_index, OriginalSkeletonNode::Binary);
+                    skeleton_tree.nodes.insert(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::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;
-                    }
+                        .insert(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() {
+                        match Layout::SubTree::should_traverse_unmodified_child(bottom_data) {
+                            UnmodifiedChildTraversal::Traverse => {
+                                next_subtrees.push(bottom_subtree);
+                            }
+                            UnmodifiedChildTraversal::Skip(unmodified_child_hash) => {
+                                skeleton_tree.nodes.insert(
+                                    path_to_bottom.bottom_index(root_index),
+                                    OriginalSkeletonNode::UnmodifiedSubTree(unmodified_child_hash),
+                                );
+                            }
+                        }
+                        continue;
+                    }
+
                     if let Some(ref mut leaves) = previous_leaves {
                         leaves.extend(
                             previously_empty_leaves_indices
@@ -119,23 +144,27 @@ 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(
+                            root_index,
+                            OriginalSkeletonNode::UnmodifiedSubTree(filled_root.hash),
+                        );
                     } else {
                         // 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);
                         }
                     }
                 }
@@ -168,7 +197,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf>(
     }
     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>(
+    fetch_nodes::<L, FactsNodeLayout>(
         &mut skeleton_tree,
         vec![main_subtree],
         storage,
@@ -202,7 +231,7 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
     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>(
+    fetch_nodes::<L, FactsNodeLayout>(
         &mut skeleton_tree,
         vec![main_subtree],
         storage,
@@ -240,19 +269,41 @@ 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_modified = !subtree.is_unmodified();
+
+    // Internal node → always traverse.
+    if !is_leaf {
         next_subtrees.push(subtree);
-    } else if subtree.is_unmodified() {
-        // Leaf sibling.
-        skeleton_tree
-            .nodes
-            .insert(subtree.root_index, OriginalSkeletonNode::UnmodifiedSubTree(subtree.root_hash));
+        return;
+    }
+
+    // Modified leaf.
+    if is_modified {
+        if should_fetch_modified_leaves {
+            next_subtrees.push(subtree);
+        }
+        return;
+    }
+
+    // Unmodified leaf sibling.
+    match SubTree::should_traverse_unmodified_child(subtree_data) {
+        UnmodifiedChildTraversal::Traverse => {
+            next_subtrees.push(subtree);
+        }
+        UnmodifiedChildTraversal::Skip(unmodified_child_hash) => {
+            skeleton_tree.nodes.insert(
+                subtree.get_root_index(),
+                OriginalSkeletonNode::UnmodifiedSubTree(unmodified_child_hash),
+            );
+        }
     }
 }
 
diff --git a/crates/starknet_committer/src/db/facts_db/db.rs b/crates/starknet_committer/src/db/facts_db/db.rs
@@ -3,8 +3,9 @@ use std::collections::HashMap;
 use async_trait::async_trait;
 use starknet_api::core::ContractAddress;
 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::LeafModifications;
+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;
@@ -24,10 +25,12 @@ use crate::block_committer::input::{
     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::forest::filled_forest::FilledForest;
 use crate::forest::forest_errors::{ForestError, ForestResult};
@@ -40,6 +43,16 @@ use crate::patricia_merkle_tree::tree::{
 };
 use crate::patricia_merkle_tree::types::CompiledClassHash;
 
+pub struct FactsNodeLayout {}
+
+impl<'a, L: Leaf> NodeLayout<'a, L> for FactsNodeLayout {
+    type NodeData = HashOutput;
+
+    type DeserializationContext = FactNodeDeserializationContext;
+
+    type SubTree = FactsSubTree<'a>;
+}
+
 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
@@ -1,8 +1,7 @@
 use std::collections::HashMap;
 
 use starknet_api::hash::HashOutput;
-use starknet_patricia::patricia_merkle_tree::filled_tree::node::{FactDbFilledNode, FilledNode};
-use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::FactNodeDeserializationContext;
+use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
 use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     NodeData,
     Preimage,
@@ -15,39 +14,36 @@ 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 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<'a, L>>(
+    subtrees: &[Layout::SubTree],
     storage: &mut impl Storage,
     key_context: &<L as HasStaticPrefix>::KeyContext,
-) -> TraversalResult<Vec<FactDbFilledNode<L>>> {
+) -> TraversalResult<Vec<FilledNode<L, Layout::NodeData>>>
+where
+    FilledNode<L, Layout::NodeData>: DBObject<DeserializeContext = Layout::DeserializationContext>,
+{
     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 = DBObject::deserialize(val, &subtree.get_root_context())?;
+        subtrees_roots.push(filled_node);
     }
     Ok(subtrees_roots)
 }
@@ -102,7 +98,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()) {
             match filled_root.data {
                 // Binary node.

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;`