starknet_committer: make fetch nodes depend on db layout

Author: ArielElp

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;

crates/starknet_committer/src/db/db_layout.rs

@@ -0,0 +1,24 @@
+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;
+
+/// Specifies the trie db layout.
+pub trait NodeLayout<'a, L: Leaf>
+where
+    FilledNode<L, Self::NodeData>: DBObject<DeserializeContext = Self::DeserializationContext>,
+{
+    /// Additional data that a node stores about its children.
+    type NodeData: Copy;
+
+    /// The context needed to deserialize the node from a raw
+    /// [starknet_patricia_storage::storage_trait::DbValue].
+    type DeserializationContext;
+
+    /// The type of the subtree that is used to traverse the trie.
+    type SubTree: SubTreeTrait<
+            'a,
+            NodeData = Self::NodeData,
+            NodeDeserializeContext = Self::DeserializationContext,
+        >;
+}

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,72 @@ 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();
+    let should_fetch_modified_leaves =
+        config.compare_modified_leaves() || previous_leaves.is_some();
     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?;
-        for (filled_root, subtree) in filled_roots.into_iter().zip(current_subtrees.iter()) {
+            get_roots_from_storage::<L, Layout>(&current_subtrees, storage, key_context).await?;
+        for (filled_root, subtree) in filled_roots.into_iter().zip(current_subtrees.into_iter()) {
+            if subtree.is_unmodified() {
+                handle_unmodified_subtree(skeleton_tree, &mut next_subtrees, filled_root, subtree);
+                continue;
+            }
+            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,
-                            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(
+                    handle_child_subtree(
                         skeleton_tree,
                         &mut next_subtrees,
                         left_subtree,
+                        left_data,
                         should_fetch_modified_leaves,
                     );
-                    handle_subtree(
+                    handle_child_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 let Some(ref mut leaves) = previous_leaves {
                         leaves.extend(
                             previously_empty_leaves_indices
@@ -115,28 +122,25 @@ async fn fetch_nodes<'a, L: Leaf>(
                         config,
                     )?;
 
-                    handle_subtree(
+                    handle_child_subtree(
                         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.")
-                    } else {
-                        // Modified leaf.
-                        if config.compare_modified_leaves()
-                            && L::compare(leaf_modifications, &subtree.root_index, &previous_leaf)?
-                        {
-                            log_trivial_modification!(subtree.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);
-                        }
+                    // Modified leaf.
+                    if config.compare_modified_leaves()
+                        && L::compare(leaf_modifications, &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(root_index, previous_leaf);
                     }
                 }
             }
@@ -168,7 +172,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 +206,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,
@@ -238,21 +242,86 @@ pub async fn get_leaves<'a, L: Leaf>(
     Ok(previous_leaves)
 }
 
+/// Adds the subtree root to the skeleton tree. If the root is an edge node, and
+/// `should_traverse_unmodified_child` is `Skip`, then the corresponding bottom node is also
+/// added to the skeleton. Otherwise, the bottom subtree is added to the next subtrees.
+fn handle_unmodified_subtree<'a, L: Leaf, SubTree: SubTreeTrait<'a>>(
+    skeleton_tree: &mut OriginalSkeletonTreeImpl<'a>,
+    next_subtrees: &mut Vec<SubTree>,
+    filled_root: FilledNode<L, SubTree::NodeData>,
+    subtree: SubTree,
+) where
+    SubTree::NodeData: Copy,
+{
+    // Sanity check.
+    assert!(subtree.is_unmodified(), "Called `handle_unmodified_subtree` for a modified subtree.");
+
+    let root_index = subtree.get_root_index();
+
+    match filled_root.data {
+        NodeData::Edge(EdgeData { bottom_data, path_to_bottom }) => {
+            // Even if a subtree rooted at an edge node is unmodified, we still need an
+            // `OriginalSkeletonNode::Edge` node in the skeleton in case we need to manipulate it
+            // later (e.g. unify the edge node with an ancestor edge node).
+            skeleton_tree.nodes.insert(root_index, OriginalSkeletonNode::Edge(path_to_bottom));
+            match SubTree::should_traverse_unmodified_child(bottom_data) {
+                UnmodifiedChildTraversal::Traverse => {
+                    let (bottom_subtree, _) =
+                        subtree.get_bottom_subtree(&path_to_bottom, bottom_data);
+                    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),
+                    );
+                }
+            }
+        }
+        NodeData::Binary(_) | NodeData::Leaf(_) => {
+            skeleton_tree
+                .nodes
+                .insert(root_index, OriginalSkeletonNode::UnmodifiedSubTree(filled_root.hash));
+        }
+    }
+}
+
 /// Handles a subtree referred by an edge or a binary node. Decides whether we deserialize the
-/// referred subtree or not.
-fn handle_subtree<'a>(
+/// referred subtree or not, and if we should continue traversing the child's direction.
+fn handle_child_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_modified = !subtree.is_unmodified();
+
+    // Internal node → always traverse.
+    if !subtree.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),
+            );
+        }
     }
 }
 

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,21 @@ use crate::patricia_merkle_tree::tree::{
 };
 use crate::patricia_merkle_tree::types::CompiledClassHash;
 
+/// Facts DB node layout.
+///
+/// In a facts DB, the storage keys are node hashes and the values are preimages. In particular,
+/// each nodes holds its child node hashes. In this layout, only once the  parent is traversed we
+/// have the db keys of its children.
+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.