starknet_committer,starknet_patricia,starknet_patricia_storage: add KeyContext

Author: ArielElp

crates/starknet_committer/src/db/external_test_utils.rs

@@ -15,6 +15,7 @@ use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::{
     UpdatedSkeletonTree,
     UpdatedSkeletonTreeImpl,
 };
+use starknet_patricia_storage::db_object::HasStaticPrefix;
 use starknet_patricia_storage::map_storage::MapStorage;
 
 use crate::db::facts_db::create_facts_tree::create_original_skeleton_tree;
@@ -27,7 +28,7 @@ pub async fn tree_computation_flow<L, TH>(
 ) -> FilledTreeImpl<L>
 where
     TH: TreeHashFunction<L> + 'static,
-    L: Leaf + 'static,
+    L: Leaf + HasStaticPrefix<KeyContext = ()> + 'static,
 {
     let mut sorted_leaf_indices: Vec<NodeIndex> = leaf_modifications.keys().copied().collect();
     let sorted_leaf_indices = SortedLeafIndices::new(&mut sorted_leaf_indices);
@@ -37,6 +38,7 @@ where
         sorted_leaf_indices,
         &config,
         &leaf_modifications,
+        &(),
     )
     .await
     .expect("Failed to create the original skeleton tree");
@@ -63,7 +65,10 @@ where
         .expect("Failed to create the filled tree")
 }
 
-pub async fn single_tree_flow_test<L: Leaf + 'static, TH: TreeHashFunction<L> + 'static>(
+pub async fn single_tree_flow_test<
+    L: Leaf + HasStaticPrefix<KeyContext = ()> + 'static,
+    TH: TreeHashFunction<L> + 'static,
+>(
     leaf_modifications: LeafModifications<L>,
     storage: &mut MapStorage,
     root_hash: HashOutput,
@@ -85,7 +90,7 @@ pub async fn single_tree_flow_test<L: Leaf + 'static, TH: TreeHashFunction<L> +
     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());
+    let json_storage = &json!(filled_tree.serialize(&()));
     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

@@ -20,6 +20,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::storage_trait::Storage;
 use tracing::warn;
 
@@ -48,13 +49,15 @@ async fn fetch_nodes<'a, L: Leaf>(
     leaf_modifications: &LeafModifications<L>,
     config: &impl OriginalSkeletonTreeConfig<L>,
     mut previous_leaves: Option<&mut HashMap<NodeIndex, L>>,
+    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> OriginalSkeletonTreeResult<()> {
     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).await?;
+        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()) {
             match filled_root.data {
                 // Binary node.
@@ -150,6 +153,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf>(
     sorted_leaf_indices: SortedLeafIndices<'a>,
     config: &impl OriginalSkeletonTreeConfig<L>,
     leaf_modifications: &LeafModifications<L>,
+    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> OriginalSkeletonTreeResult<OriginalSkeletonTreeImpl<'a>> {
     if sorted_leaf_indices.is_empty() {
         return Ok(OriginalSkeletonTreeImpl::create_unmodified(root_hash));
@@ -171,6 +175,7 @@ pub async fn create_original_skeleton_tree<'a, L: Leaf>(
         leaf_modifications,
         config,
         None,
+        key_context,
     )
     .await?;
     Ok(skeleton_tree)
@@ -182,6 +187,7 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
     sorted_leaf_indices: SortedLeafIndices<'a>,
     leaf_modifications: &LeafModifications<L>,
     config: &impl OriginalSkeletonTreeConfig<L>,
+    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);
@@ -203,6 +209,7 @@ pub async fn create_original_skeleton_tree_and_get_previous_leaves<'a, L: Leaf>(
         leaf_modifications,
         config,
         Some(&mut leaves),
+        key_context,
     )
     .await?;
     Ok((skeleton_tree, leaves))
@@ -212,6 +219,7 @@ 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 = NoCompareOriginalSkeletonTrieConfig::default();
     let leaf_modifications = LeafModifications::new();
@@ -221,6 +229,7 @@ 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/create_facts_tree_test.rs

@@ -217,6 +217,7 @@ async fn test_create_tree(
         sorted_leaf_indices,
         &config,
         &leaf_modifications,
+        &(),
     )
     .await
     .unwrap();
@@ -225,7 +226,7 @@ async fn test_create_tree(
 
 pub(crate) fn create_mock_leaf_entry(val: u128) -> (DbKey, DbValue) {
     let leaf = MockLeaf(Felt::from(val));
-    (leaf.get_db_key(&leaf.0.to_bytes_be()), leaf.serialize())
+    (leaf.get_db_key(&(), &leaf.0.to_bytes_be()), leaf.serialize())
 }
 
 fn create_mock_leaf_modifications(

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

@@ -54,6 +54,7 @@ impl<S: Storage> FactsDb<S> {
             contracts_trie_sorted_indices,
             &HashMap::new(),
             &OriginalSkeletonContractsTrieConfig::new(),
+            &(),
         )
         .await?)
     }
@@ -82,6 +83,7 @@ impl<S: Storage> FactsDb<S> {
                 *sorted_leaf_indices,
                 &config,
                 updates,
+                &(),
             )
             .await?;
             storage_tries.insert(*address, original_skeleton);
@@ -104,6 +106,7 @@ impl<S: Storage> FactsDb<S> {
             contracts_trie_sorted_indices,
             &config,
             actual_classes_updates,
+            &(),
         )
         .await?)
     }

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

@@ -10,6 +10,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::errors::StorageError;
 use starknet_patricia_storage::storage_trait::{create_db_key, DbKey, Storage};
 
@@ -23,12 +24,16 @@ pub mod traversal_test;
 pub async fn calculate_subtrees_roots<'a, L: Leaf>(
     subtrees: &[FactsSubTree<'a>],
     storage: &mut impl Storage,
+    key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> TraversalResult<Vec<FilledNode<L>>> {
     let mut subtrees_roots = vec![];
     let db_keys: Vec<DbKey> = subtrees
         .iter()
         .map(|subtree| {
-            create_db_key(subtree.get_root_prefix::<L>(), &subtree.root_hash.0.to_bytes_be())
+            create_db_key(
+                subtree.get_root_prefix::<L>(key_context),
+                &subtree.root_hash.0.to_bytes_be(),
+            )
         })
         .collect();
 
@@ -49,6 +54,7 @@ pub async fn fetch_patricia_paths<L: Leaf>(
     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();
 
@@ -58,7 +64,14 @@ pub async fn fetch_patricia_paths<L: Leaf>(
 
     let main_subtree = FactsSubTree { sorted_leaf_indices, root_index: 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)
 }
 
@@ -74,11 +87,13 @@ pub(crate) async fn fetch_patricia_paths_inner<'a, L: Leaf>(
     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 = calculate_subtrees_roots::<L>(&current_subtrees, storage).await?;
+        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()) {
             // Always insert root.
             // No need to insert an unmodified node (which is not the root), because its parent is

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

@@ -86,6 +86,7 @@ async fn test_fetch_patricia_paths_inner_impl(
         vec![main_subtree],
         &mut nodes,
         Some(&mut fetched_leaves),
+        &(),
     )
     .await
     .unwrap();

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

@@ -3,6 +3,7 @@ use starknet_patricia::patricia_merkle_tree::filled_tree::node_serde::PatriciaPr
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
 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::storage_trait::DbKeyPrefix;
 
 #[derive(Debug, PartialEq)]
@@ -34,10 +35,12 @@ impl<'a> SubTreeTrait<'a> for FactsSubTree<'a> {
     fn should_traverse_unmodified_children() -> bool {
         false
     }
-
-    fn get_root_prefix<L: Leaf>(&self) -> DbKeyPrefix {
+    fn get_root_prefix<L: Leaf>(
+        &self,
+        key_context: &<L as HasStaticPrefix>::KeyContext,
+    ) -> DbKeyPrefix {
         if self.is_leaf() {
-            PatriciaPrefix::Leaf(L::get_static_prefix()).into()
+            PatriciaPrefix::Leaf(L::get_static_prefix(key_context)).into()
         } else {
             PatriciaPrefix::InnerNode.into()
         }

crates/starknet_committer/src/forest/filled_forest.rs

@@ -59,12 +59,11 @@ impl FilledForest {
     fn serialize_trees(&self) -> DbHashMap {
         self.storage_tries
             .values()
-            .flat_map(|tree| tree.serialize().into_iter())
-            .chain(self.contracts_trie.serialize())
-            .chain(self.classes_trie.serialize())
+            .flat_map(|tree| tree.serialize(&()).into_iter())
+            .chain(self.contracts_trie.serialize(&()))
+            .chain(self.classes_trie.serialize(&()))
             .collect()
     }
-
     async fn write_hashmap_to_storage(
         &self,
         storage: &mut impl Storage,
@@ -164,8 +163,8 @@ impl FilledForest {
             contract_address_to_storage_skeleton.len(),
             contract_address_to_storage_updates.len()
         );
-        // `contract_address_to_storage_updates` includes all modified contracts, even those with
-        // unmodified storage, see StateDiff::actual_storage_updates().
+        // `contract_address_to_storage_updates` includes all modified contracts, even those
+        // with unmodified storage, see StateDiff::actual_storage_updates().
         for (contract_address, storage_updates) in contract_address_to_storage_updates {
             let node_index = contract_address_into_node_index(&contract_address);
             let original_contract_state = original_contracts_trie_leaves

crates/starknet_committer/src/forest/skeleton_forest_test.rs

@@ -50,12 +50,12 @@ macro_rules! compare_skeleton_tree {
 
 pub(crate) fn create_storage_leaf_entry(val: u128) -> (DbKey, DbValue) {
     let leaf = StarknetStorageValue(Felt::from(val));
-    (leaf.get_db_key(&leaf.0.to_bytes_be()), leaf.serialize())
+    (leaf.get_db_key(&(), &leaf.0.to_bytes_be()), leaf.serialize())
 }
 
 pub(crate) fn create_compiled_class_leaf_entry(val: u128) -> (DbKey, DbValue) {
     let leaf = CompiledClassHash(Felt::from(val));
-    (leaf.get_db_key(&leaf.0.to_bytes_be()), leaf.serialize())
+    (leaf.get_db_key(&(), &leaf.0.to_bytes_be()), leaf.serialize())
 }
 
 pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
@@ -65,7 +65,7 @@ pub(crate) fn create_contract_state_leaf_entry(val: u128) -> (DbKey, DbValue) {
         storage_root_hash: HashOutput(felt),
         class_hash: ClassHash(felt),
     };
-    (leaf.get_db_key(&felt.to_bytes_be()), leaf.serialize())
+    (leaf.get_db_key(&(), &felt.to_bytes_be()), leaf.serialize())
 }
 
 // This test uses addition hash for simplicity (i.e hash(a,b) = a + b).

crates/starknet_committer/src/patricia_merkle_tree/leaf/leaf_impl.rs

@@ -22,7 +22,8 @@ pub struct ContractState {
 }
 
 impl HasStaticPrefix for StarknetStorageValue {
-    fn get_static_prefix() -> DbKeyPrefix {
+    type KeyContext = ();
+    fn get_static_prefix(_key_context: &Self::KeyContext) -> DbKeyPrefix {
         CommitterLeafPrefix::StorageLeaf.into()
     }
 }
@@ -41,7 +42,8 @@ impl Leaf for StarknetStorageValue {
 }
 
 impl HasStaticPrefix for CompiledClassHash {
-    fn get_static_prefix() -> DbKeyPrefix {
+    type KeyContext = ();
+    fn get_static_prefix(_key_context: &Self::KeyContext) -> DbKeyPrefix {
         CommitterLeafPrefix::CompiledClassLeaf.into()
     }
 }
@@ -60,7 +62,8 @@ impl Leaf for CompiledClassHash {
 }
 
 impl HasStaticPrefix for ContractState {
-    fn get_static_prefix() -> DbKeyPrefix {
+    type KeyContext = ();
+    fn get_static_prefix(_key_context: &Self::KeyContext) -> DbKeyPrefix {
         CommitterLeafPrefix::StateTreeLeaf.into()
     }
 }