starknet_committer,starknet_patricia: layout-dependent db-key separator

Author: ArielElp

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

@@ -7,7 +7,7 @@ use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
 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;
+use starknet_patricia_storage::storage_trait::{create_db_key, DbKey, DbKeyPrefix};
 
 use crate::block_committer::input::InputContext;
 
@@ -46,19 +46,17 @@ impl<'a> SubTreeTrait<'a> for FactsSubTree<'a> {
         Self::NodeDeserializeContext { is_leaf: self.is_leaf(), node_hash: self.root_hash }
     }
 
-    fn get_root_prefix<L: Leaf>(
+    fn get_root_db_key<L: Leaf>(
         &self,
         key_context: &<L as HasStaticPrefix>::KeyContext,
-    ) -> DbKeyPrefix {
-        if self.is_leaf() {
+    ) -> DbKey {
+        let prefix: DbKeyPrefix = if self.is_leaf() {
             PatriciaPrefix::Leaf(L::get_static_prefix(key_context)).into()
         } else {
             PatriciaPrefix::InnerNode.into()
-        }
-    }
-
-    fn get_root_suffix(&self) -> Vec<u8> {
-        self.root_hash.0.to_bytes_be().to_vec()
+        };
+        let suffix = self.root_hash.0.to_bytes_be();
+        create_db_key(prefix, &suffix)
     }
 }
 /// Used for reading the roots in facts layout case.

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

@@ -8,10 +8,16 @@ use starknet_patricia_storage::db_object::{
     DBObject,
     EmptyDeserializationContext,
     EmptyKeyContext,
+    HasDynamicPrefix,
     HasStaticPrefix,
 };
 use starknet_patricia_storage::errors::{DeserializationError, SerializationError};
-use starknet_patricia_storage::storage_trait::{DbKeyPrefix, DbValue};
+use starknet_patricia_storage::storage_trait::{
+    create_db_key_no_separator,
+    DbKey,
+    DbKeyPrefix,
+    DbValue,
+};
 use starknet_types_core::felt::Felt;
 
 use crate::block_committer::input::StarknetStorageValue;
@@ -107,6 +113,11 @@ impl_leaf_for_wrappers!(
 
 impl DBObject for IndexLayoutContractState {
     type DeserializeContext = EmptyDeserializationContext;
+
+    fn get_db_key(&self, key_context: &Self::KeyContext, suffix: &[u8]) -> DbKey {
+        create_db_key_no_separator(self.get_prefix(key_context), suffix)
+    }
+
     fn serialize(&self) -> Result<DbValue, SerializationError> {
         serialize_felts(&[self.0.class_hash.0, self.0.storage_root_hash.0, self.0.nonce.0])
     }
@@ -133,6 +144,10 @@ impl DBObject for IndexLayoutContractState {
 impl DBObject for IndexLayoutCompiledClassHash {
     type DeserializeContext = EmptyDeserializationContext;
 
+    fn get_db_key(&self, key_context: &Self::KeyContext, suffix: &[u8]) -> DbKey {
+        create_db_key_no_separator(self.get_prefix(key_context), suffix)
+    }
+
     fn serialize(&self) -> Result<DbValue, SerializationError> {
         serialize_felts(&[self.0.0])
     }
@@ -150,6 +165,10 @@ impl DBObject for IndexLayoutCompiledClassHash {
 impl DBObject for IndexLayoutStarknetStorageValue {
     type DeserializeContext = EmptyDeserializationContext;
 
+    fn get_db_key(&self, key_context: &Self::KeyContext, suffix: &[u8]) -> DbKey {
+        create_db_key_no_separator(self.get_prefix(key_context), suffix)
+    }
+
     fn serialize(&self) -> Result<DbValue, SerializationError> {
         serialize_felts(&[self.0.0])
     }

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

@@ -17,10 +17,16 @@ use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::hash_functio
 use starknet_patricia_storage::db_object::{
     DBObject,
     EmptyDeserializationContext,
+    HasDynamicPrefix,
     HasStaticPrefix,
 };
 use starknet_patricia_storage::errors::{DeserializationError, SerializationResult};
-use starknet_patricia_storage::storage_trait::{DbKeyPrefix, DbValue};
+use starknet_patricia_storage::storage_trait::{
+    create_db_key_no_separator,
+    DbKey,
+    DbKeyPrefix,
+    DbValue,
+};
 use starknet_types_core::felt::Felt;
 
 use crate::hash_function::hash::TreeHashFunctionImpl;
@@ -58,6 +64,10 @@ where
 {
     type DeserializeContext = IndexNodeContext;
 
+    fn get_db_key(&self, key_context: &Self::KeyContext, suffix: &[u8]) -> DbKey {
+        create_db_key_no_separator(self.get_prefix(key_context), suffix)
+    }
+
     fn serialize(&self) -> SerializationResult<DbValue> {
         match &self.0.data {
             // Binary node - only serialize the hash.
@@ -178,14 +188,9 @@ impl<'a> SubTreeTrait<'a> for IndexLayoutSubTree<'a> {
         Self::NodeDeserializeContext { is_leaf: self.is_leaf() }
     }
 
-    fn get_root_prefix<L: Leaf>(
-        &self,
-        key_context: &<L as HasStaticPrefix>::KeyContext,
-    ) -> DbKeyPrefix {
-        L::get_static_prefix(key_context)
-    }
-
-    fn get_root_suffix(&self) -> Vec<u8> {
-        self.root_index.0.to_be_bytes().to_vec()
+    fn get_root_db_key<L: Leaf>(&self, key_context: &<L as HasStaticPrefix>::KeyContext) -> DbKey {
+        let prefix = L::get_static_prefix(key_context);
+        let suffix = self.root_index.0.to_be_bytes();
+        create_db_key_no_separator(prefix, &suffix)
     }
 }

crates/starknet_committer/src/db/trie_traversal.rs

@@ -26,7 +26,7 @@ use starknet_patricia::patricia_merkle_tree::traversal::{
 use starknet_patricia::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices};
 use starknet_patricia_storage::db_object::{DBObject, EmptyKeyContext, HasStaticPrefix};
 use starknet_patricia_storage::errors::StorageError;
-use starknet_patricia_storage::storage_trait::{create_db_key, DbKey, Storage};
+use starknet_patricia_storage::storage_trait::{DbKey, Storage};
 use tracing::warn;
 
 use crate::block_committer::input::{
@@ -245,12 +245,8 @@ pub async fn get_roots_from_storage<'a, L: Leaf, Layout: NodeLayout<'a, L>>(
     key_context: &<L as HasStaticPrefix>::KeyContext,
 ) -> TraversalResult<Vec<FilledNode<L, Layout::NodeData>>> {
     let mut subtrees_roots = vec![];
-    let db_keys: Vec<DbKey> = subtrees
-        .iter()
-        .map(|subtree| {
-            create_db_key(subtree.get_root_prefix::<L>(key_context), &subtree.get_root_suffix())
-        })
-        .collect();
+    let db_keys: Vec<DbKey> =
+        subtrees.iter().map(|subtree| subtree.get_root_db_key::<L>(key_context)).collect();
 
     let db_vals = storage.mget(&db_keys.iter().collect::<Vec<&DbKey>>()).await?;
     for ((subtree, optional_val), db_key) in subtrees.iter().zip(db_vals.iter()).zip(db_keys) {

crates/starknet_patricia/src/patricia_merkle_tree/traversal.rs

@@ -1,7 +1,7 @@
 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};
+use starknet_patricia_storage::storage_trait::{DbKey, PatriciaStorageError};
 use thiserror::Error;
 
 use crate::patricia_merkle_tree::node_data::inner_node::PathToBottom;
@@ -116,14 +116,8 @@ pub trait SubTreeTrait<'a>: Sized {
     /// [crate::patricia_merkle_tree::original_skeleton_tree::tree::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 full [DbKey] for the root node.
+    fn get_root_db_key<L: Leaf>(&self, key_context: &<L as HasStaticPrefix>::KeyContext) -> DbKey;
 
     /// Returns the `Self::NodeDeserializeContext` that's needed to deserialize the root node from a
     /// raw `DbValue`.

crates/starknet_patricia/src/patricia_merkle_tree/traversal_test.rs

@@ -2,7 +2,7 @@ use ethnum::U256;
 use pretty_assertions::assert_eq;
 use rstest::rstest;
 use starknet_patricia_storage::db_object::HasStaticPrefix;
-use starknet_patricia_storage::storage_trait::DbKeyPrefix;
+use starknet_patricia_storage::storage_trait::DbKey;
 
 use crate::patricia_merkle_tree::external_test_utils::{small_tree_index_to_full, TEST_PREFIX};
 use crate::patricia_merkle_tree::node_data::inner_node::{EdgePath, EdgePathLength, PathToBottom};
@@ -44,16 +44,12 @@ impl<'a> SubTreeTrait<'a> for TestSubTree<'a> {
 
     fn get_root_context(&self) -> Self::NodeDeserializeContext {}
 
-    fn get_root_prefix<L: Leaf>(
+    fn get_root_db_key<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(TEST_PREFIX.into())
-    }
-
-    fn get_root_suffix(&self) -> Vec<u8> {
-        vec![]
+    ) -> DbKey {
+        // Dummy key for testing purposes (we only need a key when interacting with storage).
+        DbKey(TEST_PREFIX.to_vec())
     }
 }
 

crates/starknet_patricia_storage/src/storage_trait.rs

@@ -205,11 +205,18 @@ impl Serialize for DbKey {
     }
 }
 
-/// Returns a `DbKey` from a prefix and a suffix.
+/// Returns a `DbKey` from a prefix and a suffix, with a ":" separator.
+/// This is used in the facts layout.
 pub fn create_db_key(prefix: DbKeyPrefix, suffix: &[u8]) -> DbKey {
     DbKey([prefix.to_bytes().to_vec(), b":".to_vec(), suffix.to_vec()].concat())
 }
 
+/// Returns a `DbKey` from a prefix and a suffix, without a separator.
+/// This is used in the index layout.
+pub fn create_db_key_no_separator(prefix: DbKeyPrefix, suffix: &[u8]) -> DbKey {
+    DbKey([prefix.to_bytes().to_vec(), suffix.to_vec()].concat())
+}
+
 /// Extracts the suffix from a `DbKey`. If the key doesn't match the prefix, None is returned.
 pub fn try_extract_suffix_from_db_key<'a>(
     key: &'a DbKey,