starknet_committer,starknet_committer_and_os_cli,starknet_patricia: make fetch_nodes layout dependent

Author: ArielElp

crates/starknet_committer/src/block_committer/random_structs.rs

@@ -19,7 +19,6 @@ use starknet_patricia::patricia_merkle_tree::external_test_utils::{
     get_random_u256,
     u256_try_into_felt,
 };
-use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
 use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     BinaryData,
     EdgeData,
@@ -34,6 +33,7 @@ use starknet_types_core::felt::Felt;
 use strum::IntoEnumIterator;
 
 use crate::block_committer::input::StarknetStorageValue;
+use crate::db::facts_db::types::FactDbFilledNode;
 use crate::forest::filled_forest::FilledForest;
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use crate::patricia_merkle_tree::types::{
@@ -178,7 +178,7 @@ impl RandomValue for NodeIndex {
 
 macro_rules! random_filled_node {
     ($leaf:ty) => {
-        impl RandomValue for FilledNode<$leaf> {
+        impl RandomValue for FactDbFilledNode<$leaf> {
             fn random<R: Rng>(rng: &mut R, max: Option<U256>) -> Self {
                 Self { data: NodeData::random(rng, max), hash: HashOutput::random(rng, max) }
             }
@@ -216,11 +216,13 @@ macro_rules! random_filled_tree {
                 }
                 .as_usize();
 
-                let mut nodes: Vec<(NodeIndex, FilledNode<$leaf>)> = (0..max_node_number)
-                    .map(|_| (NodeIndex::random(rng, max_size), FilledNode::random(rng, max_size)))
+                let mut nodes: Vec<(NodeIndex, FactDbFilledNode<$leaf>)> = (0..max_node_number)
+                    .map(|_| {
+                        (NodeIndex::random(rng, max_size), FactDbFilledNode::random(rng, max_size))
+                    })
                     .collect();
 
-                nodes.push((NodeIndex::ROOT, FilledNode::random(rng, max_size)));
+                nodes.push((NodeIndex::ROOT, FactDbFilledNode::random(rng, max_size)));
 
                 Self {
                     tree_map: nodes.into_iter().collect(),

crates/starknet_committer/src/db.rs

@@ -1,3 +1,4 @@
+pub mod db_layout;
 #[cfg(any(feature = "testing", test))]
 pub mod external_test_utils;
 pub mod facts_db;

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>;
+}

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::Binary(BinaryData { left_hash, right_hash }) => {
+                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_hash, right_hash);
+                        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_hash, path_to_bottom }) => {
-                    skeleton_tree
-                        .nodes
-                        .insert(subtree.root_index, OriginalSkeletonNode::Edge(path_to_bottom));
+                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() {
-                        skeleton_tree.nodes.insert(
-                            path_to_bottom.bottom_index(subtree.root_index),
-                            OriginalSkeletonNode::UnmodifiedSubTree(bottom_hash),
-                        );
+                        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;
                     }
-                    // Parse bottom.
-                    let (bottom_subtree, previously_empty_leaves_indices) =
-                        subtree.get_bottom_subtree(&path_to_bottom, bottom_hash);
+
                     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));
     }
 }
 

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

@@ -2,22 +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::NodeContext;
+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::DBObject;
+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};
@@ -30,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 = NodeContext;
+    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.