starknet_committer: enrich layout trait to abstract skeleton creation

Author: ArielElp

crates/starknet_committer/src/db/db_layout.rs

@@ -1,24 +1,38 @@
+use starknet_api::hash::HashOutput;
 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;
+use starknet_patricia_storage::db_object::{DBObject, HasStaticPrefix};
+
+use crate::db::index_db::leaves::TrieType;
 
 /// 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: Clone;
+    type NodeData: Clone + From<HashOutput>;
 
     /// The context needed to deserialize the node from a raw
     /// [starknet_patricia_storage::storage_trait::DbValue].
     type DeserializationContext;
 
+    /// The storage representation of the node.
+    type NodeDbObject: DBObject<DeserializeContext = Self::DeserializationContext>;
+
     /// The type of the subtree that is used to traverse the trie.
     type SubTree: SubTreeTrait<
             'a,
             NodeData = Self::NodeData,
             NodeDeserializeContext = Self::DeserializationContext,
         >;
+
+    /// Generates the key context for the given trie type. Used for reading nodes of a specific
+    /// tree (contracts, classes, or storage), to construct a skeleton tree.
+    fn generate_key_context(trie_type: TrieType) -> <L as HasStaticPrefix>::KeyContext;
+
+    /// Converts a node db object to a filled node. Used during the trie traversal for the skeleton
+    /// construction.
+    fn get_filled_node(node_db_object: Self::NodeDbObject) -> FilledNode<L, Self::NodeData>;
 }

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

@@ -9,7 +9,7 @@ use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::tree::{
 };
 use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait;
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
-use starknet_patricia_storage::db_object::HasStaticPrefix;
+use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::storage_trait::Storage;
 
 use crate::db::facts_db::db::FactsNodeLayout;
@@ -21,13 +21,15 @@ use crate::patricia_merkle_tree::tree::OriginalSkeletonTrieConfig;
 #[path = "create_facts_tree_test.rs"]
 pub mod create_facts_tree_test;
 
-pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
+pub async fn create_original_skeleton_tree_and_get_previous_leaves<
+    'a,
+    L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext>,
+>(
     storage: &mut impl Storage,
     root_hash: HashOutput,
     sorted_leaf_indices: SortedLeafIndices<'a>,
     leaf_modifications: &LeafModifications<L>,
     config: &impl OriginalSkeletonTreeConfig,
-    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> OriginalSkeletonTreeResult<(OriginalSkeletonTreeImpl<'a>, HashMap<NodeIndex, L>)> {
     if sorted_leaf_indices.is_empty() {
         let unmodified = OriginalSkeletonTreeImpl::create_unmodified(root_hash);
@@ -49,7 +51,7 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
         leaf_modifications,
         config,
         Some(&mut leaves),
-        key_context,
+        &EmptyKeyContext,
     )
     .await?;
     Ok((skeleton_tree, leaves))
@@ -58,11 +60,10 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
 /// Prepares the OS inputs by fetching paths to the given leaves (i.e. their induced Skeleton tree).
 /// Note that ATM, the Rust committer does not manage history and is not used for storage proofs;
 /// Thus, this function assumes facts layout.
-pub async fn get_leaves<'a, L: Leaf>(
+pub async fn get_leaves<'a, L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext>>(
     storage: &mut impl Storage,
     root_hash: HashOutput,
     sorted_leaf_indices: SortedLeafIndices<'a>,
-    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> OriginalSkeletonTreeResult<HashMap<NodeIndex, L>> {
     let config = OriginalSkeletonTrieConfig::default();
     let leaf_modifications = LeafModifications::new();
@@ -72,7 +73,6 @@ pub async fn get_leaves<'a, L: Leaf>(
         sorted_leaf_indices,
         &leaf_modifications,
         &config,
-        key_context,
     )
     .await?;
     Ok(previous_leaves)

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

@@ -3,11 +3,12 @@ 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::{FactDbFilledNode, FilledNode};
 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::{Leaf, LeafModifications};
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
-use starknet_patricia_storage::db_object::EmptyKeyContext;
+use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::errors::SerializationResult;
 use starknet_patricia_storage::map_storage::MapStorage;
 use starknet_patricia_storage::storage_trait::{DbHashMap, Storage};
@@ -16,6 +17,7 @@ use crate::block_committer::input::{ReaderConfig, StarknetStorageValue};
 use crate::db::db_layout::NodeLayout;
 use crate::db::facts_db::types::{FactsDbInitialRead, FactsSubTree};
 use crate::db::forest_trait::{ForestReader, ForestWriter};
+use crate::db::index_db::leaves::TrieType;
 use crate::db::trie_traversal::{create_classes_trie, create_contracts_trie, create_storage_tries};
 use crate::forest::filled_forest::FilledForest;
 use crate::forest::forest_errors::ForestResult;
@@ -30,12 +32,25 @@ use crate::patricia_merkle_tree::types::CompiledClassHash;
 /// have the db keys of its children.
 pub struct FactsNodeLayout {}
 
-impl<'a, L: Leaf> NodeLayout<'a, L> for FactsNodeLayout {
+impl<'a, L: Leaf> NodeLayout<'a, L> for FactsNodeLayout
+where
+    L: HasStaticPrefix<KeyContext = EmptyKeyContext>,
+{
     type NodeData = HashOutput;
 
+    type NodeDbObject = FactDbFilledNode<L>;
+
     type DeserializationContext = FactNodeDeserializationContext;
 
     type SubTree = FactsSubTree<'a>;
+
+    fn generate_key_context(_trie_type: TrieType) -> <L as HasStaticPrefix>::KeyContext {
+        EmptyKeyContext
+    }
+
+    fn get_filled_node(node_db_object: Self::NodeDbObject) -> FilledNode<L, Self::NodeData> {
+        node_db_object
+    }
 }
 
 pub struct FactsDb<S: Storage> {

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

@@ -9,7 +9,7 @@ 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::HasStaticPrefix;
+use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::storage_trait::Storage;
 
 use crate::db::facts_db::db::FactsNodeLayout;
@@ -24,12 +24,11 @@ pub mod traversal_test;
 /// given tree according to the `root_hash`.
 /// If `leaves` is not `None`, it also fetches the modified leaves and inserts them into the
 /// provided map.
-pub async fn fetch_patricia_paths<L: Leaf>(
+pub async fn fetch_patricia_paths<L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext>>(
     storage: &mut impl Storage,
     root_hash: HashOutput,
     sorted_leaf_indices: SortedLeafIndices<'_>,
     leaves: Option<&mut HashMap<NodeIndex, L>>,
-    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> TraversalResult<PreimageMap> {
     let mut witnesses = PreimageMap::new();
 
@@ -39,14 +38,7 @@ pub async fn fetch_patricia_paths<L: Leaf>(
 
     let main_subtree = FactsSubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash);
 
-    fetch_patricia_paths_inner::<L>(
-        storage,
-        vec![main_subtree],
-        &mut witnesses,
-        leaves,
-        key_context,
-    )
-    .await?;
+    fetch_patricia_paths_inner::<L>(storage, vec![main_subtree], &mut witnesses, leaves).await?;
     Ok(witnesses)
 }
 
@@ -59,19 +51,24 @@ pub async fn fetch_patricia_paths<L: Leaf>(
 /// inner nodes in their paths and their siblings.
 /// If `leaves` is not `None`, it also fetches the modified leaves and inserts them into the
 /// provided map.
-pub(crate) async fn fetch_patricia_paths_inner<'a, L: Leaf>(
+pub(crate) async fn fetch_patricia_paths_inner<
+    'a,
+    L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext>,
+>(
     storage: &mut impl Storage,
     subtrees: Vec<FactsSubTree<'a>>,
     witnesses: &mut PreimageMap,
     mut leaves: Option<&mut HashMap<NodeIndex, L>>,
-    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> TraversalResult<()> {
     let mut current_subtrees = subtrees;
     let mut next_subtrees = Vec::new();
     while !current_subtrees.is_empty() {
-        let filled_roots =
-            get_roots_from_storage::<L, FactsNodeLayout>(&current_subtrees, storage, key_context)
-                .await?;
+        let filled_roots = get_roots_from_storage::<L, FactsNodeLayout>(
+            &current_subtrees,
+            storage,
+            &EmptyKeyContext,
+        )
+        .await?;
         for (filled_root, subtree) in filled_roots.into_iter().zip(current_subtrees.iter()) {
             match filled_root.data {
                 // Binary node.

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

@@ -27,7 +27,6 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
 };
 use starknet_patricia::patricia_merkle_tree::traversal::SubTreeTrait;
 use starknet_patricia::patricia_merkle_tree::types::{SortedLeafIndices, SubTreeHeight};
-use starknet_patricia_storage::db_object::EmptyKeyContext;
 use starknet_patricia_storage::map_storage::MapStorage;
 use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue};
 use starknet_types_core::felt::Felt;
@@ -88,7 +87,6 @@ async fn test_fetch_patricia_paths_inner_impl(
         vec![main_subtree],
         &mut nodes,
         Some(&mut fetched_leaves),
-        &EmptyKeyContext,
     )
     .await
     .unwrap();

crates/starknet_committer/src/db/index_db/serde_tests.rs

@@ -26,7 +26,7 @@ use crate::db::index_db::leaves::{
     IndexLayoutContractState,
     IndexLayoutStarknetStorageValue,
 };
-use crate::db::index_db::types::{IndexFilledNode, IndexNodeContext};
+use crate::db::index_db::types::{EmptyNodeData, IndexFilledNode, IndexNodeContext};
 use crate::hash_function::hash::TreeHashFunctionImpl;
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use crate::patricia_merkle_tree::types::CompiledClassHash;
@@ -35,7 +35,9 @@ use crate::patricia_merkle_tree::types::CompiledClassHash;
 ///
 /// Useful for using the same data for computing the hash and creating an index layout leaf
 /// instance.
-fn index_leaf_data_from_hash_data<L: Leaf>(data: NodeData<L, HashOutput>) -> NodeData<L, ()> {
+fn index_leaf_data_from_hash_data<L: Leaf>(
+    data: NodeData<L, HashOutput>,
+) -> NodeData<L, EmptyNodeData> {
     match data {
         NodeData::Binary(_) | NodeData::Edge(_) => {
             unreachable!("this helper is intended for leaf-only test data")
@@ -85,15 +87,15 @@ fn starknet_storage_value_leaf_136_bits() -> IndexLayoutStarknetStorageValue {
 fn binary_node() -> IndexFilledNode<IndexLayoutContractState> {
     IndexFilledNode(FilledNode {
         hash: HashOutput(Felt::from(1)),
-        data: NodeData::Binary(BinaryData { left_data: (), right_data: () }),
+        data: NodeData::Binary(BinaryData { left_data: EmptyNodeData, right_data: EmptyNodeData }),
     })
 }
 
 fn edge_node_short_path_len_3() -> IndexFilledNode<IndexLayoutContractState> {
     IndexFilledNode(FilledNode {
         hash: HashOutput(Felt::from(1)),
         data: NodeData::Edge(EdgeData {
-            bottom_data: (),
+            bottom_data: EmptyNodeData,
             // 110, right, right, left
             path_to_bottom: PathToBottom::new(
                 EdgePath(U256::from(6_u128)),
@@ -108,7 +110,7 @@ fn edge_node_short_path_len_10() -> IndexFilledNode<IndexLayoutContractState> {
     IndexFilledNode(FilledNode {
         hash: HashOutput(Felt::from(1)),
         data: NodeData::Edge(EdgeData {
-            bottom_data: (),
+            bottom_data: EmptyNodeData,
             // 0...0 seven times followed by 110
             path_to_bottom: PathToBottom::new(
                 EdgePath(U256::from(6_u128)),
@@ -123,7 +125,7 @@ fn edge_node_path_divisible_by_8() -> IndexFilledNode<IndexLayoutContractState>
     IndexFilledNode(FilledNode {
         hash: HashOutput(Felt::from(1)),
         data: NodeData::Edge(EdgeData {
-            bottom_data: (),
+            bottom_data: EmptyNodeData,
             path_to_bottom: PathToBottom::new(
                 // 1...1 24 times
                 EdgePath(U256::from((1_u128 << 24) - 1)),
@@ -138,7 +140,7 @@ fn edge_node_path_not_divisible_by_8() -> IndexFilledNode<IndexLayoutContractSta
     IndexFilledNode(FilledNode {
         hash: HashOutput(Felt::from(1)),
         data: NodeData::Edge(EdgeData {
-            bottom_data: (),
+            bottom_data: EmptyNodeData,
             // 000 followed by 1...1 19 times
             path_to_bottom: PathToBottom::new(
                 EdgePath(U256::from((1_u128 << 19) - 1)),
@@ -153,7 +155,7 @@ fn edge_node_long_zero_path() -> IndexFilledNode<IndexLayoutContractState> {
     IndexFilledNode(FilledNode {
         hash: HashOutput(Felt::from(1)),
         data: NodeData::Edge(EdgeData {
-            bottom_data: (),
+            bottom_data: EmptyNodeData,
             // 0...0 path of length 250
             path_to_bottom: PathToBottom::new(
                 EdgePath(U256::ZERO),
@@ -168,7 +170,7 @@ fn edge_node_long_non_zero_path() -> IndexFilledNode<IndexLayoutContractState> {
     IndexFilledNode(FilledNode {
         hash: HashOutput(Felt::from(1)),
         data: NodeData::Edge(EdgeData {
-            bottom_data: (),
+            bottom_data: EmptyNodeData,
             // 1 followed by 250 zeros path of length 251
             path_to_bottom: PathToBottom::new(
                 EdgePath(U256::from(1u8) << 250),

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

@@ -28,8 +28,17 @@ use crate::hash_function::hash::TreeHashFunctionImpl;
 // In index layout, for binary nodes, only the hash is stored.
 pub(crate) const INDEX_LAYOUT_BINARY_BYTES: usize = SERIALIZE_HASH_BYTES;
 
+#[derive(Clone, Debug, PartialEq, Eq)]
+pub struct EmptyNodeData;
+
+impl From<HashOutput> for EmptyNodeData {
+    fn from(_hash: HashOutput) -> Self {
+        Self
+    }
+}
+
 #[derive(PartialEq, Debug)]
-pub struct IndexFilledNode<L: Leaf>(pub FilledNode<L, ()>);
+pub struct IndexFilledNode<L: Leaf>(pub FilledNode<L, EmptyNodeData>);
 
 pub struct IndexNodeContext {
     pub is_leaf: bool,
@@ -82,7 +91,10 @@ where
         } else if value.0.len() == INDEX_LAYOUT_BINARY_BYTES {
             Ok(Self(FilledNode {
                 hash: HashOutput(Felt::from_bytes_be_slice(&value.0)),
-                data: NodeData::Binary(BinaryData { left_data: (), right_data: () }),
+                data: NodeData::Binary(BinaryData {
+                    left_data: EmptyNodeData,
+                    right_data: EmptyNodeData,
+                }),
             }))
         }
         // Edge nodes are always serailized to more than INDEX_LAYOUT_BINARY_BYTES bytes.
@@ -112,8 +124,8 @@ where
 
             Ok(Self(FilledNode {
                 hash: node_hash,
-                data: NodeData::Edge(EdgeData::<()> {
-                    bottom_data: (),
+                data: NodeData::Edge(EdgeData {
+                    bottom_data: EmptyNodeData,
                     path_to_bottom: PathToBottom::new(
                         EdgePath(path),
                         EdgePathLength::new(bit_len)

crates/starknet_committer/src/db/trie_traversal.rs

@@ -288,7 +288,10 @@ pub(crate) fn log_warning_for_empty_leaves<L: Leaf, T: Borrow<NodeIndex> + Debug
     Ok(())
 }
 
-pub async fn create_original_skeleton_tree<'a, L: Leaf>(
+pub async fn create_original_skeleton_tree<
+    'a,
+    L: Leaf + HasStaticPrefix<KeyContext = EmptyKeyContext>,
+>(
     storage: &mut impl Storage,
     root_hash: HashOutput,
     sorted_leaf_indices: SortedLeafIndices<'a>,
@@ -368,7 +371,6 @@ pub async fn create_contracts_trie<'a>(
         contracts_trie_sorted_indices,
         &HashMap::new(),
         &OriginalSkeletonTrieConfig::new_for_contracts_trie(),
-        &EmptyKeyContext,
     )
     .await?)
 }