From 482f825e07a2d201362604eda05f9f9facafaaac Mon Sep 17 00:00:00 2001
From: AvivYossef-starkware
 <141143145+AvivYossef-starkware@users.noreply.github.com>
Date: Thu, 18 Apr 2024 17:57:06 +0300
Subject: [PATCH] build: impl db key to filled node (#49)

---
 .../src/patricia_merkle_tree/filled_node.rs   | 30 +++++++++++++++++--
 .../patricia_merkle_tree/serialized_node.rs   |  6 ++++
 crates/committer/src/storage/storage_trait.rs |  5 ++++
 3 files changed, 39 insertions(+), 2 deletions(-)

diff --git a/crates/committer/src/patricia_merkle_tree/filled_node.rs b/crates/committer/src/patricia_merkle_tree/filled_node.rs
index 4931f437..b1e5f280 100644
--- a/crates/committer/src/patricia_merkle_tree/filled_node.rs
+++ b/crates/committer/src/patricia_merkle_tree/filled_node.rs
@@ -3,12 +3,13 @@ use crate::patricia_merkle_tree::errors::FilledTreeError;
 use crate::patricia_merkle_tree::filled_tree::FilledTreeResult;
 use crate::patricia_merkle_tree::original_skeleton_tree::OriginalSkeletonTreeResult;
 use crate::patricia_merkle_tree::serialized_node::{
-    LeafCompiledClassToSerialize, SerializeNode, SERIALIZE_HASH_BYTES,
+    LeafCompiledClassToSerialize, SerializeNode, COMPLIED_CLASS_PREFIX, INNER_NODE_PREFIX,
+    SERIALIZE_HASH_BYTES, STATE_TREE_LEAF_PREFIX, STORAGE_LEAF_PREFIX,
 };
 use crate::patricia_merkle_tree::serialized_node::{BINARY_BYTES, EDGE_BYTES, EDGE_PATH_BYTES};
 use crate::patricia_merkle_tree::types::{EdgeData, LeafDataTrait};
 use crate::patricia_merkle_tree::types::{EdgePath, EdgePathLength, PathToBottom};
-use crate::storage::storage_trait::{StorageKey, StorageValue};
+use crate::storage::storage_trait::{create_db_key, StorageKey, StorageValue};
 use crate::types::Felt;
 
 // TODO(Nimrod, 1/6/2024): Swap to starknet-types-core types once implemented.
@@ -193,4 +194,29 @@ impl FilledNode<LeafData> {
             NodeData::Leaf(leaf_data) => leaf_data.serialize(),
         }
     }
+
+    /// Returns the suffix of the filled node, represented by its hash as a byte array.
+    #[allow(dead_code)]
+    pub(crate) fn suffix(&self) -> [u8; SERIALIZE_HASH_BYTES] {
+        self.hash.0.as_bytes()
+    }
+
+    /// Returns the db key of the filled node - [prefix + b":" + suffix].
+    #[allow(dead_code)]
+    pub(crate) fn db_key(&self) -> StorageKey {
+        let suffix = self.suffix();
+
+        match &self.data {
+            NodeData::Binary(_) | NodeData::Edge(_) => create_db_key(INNER_NODE_PREFIX, &suffix),
+            NodeData::Leaf(LeafData::StorageValue(_)) => {
+                create_db_key(STORAGE_LEAF_PREFIX, &suffix)
+            }
+            NodeData::Leaf(LeafData::CompiledClassHash(_)) => {
+                create_db_key(COMPLIED_CLASS_PREFIX, &suffix)
+            }
+            NodeData::Leaf(LeafData::StateTreeTuple { .. }) => {
+                create_db_key(STATE_TREE_LEAF_PREFIX, &suffix)
+            }
+        }
+    }
 }
diff --git a/crates/committer/src/patricia_merkle_tree/serialized_node.rs b/crates/committer/src/patricia_merkle_tree/serialized_node.rs
index 65fbd28f..03ab1b36 100644
--- a/crates/committer/src/patricia_merkle_tree/serialized_node.rs
+++ b/crates/committer/src/patricia_merkle_tree/serialized_node.rs
@@ -17,6 +17,12 @@ pub(crate) const STORAGE_LEAF_SIZE: usize = SERIALIZE_HASH_BYTES;
 // TODO(Aviv, 17/4/2024): add CompiledClassLeaf size.
 // TODO(Aviv, 17/4/2024): add StateTreeLeaf size.
 
+// Const describe the prefix of the serialized node.
+pub(crate) const STORAGE_LEAF_PREFIX: &[u8; 21] = b"starknet_storage_leaf";
+pub(crate) const STATE_TREE_LEAF_PREFIX: &[u8; 14] = b"contract_state";
+pub(crate) const COMPLIED_CLASS_PREFIX: &[u8; 19] = b"contract_class_leaf";
+pub(crate) const INNER_NODE_PREFIX: &[u8; 13] = b"patricia_node";
+
 /// Enum to describe the serialized node.
 #[allow(dead_code)]
 pub(crate) enum SerializeNode {
diff --git a/crates/committer/src/storage/storage_trait.rs b/crates/committer/src/storage/storage_trait.rs
index 807d62e9..f143b501 100644
--- a/crates/committer/src/storage/storage_trait.rs
+++ b/crates/committer/src/storage/storage_trait.rs
@@ -53,3 +53,8 @@ impl From<Felt> for StorageKey {
         StorageKey(value.to_bytes_be().to_vec())
     }
 }
+
+/// Returns a `StorageKey` from a prefix and a suffix.
+pub(crate) fn create_db_key(prefix: &[u8], suffix: &[u8]) -> StorageKey {
+    StorageKey([prefix.to_vec(), b":".to_vec(), suffix.to_vec()].concat())
+}