starknet_committer, starknet_patricia: refactor key context (#11592)

Author: ArielElp

crates/starknet_committer/src/db/external_test_utils.rs

@@ -4,8 +4,8 @@ use serde_json::json;
 use starknet_api::hash::HashOutput;
 use starknet_patricia::patricia_merkle_tree::filled_tree::tree::{FilledTree, FilledTreeImpl};
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{
+    Leaf,
     LeafModifications,
-    LeafWithEmptyKeyContext,
     SkeletonLeaf,
 };
 use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::config::OriginalSkeletonTreeConfig;
@@ -15,7 +15,7 @@ use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::{
     UpdatedSkeletonTree,
     UpdatedSkeletonTreeImpl,
 };
-use starknet_patricia_storage::db_object::EmptyKeyContext;
+use starknet_patricia_storage::db_object::HasStaticPrefix;
 use starknet_patricia_storage::map_storage::MapStorage;
 
 use crate::db::facts_db::db::FactsNodeLayout;
@@ -27,10 +27,11 @@ pub async fn tree_computation_flow<L, TH>(
     storage: &mut MapStorage,
     root_hash: HashOutput,
     config: impl OriginalSkeletonTreeConfig,
+    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> FilledTreeImpl<L>
 where
     TH: TreeHashFunction<L> + 'static,
-    L: LeafWithEmptyKeyContext + 'static,
+    L: Leaf + 'static,
 {
     let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
     let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices);
@@ -41,7 +42,7 @@ where
         &config,
         &leaf_modifications,
         None,
-        &EmptyKeyContext,
+        key_context,
     )
     .await
     .expect("Failed to create the original skeleton tree");
@@ -68,14 +69,12 @@ where
         .expect("Failed to create the filled tree")
 }
 
-pub async fn single_tree_flow_test<
-    L: LeafWithEmptyKeyContext + 'static,
-    TH: TreeHashFunction<L> + 'static,
->(
+pub async fn single_tree_flow_test<L: Leaf + 'static, TH: TreeHashFunction<L> + 'static>(
     leaf_modifications: LeafModifications<L>,
     storage: &mut MapStorage,
     root_hash: HashOutput,
     config: impl OriginalSkeletonTreeConfig,
+    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> String {
     // Move from leaf number to actual index.
     let leaf_modifications = leaf_modifications
@@ -84,7 +83,8 @@ pub async fn single_tree_flow_test<
         .collect::<LeafModifications<L>>();
 
     let filled_tree =
-        tree_computation_flow::<L, TH>(leaf_modifications, storage, root_hash, config).await;
+        tree_computation_flow::<L, TH>(leaf_modifications, storage, root_hash, config, key_context)
+            .await;
 
     let hash_result = filled_tree.get_root_hash();
 
@@ -93,7 +93,7 @@ pub async fn single_tree_flow_test<
     let json_hash = &json!(hash_result.0.to_hex_string());
     result_map.insert("root_hash", json_hash);
     // Serlialize the storage modifications.
-    let json_storage = &json!(filled_tree.serialize(&EmptyKeyContext).unwrap());
+    let json_storage = &json!(filled_tree.serialize(key_context).unwrap());
     result_map.insert("storage_changes", json_storage);
     serde_json::to_string(&result_map).expect("serialization failed")
 }

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

@@ -1,13 +1,10 @@
 use std::collections::HashMap;
 
 use starknet_api::hash::HashOutput;
-use starknet_patricia::patricia_merkle_tree::node_data::leaf::{
-    LeafModifications,
-    LeafWithEmptyKeyContext,
-};
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
 use starknet_patricia::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonTreeResult;
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
-use starknet_patricia_storage::db_object::EmptyKeyContext;
+use starknet_patricia_storage::db_object::HasStaticPrefix;
 use starknet_patricia_storage::storage_trait::Storage;
 
 use crate::db::facts_db::db::FactsNodeLayout;
@@ -21,10 +18,11 @@ pub mod create_facts_tree_test;
 /// 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: LeafWithEmptyKeyContext>(
+pub async fn get_leaves<'a, L: Leaf>(
     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();
@@ -36,7 +34,7 @@ pub async fn get_leaves<'a, L: LeafWithEmptyKeyContext>(
         &config,
         &leaf_modifications,
         Some(&mut previous_leaves),
-        &EmptyKeyContext,
+        key_context,
     )
     .await?;
     Ok(previous_leaves)

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

@@ -3,15 +3,11 @@ use std::collections::HashMap;
 use async_trait::async_trait;
 use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
-use starknet_patricia::db_layout::{NodeLayout, NodeLayoutFor, TrieType};
+use starknet_patricia::db_layout::{NodeLayout, NodeLayoutFor};
 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,
-    LeafWithEmptyKeyContext,
-};
+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::{DBObject, EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::errors::SerializationResult;
@@ -34,7 +30,7 @@ use crate::patricia_merkle_tree::types::CompiledClassHash;
 /// have the db keys of its children.
 pub struct FactsNodeLayout {}
 
-impl<'a, L: LeafWithEmptyKeyContext> NodeLayout<'a, L> for FactsNodeLayout {
+impl<'a, L: Leaf> NodeLayout<'a, L> for FactsNodeLayout {
     type NodeData = HashOutput;
 
     type NodeDbObject = FactDbFilledNode<L>;
@@ -43,10 +39,6 @@ impl<'a, L: LeafWithEmptyKeyContext> NodeLayout<'a, L> for FactsNodeLayout {
 
     type SubTree = FactsSubTree<'a>;
 
-    fn generate_key_context(_trie_type: TrieType) -> <L as HasStaticPrefix>::KeyContext {
-        EmptyKeyContext
-    }
-
     fn get_db_object<LeafBase: Leaf + Into<L>>(
         _node_index: NodeIndex,
         key_context: &<L as HasStaticPrefix>::KeyContext,
@@ -122,8 +114,8 @@ impl<S: Storage> ForestWriter for FactsDb<S> {
         let mut serialized_forest = DbHashMap::new();
 
         // Storage tries.
-        for tree in filled_forest.storage_tries.values() {
-            serialized_forest.extend(tree.serialize(&EmptyKeyContext)?);
+        for (contract_address, tree) in &filled_forest.storage_tries {
+            serialized_forest.extend(tree.serialize(contract_address)?);
         }
 
         // Contracts and classes tries.

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

@@ -6,10 +6,10 @@ use starknet_patricia::patricia_merkle_tree::node_data::inner_node::{
     Preimage,
     PreimageMap,
 };
-use starknet_patricia::patricia_merkle_tree::node_data::leaf::LeafWithEmptyKeyContext;
+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::EmptyKeyContext;
+use starknet_patricia_storage::db_object::HasStaticPrefix;
 use starknet_patricia_storage::storage_trait::Storage;
 
 use crate::db::facts_db::db::FactsNodeLayout;
@@ -24,11 +24,12 @@ 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: LeafWithEmptyKeyContext>(
+pub async fn fetch_patricia_paths<L: Leaf>(
     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();
 
@@ -38,7 +39,14 @@ pub async fn fetch_patricia_paths<L: LeafWithEmptyKeyContext>(
 
     let main_subtree = FactsSubTree::create(sorted_leaf_indices, NodeIndex::ROOT, root_hash);
 
-    fetch_patricia_paths_inner::<L>(storage, vec![main_subtree], &mut witnesses, leaves).await?;
+    fetch_patricia_paths_inner::<L>(
+        storage,
+        vec![main_subtree],
+        &mut witnesses,
+        leaves,
+        key_context,
+    )
+    .await?;
     Ok(witnesses)
 }
 
@@ -51,21 +59,19 @@ pub async fn fetch_patricia_paths<L: LeafWithEmptyKeyContext>(
 /// 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: LeafWithEmptyKeyContext>(
+pub(crate) async fn fetch_patricia_paths_inner<'a, L: Leaf>(
     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,
-            &EmptyKeyContext,
-        )
-        .await?;
+        let filled_roots =
+            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.

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

@@ -27,6 +27,7 @@ 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;
@@ -87,6 +88,7 @@ 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/forest_trait.rs

@@ -6,6 +6,7 @@ use starknet_api::core::ContractAddress;
 use starknet_patricia::db_layout::NodeLayoutFor;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::LeafModifications;
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
+use starknet_patricia_storage::db_object::{EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::errors::SerializationResult;
 use starknet_patricia_storage::storage_trait::{DbHashMap, DbKey, DbValue, Storage};
 
@@ -82,6 +83,11 @@ where
     Layout: NodeLayoutFor<StarknetStorageValue>
         + NodeLayoutFor<ContractState>
         + NodeLayoutFor<CompiledClassHash>,
+    <Layout as NodeLayoutFor<StarknetStorageValue>>::DbLeaf:
+        HasStaticPrefix<KeyContext = ContractAddress>,
+    <Layout as NodeLayoutFor<ContractState>>::DbLeaf: HasStaticPrefix<KeyContext = EmptyKeyContext>,
+    <Layout as NodeLayoutFor<CompiledClassHash>>::DbLeaf:
+        HasStaticPrefix<KeyContext = EmptyKeyContext>,
 {
     let (contracts_trie, original_contracts_trie_leaves) = create_contracts_trie::<Layout>(
         storage,

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

@@ -3,7 +3,7 @@ use std::collections::HashMap;
 use async_trait::async_trait;
 use starknet_api::core::ContractAddress;
 use starknet_api::hash::HashOutput;
-use starknet_patricia::db_layout::{NodeLayout, NodeLayoutFor, TrieType};
+use starknet_patricia::db_layout::{NodeLayout, NodeLayoutFor};
 use starknet_patricia::patricia_merkle_tree::filled_tree::node::FilledNode;
 use starknet_patricia::patricia_merkle_tree::filled_tree::tree::FilledTree;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
@@ -48,18 +48,14 @@ pub struct IndexNodeLayout {}
 
 impl<'a, L> NodeLayout<'a, L> for IndexNodeLayout
 where
-    L: Leaf + HasStaticPrefix<KeyContext = TrieType>,
+    L: Leaf,
     TreeHashFunctionImpl: TreeHashFunction<L>,
 {
     type NodeData = EmptyNodeData;
     type NodeDbObject = IndexFilledNode<L>;
     type DeserializationContext = IndexNodeContext;
     type SubTree = IndexLayoutSubTree<'a>;
 
-    fn generate_key_context(trie_type: TrieType) -> <L as HasStaticPrefix>::KeyContext {
-        trie_type
-    }
-
     fn get_db_object<LeafBase: Leaf + Into<L>>(
         node_index: NodeIndex,
         key_context: &<L as HasStaticPrefix>::KeyContext,
@@ -120,9 +116,8 @@ impl<S: Storage> ForestWriter for IndexDb<S> {
     fn serialize_forest(filled_forest: &FilledForest) -> SerializationResult<DbHashMap> {
         let mut serialized_forest = DbHashMap::new();
 
-        // TODO(Ariel): use a different key context when FilledForest is generic over leaf types.
-        for tree in filled_forest.storage_tries.values() {
-            serialized_forest.extend(tree.serialize(&EmptyKeyContext)?);
+        for (contract_address, tree) in &filled_forest.storage_tries {
+            serialized_forest.extend(tree.serialize(contract_address)?);
         }
 
         // Contracts and classes tries.

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

@@ -1,13 +1,13 @@
 use std::sync::LazyLock;
 
-use starknet_api::core::{ClassHash, Nonce, PATRICIA_KEY_UPPER_BOUND};
+use starknet_api::core::{ClassHash, ContractAddress, Nonce, PATRICIA_KEY_UPPER_BOUND};
 use starknet_api::hash::HashOutput;
-use starknet_patricia::db_layout::TrieType;
 use starknet_patricia::patricia_merkle_tree::node_data::errors::LeafResult;
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
 use starknet_patricia_storage::db_object::{
     DBObject,
     EmptyDeserializationContext,
+    EmptyKeyContext,
     HasStaticPrefix,
 };
 use starknet_patricia_storage::errors::{DeserializationError, SerializationError};
@@ -48,34 +48,32 @@ static CONTRACTS_TREE_PREFIX: LazyLock<[u8; 32]> =
 static CLASSES_TREE_PREFIX: LazyLock<[u8; 32]> =
     LazyLock::new(|| (*FIRST_AVAILABLE_PREFIX_FELT + Felt::ONE).to_bytes_be());
 
-fn trie_type_db_prefix(trie_type: &TrieType) -> DbKeyPrefix {
-    match trie_type {
-        TrieType::ContractsTrie => DbKeyPrefix::new((&CONTRACTS_TREE_PREFIX[..]).into()),
-        TrieType::ClassesTrie => DbKeyPrefix::new((&CLASSES_TREE_PREFIX[..]).into()),
-        TrieType::StorageTrie(contract_address) => {
-            let prefix = contract_address.to_bytes_be().to_vec();
-            DbKeyPrefix::new(prefix.into())
-        }
-    }
-}
-
-macro_rules! impl_has_static_prefix_for_index_layouts {
-    ($($ty:ty),* $(,)?) => {
+macro_rules! impl_has_static_prefix_empty_context {
+    ($($ty:ty => $prefix:expr),* $(,)?) => {
         $(
             impl HasStaticPrefix for $ty {
-                type KeyContext = TrieType;
-                fn get_static_prefix(key_context: &Self::KeyContext) -> DbKeyPrefix {
-                    trie_type_db_prefix(key_context)
+                type KeyContext = EmptyKeyContext;
+                fn get_static_prefix(_key_context: &Self::KeyContext) -> DbKeyPrefix {
+                    DbKeyPrefix::new((&$prefix[..]).into())
                 }
             }
         )*
     };
 }
 
-impl_has_static_prefix_for_index_layouts! {
-    IndexLayoutContractState,
-    IndexLayoutCompiledClassHash,
-    IndexLayoutStarknetStorageValue,
+impl_has_static_prefix_empty_context! {
+    IndexLayoutContractState => CONTRACTS_TREE_PREFIX,
+    IndexLayoutCompiledClassHash => CLASSES_TREE_PREFIX,
+}
+
+impl HasStaticPrefix for IndexLayoutStarknetStorageValue {
+    type KeyContext = ContractAddress;
+
+    /// Returns the contract address, which is the storage trie prefix in index layout.
+    fn get_static_prefix(key_context: &Self::KeyContext) -> DbKeyPrefix {
+        let prefix = key_context.to_bytes_be().to_vec();
+        DbKeyPrefix::new(prefix.into())
+    }
 }
 
 macro_rules! impl_leaf_for_wrappers {