refactor: leafdata implementations

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

@@ -4,10 +4,11 @@ use crate::patricia_merkle_tree::filled_tree::node::FilledNode;
 use crate::patricia_merkle_tree::node_data::inner_node::{
     BinaryData, EdgeData, EdgePath, EdgePathLength, NodeData, PathToBottom,
 };
-use crate::patricia_merkle_tree::node_data::leaf::LeafDataImpl;
+use crate::patricia_merkle_tree::node_data::leaf::{LeafData, LeafDataImpl};
 use crate::storage::errors::{DeserializationError, SerializationError};
-use crate::storage::serde_trait::{Deserializable, Serializable};
-use crate::storage::storage_trait::{create_db_key, StorageKey, StoragePrefix, StorageValue};
+use crate::storage::serde_trait::{DBObject, Deserializable};
+use crate::storage::storage_trait::{StorageKey, StoragePrefix, StorageValue};
+
 use serde::{Deserialize, Serialize};
 
 // Const describe the size of the serialized node.
@@ -31,13 +32,18 @@ pub(crate) struct LeafCompiledClassToSerialize {
 type FilledNodeSerializationResult = Result<StorageValue, SerializationError>;
 type FilledNodeDeserializationResult = Result<FilledNode<LeafDataImpl>, DeserializationError>;
 
-impl FilledNode<LeafDataImpl> {
-    pub(crate) fn suffix(&self) -> [u8; SERIALIZE_HASH_BYTES] {
+#[allow(dead_code)]
+impl<L: LeafData> FilledNode<L> {
+    pub fn suffix(&self) -> [u8; SERIALIZE_HASH_BYTES] {
         self.hash.0.to_bytes_be()
     }
+
+    pub fn db_key(&self) -> StorageKey {
+        self.get_db_key(&self.suffix())
+    }
 }
 
-impl Serializable for FilledNode<LeafDataImpl> {
+impl<L: LeafData> DBObject for FilledNode<L> {
     /// This method serializes the filled node into a byte vector, where:
     /// - For binary nodes: Concatenates left and right hashes.
     /// - For edge nodes: Concatenates bottom hash, path, and path length.
@@ -75,23 +81,10 @@ impl Serializable for FilledNode<LeafDataImpl> {
         }
     }
 
-    /// Returns the db key of the filled node - [prefix + b":" + suffix].
-    fn db_key(&self) -> StorageKey {
-        let suffix = self.suffix();
-
+    fn get_prefix(&self) -> StoragePrefix {
         match &self.data {
-            NodeData::Binary(_) | NodeData::Edge(_) => {
-                create_db_key(StoragePrefix::InnerNode, &suffix)
-            }
-            NodeData::Leaf(LeafDataImpl::StorageValue(_)) => {
-                create_db_key(StoragePrefix::StorageLeaf, &suffix)
-            }
-            NodeData::Leaf(LeafDataImpl::CompiledClassHash(_)) => {
-                create_db_key(StoragePrefix::CompiledClassLeaf, &suffix)
-            }
-            NodeData::Leaf(LeafDataImpl::ContractState { .. }) => {
-                create_db_key(StoragePrefix::StateTreeLeaf, &suffix)
-            }
+            NodeData::Binary(_) | NodeData::Edge(_) => StoragePrefix::InnerNode,
+            NodeData::Leaf(leaf_data) => leaf_data.get_prefix(),
         }
     }
 }

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

@@ -4,8 +4,9 @@ use crate::felt::Felt;
 use crate::hash::hash_trait::HashOutput;
 use crate::patricia_merkle_tree::filled_tree::node::{ClassHash, Nonce};
 use crate::patricia_merkle_tree::types::NodeIndex;
+use crate::storage::serde_trait::DBObject;
 
-pub trait LeafData: Clone + Sync + Send {
+pub trait LeafData: Clone + Sync + Send + DBObject {
     /// Returns true if leaf is empty.
     fn is_empty(&self) -> bool;
 }

crates/committer/src/patricia_merkle_tree/node_data/leaf_serde.rs

@@ -1,13 +1,14 @@
 use crate::patricia_merkle_tree::node_data::leaf::{ContractState, LeafDataImpl};
-use crate::storage::storage_trait::StorageValue;
+use crate::storage::serde_trait::DBObject;
+use crate::storage::storage_trait::{StoragePrefix, StorageValue};
 
-impl LeafDataImpl {
+impl DBObject for LeafDataImpl {
     /// Serializes the leaf data into a byte vector.
     /// The serialization is done as follows:
     /// - For storage values: serializes the value into a 32-byte vector.
     /// - For compiled class hashes or state tree tuples: creates a  json string
     ///   describing the leaf and cast it into a byte vector.
-    pub fn serialize(&self) -> StorageValue {
+    fn serialize(&self) -> StorageValue {
         match &self {
             LeafDataImpl::StorageValue(value) => StorageValue(value.to_bytes_be().to_vec()),
 
@@ -33,4 +34,12 @@ impl LeafDataImpl {
             }
         }
     }
+
+    fn get_prefix(&self) -> StoragePrefix {
+        match self {
+            LeafDataImpl::StorageValue(_) => StoragePrefix::StorageLeaf,
+            LeafDataImpl::CompiledClassHash(_) => StoragePrefix::CompiledClassLeaf,
+            LeafDataImpl::ContractState { .. } => StoragePrefix::StateTreeLeaf,
+        }
+    }
 }

crates/committer/src/storage/serde_trait.rs

@@ -1,11 +1,24 @@
 use crate::storage::errors::DeserializationError;
-use crate::storage::storage_trait::{StorageKey, StorageValue};
+use crate::storage::storage_trait::{StorageKey, StoragePrefix, StorageValue};
 
-pub trait Serializable {
+pub trait DBObject {
     /// Serializes the given value.
     fn serialize(&self) -> StorageValue;
-    /// Returns the key used to store self in storage.
-    fn db_key(&self) -> StorageKey;
+
+    /// Returns the storage key prefix of the DB object.
+    fn get_prefix(&self) -> StoragePrefix;
+
+    /// Returns a `StorageKey` from a prefix and a suffix.
+    fn get_db_key(&self, suffix: &[u8]) -> StorageKey {
+        StorageKey(
+            [
+                self.get_prefix().to_bytes().to_vec(),
+                b":".to_vec(),
+                suffix.to_vec(),
+            ]
+            .concat(),
+        )
+    }
 }
 
 pub(crate) trait Deserializable: Sized {

crates/committer/src/storage/storage_trait.rs

@@ -34,7 +34,7 @@ pub trait Storage: From<HashMap<StorageKey, StorageValue>> {
     fn delete(&mut self, key: &StorageKey) -> Option<StorageValue>;
 }
 
-pub(crate) enum StoragePrefix {
+pub enum StoragePrefix {
     InnerNode,
     StorageLeaf,
     StateTreeLeaf,

crates/committer_cli/src/tests/python_tests.rs

@@ -16,7 +16,7 @@ use committer::patricia_merkle_tree::updated_skeleton_tree::hash_function::{
 };
 use committer::storage::errors::{DeserializationError, SerializationError};
 use committer::storage::map_storage::MapStorage;
-use committer::storage::serde_trait::Serializable;
+use committer::storage::serde_trait::DBObject;
 use committer::storage::storage_trait::{Storage, StorageKey, StorageValue};
 use std::fmt::Debug;
 use std::{collections::HashMap, io};
@@ -195,7 +195,7 @@ pub(crate) fn test_binary_serialize_test(binary_input: HashMap<String, u128>) ->
     };
 
     // Create a filled node with binary data and zero hash.
-    let filled_node = FilledNode {
+    let filled_node: FilledNode<LeafDataImpl> = FilledNode {
         data: NodeData::Binary(binary_data),
         hash: HashOutput(Felt::ZERO),
     };
@@ -386,50 +386,46 @@ pub(crate) fn test_node_db_key() -> String {
     // Generate keys for different node types.
     let hash = HashOutput(zero);
 
-    let binary_node_key = FilledNode {
+    let binary_node: FilledNode<LeafDataImpl> = FilledNode {
         data: NodeData::Binary(BinaryData {
             left_hash: hash,
             right_hash: hash,
         }),
         hash,
-    }
-    .db_key()
-    .0;
+    };
+    let binary_node_key = binary_node.db_key().0;
 
-    let edge_node_key = FilledNode {
+    let edge_node: FilledNode<LeafDataImpl> = FilledNode {
         data: NodeData::Edge(EdgeData {
             bottom_hash: hash,
             path_to_bottom: Default::default(),
         }),
         hash,
-    }
-    .db_key()
-    .0;
+    };
+
+    let edge_node_key = edge_node.db_key().0;
 
-    let storage_leaf_key = FilledNode {
+    let storage_leaf = FilledNode {
         data: NodeData::Leaf(LeafDataImpl::StorageValue(zero)),
         hash,
-    }
-    .db_key()
-    .0;
+    };
+    let storage_leaf_key = storage_leaf.db_key().0;
 
-    let state_tree_leaf_key = FilledNode {
+    let state_tree_leaf = FilledNode {
         data: NodeData::Leaf(LeafDataImpl::ContractState(ContractState {
             class_hash: ClassHash(zero),
             storage_root_hash: HashOutput(zero),
             nonce: Nonce(zero),
         })),
         hash,
-    }
-    .db_key()
-    .0;
+    };
+    let state_tree_leaf_key = state_tree_leaf.db_key().0;
 
-    let compiled_class_leaf_key = FilledNode {
+    let compiled_class_leaf = FilledNode {
         data: NodeData::Leaf(LeafDataImpl::CompiledClassHash(ClassHash(zero))),
         hash,
-    }
-    .db_key()
-    .0;
+    };
+    let compiled_class_leaf_key = compiled_class_leaf.db_key().0;
 
     // Store keys in a HashMap.
     let mut map: HashMap<String, Vec<u8>> = HashMap::new();
@@ -499,7 +495,7 @@ fn test_storage_node(data: HashMap<String, String>) -> Result<String, PythonTest
     let binary_data: HashMap<String, u128> = serde_json::from_str(binary_json)?;
 
     // Create a binary node from the parsed data.
-    let binary_rust = FilledNode {
+    let binary_rust: FilledNode<LeafDataImpl> = FilledNode {
         data: NodeData::Binary(BinaryData {
             left_hash: HashOutput(Felt::from(*get_or_key_not_found(&binary_data, "left")?)),
             right_hash: HashOutput(Felt::from(*get_or_key_not_found(&binary_data, "right")?)),
@@ -515,7 +511,7 @@ fn test_storage_node(data: HashMap<String, String>) -> Result<String, PythonTest
     let edge_data: HashMap<String, u128> = serde_json::from_str(edge_json)?;
 
     // Create an edge node from the parsed data.
-    let edge_rust = FilledNode {
+    let edge_rust: FilledNode<LeafDataImpl> = FilledNode {
         data: NodeData::Edge(EdgeData {
             bottom_hash: HashOutput(Felt::from(*get_or_key_not_found(&edge_data, "bottom")?)),
             path_to_bottom: PathToBottom {