feat(starknet_patricia,starknet_committer): abstract storage prefix t…

…rait

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

@@ -5,7 +5,9 @@ use starknet_patricia::patricia_merkle_tree::node_data::errors::{LeafError, Leaf
 use starknet_patricia::patricia_merkle_tree::node_data::leaf::{Leaf, LeafModifications};
 use starknet_patricia::patricia_merkle_tree::types::NodeIndex;
 use starknet_patricia::patricia_merkle_tree::updated_skeleton_tree::tree::UpdatedSkeletonTreeImpl;
+use starknet_patricia::storage::storage_trait::StoragePrefix;
 
+use super::leaf_serde::CommitterLeafPrefix;
 use crate::block_committer::input::StarknetStorageValue;
 use crate::hash_function::hash::TreeHashFunctionImpl;
 use crate::patricia_merkle_tree::types::{ClassHash, CompiledClassHash, Nonce};
@@ -21,6 +23,10 @@ impl Leaf for StarknetStorageValue {
     type Input = Self;
     type Output = ();
 
+    fn storage_prefix() -> impl StoragePrefix {
+        CommitterLeafPrefix::StorageLeaf
+    }
+
     fn is_empty(&self) -> bool {
         self.0 == Felt::ZERO
     }
@@ -34,6 +40,10 @@ impl Leaf for CompiledClassHash {
     type Input = Self;
     type Output = ();
 
+    fn storage_prefix() -> impl StoragePrefix {
+        CommitterLeafPrefix::CompiledClassLeaf
+    }
+
     fn is_empty(&self) -> bool {
         self.0 == Felt::ZERO
     }
@@ -47,6 +57,10 @@ impl Leaf for ContractState {
     type Input = ContractStateInput;
     type Output = FilledTreeImpl<StarknetStorageValue>;
 
+    fn storage_prefix() -> impl StoragePrefix {
+        CommitterLeafPrefix::StateTreeLeaf
+    }
+
     fn is_empty(&self) -> bool {
         self.nonce.0 == Felt::ZERO
             && self.class_hash.0 == Felt::ZERO

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

@@ -3,23 +3,41 @@ use std::collections::HashMap;
 use serde_json::Value;
 use starknet_patricia::felt::Felt;
 use starknet_patricia::hash::hash_trait::HashOutput;
+use starknet_patricia::patricia_merkle_tree::node_data::leaf::Leaf;
 use starknet_patricia::patricia_merkle_tree::types::SubTreeHeight;
 use starknet_patricia::storage::db_object::{DBObject, Deserializable};
 use starknet_patricia::storage::errors::DeserializationError;
-use starknet_patricia::storage::storage_trait::{StarknetPrefix, StorageValue};
+use starknet_patricia::storage::storage_trait::{StoragePrefix, StorageValue};
 
 use crate::block_committer::input::StarknetStorageValue;
 use crate::patricia_merkle_tree::leaf::leaf_impl::ContractState;
 use crate::patricia_merkle_tree::types::{ClassHash, CompiledClassHash, Nonce};
 
+#[derive(Clone, Debug)]
+pub enum CommitterLeafPrefix {
+    StorageLeaf,
+    StateTreeLeaf,
+    CompiledClassLeaf,
+}
+
+impl StoragePrefix for CommitterLeafPrefix {
+    fn to_bytes(&self) -> &'static [u8] {
+        match self {
+            Self::StorageLeaf => b"starknet_storage_leaf",
+            Self::StateTreeLeaf => b"contract_state",
+            Self::CompiledClassLeaf => b"contract_class_leaf",
+        }
+    }
+}
+
 impl DBObject for StarknetStorageValue {
     /// Serializes the value into a 32-byte vector.
     fn serialize(&self) -> StorageValue {
         StorageValue(self.0.to_bytes_be().to_vec())
     }
 
-    fn get_prefix(&self) -> Vec<u8> {
-        StarknetPrefix::StorageLeaf.to_storage_prefix()
+    fn get_prefix(&self) -> impl StoragePrefix {
+        Self::storage_prefix()
     }
 }
 
@@ -30,8 +48,8 @@ impl DBObject for CompiledClassHash {
         StorageValue(json_string.into_bytes())
     }
 
-    fn get_prefix(&self) -> Vec<u8> {
-        StarknetPrefix::CompiledClassLeaf.to_storage_prefix()
+    fn get_prefix(&self) -> impl StoragePrefix {
+        Self::storage_prefix()
     }
 }
 
@@ -48,19 +66,15 @@ impl DBObject for ContractState {
         StorageValue(json_string.into_bytes())
     }
 
-    fn get_prefix(&self) -> Vec<u8> {
-        StarknetPrefix::StateTreeLeaf.to_storage_prefix()
+    fn get_prefix(&self) -> impl StoragePrefix {
+        Self::storage_prefix()
     }
 }
 
 impl Deserializable for StarknetStorageValue {
     fn deserialize(value: &StorageValue) -> Result<Self, DeserializationError> {
         Ok(Self(Felt::from_bytes_be_slice(&value.0)))
     }
-
-    fn prefix() -> Vec<u8> {
-        StarknetPrefix::StorageLeaf.to_storage_prefix()
-    }
 }
 
 impl Deserializable for CompiledClassHash {
@@ -72,10 +86,6 @@ impl Deserializable for CompiledClassHash {
             .ok_or(DeserializationError::NonExistingKey("compiled_class_hash".to_string()))?;
         Ok(Self::from_hex(hash_as_hex)?)
     }
-
-    fn prefix() -> Vec<u8> {
-        StarknetPrefix::CompiledClassLeaf.to_storage_prefix()
-    }
 }
 
 impl Deserializable for ContractState {
@@ -102,10 +112,6 @@ impl Deserializable for ContractState {
             class_hash: ClassHash::from_hex(&class_hash_as_hex)?,
         })
     }
-
-    fn prefix() -> Vec<u8> {
-        StarknetPrefix::StateTreeLeaf.to_storage_prefix()
-    }
 }
 
 fn get_key_from_map<'a>(map: &'a Value, key: &str) -> Result<&'a Value, DeserializationError> {

crates/starknet_patricia/src/patricia_merkle_tree/external_test_utils.rs

@@ -4,6 +4,7 @@ use ethnum::U256;
 use rand::Rng;
 use serde_json::json;
 
+use super::filled_tree::node_serde::PatriciaPrefix;
 use super::filled_tree::tree::{FilledTree, FilledTreeImpl};
 use super::node_data::inner_node::{EdgePathLength, PathToBottom};
 use super::node_data::leaf::{Leaf, LeafModifications, SkeletonLeaf};
@@ -17,7 +18,14 @@ use crate::felt::Felt;
 use crate::hash::hash_trait::HashOutput;
 use crate::patricia_merkle_tree::errors::TypesError;
 use crate::storage::map_storage::MapStorage;
-use crate::storage::storage_trait::{create_db_key, StarknetPrefix, StorageKey, StorageValue};
+use crate::storage::storage_trait::{create_db_key, StorageKey, StoragePrefix, StorageValue};
+
+pub struct MockLeafPrefix;
+impl StoragePrefix for MockLeafPrefix {
+    fn to_bytes(&self) -> &'static [u8] {
+        &[0]
+    }
+}
 
 impl TryFrom<&U256> for Felt {
     type Error = TypesError<U256>;
@@ -139,7 +147,7 @@ pub fn create_32_bytes_entry(simple_val: u128) -> [u8; 32] {
 }
 
 fn create_patricia_key(val: u128) -> StorageKey {
-    create_db_key(StarknetPrefix::InnerNode.to_storage_prefix(), &U256::from(val).to_be_bytes())
+    create_db_key(PatriciaPrefix::<MockLeafPrefix>::InnerNode, &U256::from(val).to_be_bytes())
 }
 
 fn create_binary_val(left: u128, right: u128) -> StorageValue {
@@ -201,7 +209,7 @@ pub fn create_root_edge_entry(
     let length = SubTreeHeight::ACTUAL_HEIGHT.0 - subtree_height.0;
     let new_root = old_root + u128::from(length);
     let key = create_db_key(
-        StarknetPrefix::InnerNode.to_storage_prefix(),
+        PatriciaPrefix::<MockLeafPrefix>::InnerNode,
         &Felt::from(new_root).to_bytes_be(),
     );
     let value = StorageValue(

crates/starknet_patricia/src/patricia_merkle_tree/filled_tree/node_serde.rs

@@ -14,7 +14,7 @@ use crate::patricia_merkle_tree::node_data::inner_node::{
 use crate::patricia_merkle_tree::node_data::leaf::Leaf;
 use crate::storage::db_object::DBObject;
 use crate::storage::errors::DeserializationError;
-use crate::storage::storage_trait::{StarknetPrefix, StorageKey, StorageValue};
+use crate::storage::storage_trait::{StorageKey, StoragePrefix, StorageValue};
 
 // Const describe the size of the serialized node.
 pub(crate) const SERIALIZE_HASH_BYTES: usize = 32;
@@ -25,6 +25,21 @@ pub(crate) const EDGE_BYTES: usize = SERIALIZE_HASH_BYTES + EDGE_PATH_BYTES + ED
 #[allow(dead_code)]
 pub(crate) const STORAGE_LEAF_SIZE: usize = SERIALIZE_HASH_BYTES;
 
+#[derive(Clone, Debug)]
+pub enum PatriciaPrefix<P: StoragePrefix> {
+    InnerNode,
+    Leaf(P),
+}
+
+impl<L: StoragePrefix> StoragePrefix for PatriciaPrefix<L> {
+    fn to_bytes(&self) -> &'static [u8] {
+        match self {
+            PatriciaPrefix::InnerNode => b"patricia_node",
+            PatriciaPrefix::Leaf(prefix) => prefix.to_bytes(),
+        }
+    }
+}
+
 /// Temporary struct to serialize the leaf CompiledClass.
 /// Required to comply to existing storage layout.
 #[derive(Serialize, Deserialize)]
@@ -75,12 +90,10 @@ impl<L: Leaf> DBObject for FilledNode<L> {
         }
     }
 
-    fn get_prefix(&self) -> Vec<u8> {
+    fn get_prefix(&self) -> impl StoragePrefix {
         match &self.data {
-            NodeData::Binary(_) | NodeData::Edge(_) => {
-                StarknetPrefix::InnerNode.to_storage_prefix()
-            }
-            NodeData::Leaf(leaf_data) => leaf_data.get_prefix(),
+            NodeData::Binary(_) | NodeData::Edge(_) => PatriciaPrefix::InnerNode,
+            NodeData::Leaf(_) => PatriciaPrefix::Leaf(L::storage_prefix()),
         }
     }
 }

crates/starknet_patricia/src/patricia_merkle_tree/internal_test_utils.rs

@@ -5,7 +5,7 @@ use rstest::{fixture, rstest};
 use crate::felt::Felt;
 use crate::generate_trie_config;
 use crate::hash::hash_trait::HashOutput;
-use crate::patricia_merkle_tree::external_test_utils::get_random_u256;
+use crate::patricia_merkle_tree::external_test_utils::{get_random_u256, MockLeafPrefix};
 use crate::patricia_merkle_tree::filled_tree::tree::FilledTreeImpl;
 use crate::patricia_merkle_tree::node_data::errors::{LeafError, LeafResult};
 use crate::patricia_merkle_tree::node_data::inner_node::{EdgePathLength, NodeData, PathToBottom};
@@ -20,7 +20,7 @@ use crate::patricia_merkle_tree::updated_skeleton_tree::hash_function::{
 use crate::patricia_merkle_tree::updated_skeleton_tree::node::UpdatedSkeletonNode;
 use crate::patricia_merkle_tree::updated_skeleton_tree::tree::UpdatedSkeletonTreeImpl;
 use crate::storage::db_object::{DBObject, Deserializable};
-use crate::storage::storage_trait::StorageValue;
+use crate::storage::storage_trait::{StoragePrefix, StorageValue};
 
 #[derive(Debug, PartialEq, Clone, Copy, Default, Eq)]
 pub struct MockLeaf(pub(crate) Felt);
@@ -30,8 +30,8 @@ impl DBObject for MockLeaf {
         StorageValue(self.0.to_bytes_be().to_vec())
     }
 
-    fn get_prefix(&self) -> Vec<u8> {
-        vec![0]
+    fn get_prefix(&self) -> impl StoragePrefix {
+        MockLeafPrefix
     }
 }
 
@@ -41,16 +41,16 @@ impl Deserializable for MockLeaf {
     ) -> Result<Self, crate::storage::errors::DeserializationError> {
         Ok(Self(Felt::from_bytes_be_slice(&value.0)))
     }
-
-    fn prefix() -> Vec<u8> {
-        vec![0]
-    }
 }
 
 impl Leaf for MockLeaf {
     type Input = Felt;
     type Output = String;
 
+    fn storage_prefix() -> impl StoragePrefix {
+        MockLeafPrefix
+    }
+
     fn is_empty(&self) -> bool {
         self.0 == Felt::ZERO
     }

crates/starknet_patricia/src/patricia_merkle_tree/node_data/leaf.rs

@@ -8,6 +8,7 @@ use crate::patricia_merkle_tree::original_skeleton_tree::errors::OriginalSkeleto
 use crate::patricia_merkle_tree::original_skeleton_tree::tree::OriginalSkeletonTreeResult;
 use crate::patricia_merkle_tree::types::NodeIndex;
 use crate::storage::db_object::{DBObject, Deserializable};
+use crate::storage::storage_trait::StoragePrefix;
 
 pub trait Leaf: Clone + Sync + Send + DBObject + Deserializable + Default + Debug + Eq {
     // TODO(Amos, 1/1/2025): When default values for associated types are stable - use them, and
@@ -16,6 +17,8 @@ pub trait Leaf: Clone + Sync + Send + DBObject + Deserializable + Default + Debu
     type Input: Send + Sync + 'static;
     type Output: Send + Debug + 'static;
 
+    fn storage_prefix() -> impl StoragePrefix;
+
     /// Returns true if leaf is empty.
     fn is_empty(&self) -> bool;
 

crates/starknet_patricia/src/patricia_merkle_tree/original_skeleton_tree/create_tree.rs

@@ -6,6 +6,7 @@ use tracing::warn;
 
 use crate::hash::hash_trait::HashOutput;
 use crate::patricia_merkle_tree::filled_tree::node::FilledNode;
+use crate::patricia_merkle_tree::filled_tree::node_serde::PatriciaPrefix;
 use crate::patricia_merkle_tree::node_data::inner_node::{
     BinaryData,
     EdgeData,
@@ -22,7 +23,7 @@ use crate::patricia_merkle_tree::original_skeleton_tree::tree::{
 use crate::patricia_merkle_tree::original_skeleton_tree::utils::split_leaves;
 use crate::patricia_merkle_tree::types::{NodeIndex, SortedLeafIndices, SubTreeHeight};
 use crate::storage::errors::StorageError;
-use crate::storage::storage_trait::{create_db_key, StarknetPrefix, Storage, StorageKey};
+use crate::storage::storage_trait::{create_db_key, Storage, StorageKey};
 
 #[cfg(test)]
 #[path = "create_tree_test.rs"]
@@ -222,9 +223,9 @@ impl<'a> OriginalSkeletonTreeImpl<'a> {
             .map(|subtree| {
                 create_db_key(
                     if subtree.is_leaf() {
-                        L::prefix()
+                        PatriciaPrefix::Leaf(L::storage_prefix())
                     } else {
-                        StarknetPrefix::InnerNode.to_storage_prefix()
+                        PatriciaPrefix::InnerNode
                     },
                     &subtree.root_hash.0.to_bytes_be(),
                 )

crates/starknet_patricia/src/storage/db_object.rs

@@ -1,25 +1,20 @@
 use crate::storage::errors::DeserializationError;
-use crate::storage::storage_trait::{StorageKey, StorageValue};
+use crate::storage::storage_trait::{create_db_key, StorageKey, StoragePrefix, StorageValue};
 
 pub trait DBObject {
     /// Serializes the given value.
     fn serialize(&self) -> StorageValue;
 
-    // TODO(Aviv, 17/07/2024): Define a trait `T` for storage prefix and return `impl T` here.
     /// Returns the storage key prefix of the DB object.
-    fn get_prefix(&self) -> Vec<u8>;
+    fn get_prefix(&self) -> impl StoragePrefix;
 
     /// Returns a `StorageKey` from a prefix and a suffix.
     fn get_db_key(&self, suffix: &[u8]) -> StorageKey {
-        StorageKey([self.get_prefix(), b":".to_vec(), suffix.to_vec()].concat())
+        create_db_key(self.get_prefix(), suffix)
     }
 }
 
 pub trait Deserializable: Sized {
     /// Deserializes the given value.
     fn deserialize(value: &StorageValue) -> Result<Self, DeserializationError>;
-
-    // TODO(Aviv, 17/07/2024): Define a trait `T` for storage prefix and return `impl T` here.
-    /// The prefix used to store in DB.
-    fn prefix() -> Vec<u8>;
 }

crates/starknet_patricia/src/storage/storage_trait.rs

@@ -31,30 +31,8 @@ pub trait Storage: From<HashMap<StorageKey, StorageValue>> {
     fn delete(&mut self, key: &StorageKey) -> Option<StorageValue>;
 }
 
-// TODO(Aviv, 17/07/2024): Split between Storage prefix representation (trait) and node
-// specific implementation (enum).
-#[derive(Clone, Debug)]
-pub enum StarknetPrefix {
-    InnerNode,
-    StorageLeaf,
-    StateTreeLeaf,
-    CompiledClassLeaf,
-}
-
-/// Describes a storage prefix as used in Aerospike DB.
-impl StarknetPrefix {
-    pub fn to_bytes(&self) -> &'static [u8] {
-        match self {
-            Self::InnerNode => b"patricia_node",
-            Self::StorageLeaf => b"starknet_storage_leaf",
-            Self::StateTreeLeaf => b"contract_state",
-            Self::CompiledClassLeaf => b"contract_class_leaf",
-        }
-    }
-
-    pub fn to_storage_prefix(&self) -> Vec<u8> {
-        self.to_bytes().to_vec()
-    }
+pub trait StoragePrefix {
+    fn to_bytes(&self) -> &'static [u8];
 }
 
 impl From<Felt> for StorageKey {
@@ -78,6 +56,6 @@ impl Serialize for StorageKey {
 }
 
 /// Returns a `StorageKey` from a prefix and a suffix.
-pub(crate) fn create_db_key(prefix: Vec<u8>, suffix: &[u8]) -> StorageKey {
-    StorageKey([prefix, b":".to_vec(), suffix.to_vec()].concat())
+pub fn create_db_key(prefix: impl StoragePrefix, suffix: &[u8]) -> StorageKey {
+    StorageKey([prefix.to_bytes().to_vec(), b":".to_vec(), suffix.to_vec()].concat())
 }