starknet_committer,starknet_patricia: enrich subtree trait to support fetch nodes

Author: ArielElp

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

@@ -1,7 +1,10 @@
 use starknet_api::hash::HashOutput;
-use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::PatriciaPrefix;
+use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::{
+    FactNodeDeserializationContext,
+    PatriciaPrefix,
+};
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
-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::storage_trait::DbKeyPrefix;
@@ -15,6 +18,7 @@ pub struct FactsSubTree<'a> {
 
 impl<'a> SubTreeTrait<'a> for FactsSubTree<'a> {
     type NodeData = HashOutput;
+    type NodeDeserializeContext = FactNodeDeserializationContext;
 
     fn create(
         sorted_leaf_indices: SortedLeafIndices<'a>,
@@ -32,8 +36,12 @@ impl<'a> SubTreeTrait<'a> for FactsSubTree<'a> {
         &self.sorted_leaf_indices
     }
 
-    fn should_traverse_unmodified_children() -> bool {
-        false
+    fn should_traverse_unmodified_child(data: Self::NodeData) -> UnmodifiedChildTraversal {
+        UnmodifiedChildTraversal::Skip(data)
+    }
+
+    fn get_root_context(&self) -> Self::NodeDeserializeContext {
+        Self::NodeDeserializeContext { is_leaf: self.is_leaf(), node_hash: self.root_hash }
     }
 
     fn get_root_prefix<L: Leaf>(
@@ -46,4 +54,8 @@ impl<'a> SubTreeTrait<'a> for FactsSubTree<'a> {
             PatriciaPrefix::InnerNode.into()
         }
     }
+
+    fn get_root_suffix(&self) -> Vec<u8> {
+        self.root_hash.0.to_bytes_be().to_vec()
+    }
 }

crates/starknet_patricia/src/patricia_merkle_tree/traversal.rs

@@ -1,3 +1,4 @@
+use starknet_api::hash::HashOutput;
 use starknet_patricia_storage::db_object::HasStaticPrefix;
 use starknet_patricia_storage::errors::{DeserializationError, StorageError};
 use starknet_patricia_storage::storage_trait::{DbKeyPrefix, PatriciaStorageError};
@@ -22,6 +23,16 @@ pub enum TraversalError {
     PatriciaStorage(#[from] PatriciaStorageError),
 }
 
+/// An enum that specifies how to treat unmodified children during the construction of the
+/// [OriginalSkeletonTree].
+pub enum UnmodifiedChildTraversal {
+    // Indicates that the child must be read since we don't have the data to create an original
+    // tree node.
+    Traverse,
+    // Indicates that the child should be skipped as its unmodified and we have its hash.
+    Skip(HashOutput),
+}
+
 pub type TraversalResult<T> = Result<T, TraversalError>;
 
 /// A trait that allows traversing a trie without knowledge of the concrete node types and data or
@@ -30,6 +41,10 @@ pub trait SubTreeTrait<'a>: Sized {
     /// Extra data a node can carry (e.g. its hash).
     type NodeData;
 
+    /// Extra context needed to deserialize a node from a raw DbValue. For more info, see
+    /// `DeserializeContext` in the `DBObject` trait.
+    type NodeDeserializeContext;
+
     /// Creates a concrete child node given its index and data.
     fn create(
         sorted_leaf_indices: SortedLeafIndices<'a>,
@@ -97,13 +112,20 @@ pub trait SubTreeTrait<'a>: Sized {
         self.get_root_index().is_leaf()
     }
 
-    /// Indicates whether unmodified children should be traversed during the construction of the
-    /// original skeleton tree.
-    fn should_traverse_unmodified_children() -> bool;
+    /// Decide whether to traverse an unmodified child during the construction of the
+    /// `OriginalSkeletonTree`.
+    fn should_traverse_unmodified_child(data: Self::NodeData) -> UnmodifiedChildTraversal;
 
     /// Returns the [DbKeyPrefix] of the root node.
     fn get_root_prefix<L: Leaf>(
         &self,
         key_context: &<L as HasStaticPrefix>::KeyContext,
     ) -> DbKeyPrefix;
+
+    /// Returns the suffix of the root's db key.
+    fn get_root_suffix(&self) -> Vec<u8>;
+
+    /// Returns the `Self::NodeDeserializeContext` that's needed to deserialize the root node from a
+    /// raw `DbValue`.
+    fn get_root_context(&self) -> Self::NodeDeserializeContext;
 }

crates/starknet_patricia/src/patricia_merkle_tree/traversal_test.rs

@@ -7,7 +7,7 @@ use starknet_patricia_storage::storage_trait::DbKeyPrefix;
 use crate::patricia_merkle_tree::external_test_utils::small_tree_index_to_full;
 use crate::patricia_merkle_tree::node_data::inner_node::{EdgePath, EdgePathLength, PathToBottom};
 use crate::patricia_merkle_tree::node_data::leaf::Leaf;
-use crate::patricia_merkle_tree::traversal::SubTreeTrait;
+use crate::patricia_merkle_tree::traversal::{SubTreeTrait, UnmodifiedChildTraversal};
 use crate::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices, SubTreeHeight};
 
 /// Test implementation for [SubTreeTrait]. It should only be used for child creation purposes,
@@ -20,6 +20,7 @@ struct TestSubTree<'a> {
 
 impl<'a> SubTreeTrait<'a> for TestSubTree<'a> {
     type NodeData = ();
+    type NodeDeserializeContext = ();
 
     fn create(
         sorted_leaf_indices: SortedLeafIndices<'a>,
@@ -37,17 +38,23 @@ impl<'a> SubTreeTrait<'a> for TestSubTree<'a> {
         &self.sorted_leaf_indices
     }
 
-    fn should_traverse_unmodified_children() -> bool {
-        false
+    fn should_traverse_unmodified_child(_child_data: Self::NodeData) -> UnmodifiedChildTraversal {
+        UnmodifiedChildTraversal::Traverse
     }
 
+    fn get_root_context(&self) -> Self::NodeDeserializeContext {}
+
     fn get_root_prefix<L: Leaf>(
         &self,
         _key_context: &<L as HasStaticPrefix>::KeyContext,
     ) -> DbKeyPrefix {
         // Dummy prefix for testing purposes (we only need a prefix when interacting with storage).
         DbKeyPrefix::new(&[0])
     }
+
+    fn get_root_suffix(&self) -> Vec<u8> {
+        vec![]
+    }
 }
 
 /// case::single_right_child