feat: refactor Digest and add basic tests

crates/common/src/hashchain.rs

@@ -9,7 +9,7 @@ use std::{
 
 use crate::{
     operation::{CreateAccountArgs, Operation, PublicKey, ServiceChallengeInput},
-    tree::{hash, Digest, Hasher},
+    tree::{Digest, Hasher},
 };
 
 #[derive(Clone, Serialize, Deserialize, Debug, PartialEq)]
@@ -280,12 +280,11 @@ pub struct HashchainEntry {
 
 impl HashchainEntry {
     pub fn new(operation: Operation, previous_hash: Digest) -> Self {
-        let hash = {
-            let mut data = Vec::new();
-            data.extend_from_slice(operation.to_string().as_bytes());
-            data.extend_from_slice(previous_hash.as_ref());
-            hash(&data)
-        };
+        let mut data = Vec::new();
+        data.extend_from_slice(operation.to_string().as_bytes());
+        data.extend_from_slice(previous_hash.as_ref());
+        let hash: Digest = data.as_slice().into();
+
         Self {
             hash,
             previous_hash,

crates/common/src/tree.rs

@@ -8,6 +8,7 @@ use jmt::{
 };
 use prism_errors::DatabaseError;
 use serde::{ser::SerializeTupleStruct, Deserialize, Serialize};
+use std::convert::{From, Into};
 use std::sync::Arc;
 
 use crate::{
@@ -70,17 +71,19 @@ impl SimpleHasher for Hasher {
     }
 }
 
-pub fn hash(data: &[u8]) -> Digest {
-    let mut hasher = Hasher::new();
-    hasher.update(data);
-    Digest(hasher.finalize())
-}
-
 #[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq, Copy)]
 pub struct Digest([u8; 32]);
 
-impl Digest {
-    pub fn to_bytes(&self) -> [u8; 32] {
+impl From<&[u8]> for Digest {
+    fn from(data: &[u8]) -> Self {
+        let mut hasher = Hasher::new();
+        hasher.update(data);
+        Self(hasher.finalize())
+    }
+}
+
+impl Into<[u8; 32]> for &Digest {
+    fn into(self) -> [u8; 32] {
         self.0
     }
 }
@@ -316,7 +319,7 @@ where
         match operation {
             Operation::AddKey(KeyOperationArgs { id, .. })
             | Operation::RevokeKey(KeyOperationArgs { id, .. }) => {
-                let hashed_id = hash(id.as_bytes());
+                let hashed_id: Digest = id.as_bytes().into();
                 let key_hash = KeyHash::with::<Hasher>(hashed_id);
 
                 let mut current_chain = self
@@ -337,7 +340,7 @@ where
                 service_id,
                 challenge,
             }) => {
-                let hashed_id = hash(id.as_bytes());
+                let hashed_id: Digest = id.as_bytes().into();
                 let key_hash = KeyHash::with::<Hasher>(hashed_id);
 
                 match &challenge {

crates/da/src/celestia.rs

@@ -16,8 +16,6 @@ use std::{
 };
 use tokio::{sync::broadcast, task::spawn};
 
-use tokio::{sync::broadcast, task::spawn};
-
 use bincode;
 
 impl TryFrom<&Blob> for FinalizedEpoch {

crates/prism/src/node_types/sequencer.rs

@@ -6,7 +6,7 @@ use jmt::KeyHash;
 use prism_common::{
     hashchain::Hashchain,
     tree::{
-        hash, Batch, Digest, Hasher, KeyDirectoryTree, NonMembershipProof, Proof, SnarkableTree,
+         Batch, Digest, Hasher, KeyDirectoryTree, NonMembershipProof, Proof, SnarkableTree,
     },
 };
 use std::{self, collections::VecDeque, sync::Arc};
@@ -359,7 +359,7 @@ impl Sequencer {
         id: &String,
     ) -> Result<Result<Hashchain, NonMembershipProof>> {
         let tree = self.tree.read().await;
-        let hashed_id = hash(id.as_bytes());
+        let hashed_id: Digest = id.as_bytes().into();
         let key_hash = KeyHash::with::<Hasher>(hashed_id);
 
         tree.get(key_hash)

crates/storage/Cargo.toml

@@ -28,3 +28,7 @@ prism-errors = { workspace = true }
 prism-common = { workspace = true }
 auto_impl = { workspace = true }
 rocksdb = { workspace = true }
+
+
+[dev-dependencies]
+tempfile = "3"

crates/storage/src/rocksdb.rs

@@ -36,9 +36,9 @@ impl Database for RocksDBConnection {
         epoch: &u64,
         commitment: &prism_common::tree::Digest,
     ) -> anyhow::Result<()> {
-        Ok(self.connection.put(
+        Ok(self.connection.put::<&[u8], [u8; 32]>(
             format!("commitments:epoch_{}", epoch).as_bytes(),
-            commitment.to_bytes(),
+            commitment.into(),
         )?)
     }
 
@@ -48,17 +48,19 @@ impl Database for RocksDBConnection {
             .get(b"app_state:epoch")?
             .ok_or_else(|| DatabaseError::NotFoundError("current epoch".to_string()))?;
 
-        Ok(u64::from_be_bytes(
-            res.try_into()
-                .map_err(|e| DatabaseError::ReadError(e.to_string()))?,
-        ))
+        Ok(u64::from_be_bytes(res.try_into().unwrap()))
     }
 
     fn set_epoch(&self, epoch: &u64) -> anyhow::Result<()> {
         Ok(self
             .connection
             .put(b"app_state:epoch", epoch.to_be_bytes())?)
     }
+
+    #[cfg(test)]
+    fn flush_database(&self) -> Result<()> {
+        todo!()
+    }
 }
 
 impl TreeWriter for RocksDBConnection {
@@ -84,3 +86,62 @@ impl TreeReader for RocksDBConnection {
         todo!()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use prism_common::tree::Digest;
+    use tempfile::TempDir;
+
+    #[test]
+    fn test_get_commitment() {
+        let temp_dir = TempDir::new().unwrap();
+        let db = RocksDBConnection::new(temp_dir.path().to_str().unwrap()).unwrap();
+
+        let epoch = 1;
+        let commitment = "some_commitment";
+        db.set_commitment(&epoch, &Digest::from(commitment.as_bytes()))
+            .unwrap();
+
+        let result = db.get_commitment(&epoch).unwrap();
+        assert_eq!(result, commitment.to_string());
+    }
+
+    #[test]
+    fn test_set_commitment() {
+        let temp_dir = TempDir::new().unwrap();
+        let db = RocksDBConnection::new(temp_dir.path().to_str().unwrap()).unwrap();
+
+        let epoch = 1;
+        let commitment = "some_commitment";
+        db.set_commitment(&epoch, &Digest::from(commitment.as_bytes()))
+            .unwrap();
+
+        let result = db.get_commitment(&epoch).unwrap();
+        assert_eq!(result, commitment.to_string());
+    }
+
+    #[test]
+    fn test_get_epoch() {
+        let temp_dir = TempDir::new().unwrap();
+        let db = RocksDBConnection::new(temp_dir.path().to_str().unwrap()).unwrap();
+
+        let epoch = 1;
+        db.set_epoch(&epoch).unwrap();
+
+        let result = db.get_epoch().unwrap();
+        assert_eq!(result, epoch);
+    }
+
+    #[test]
+    fn test_set_epoch() {
+        let temp_dir = TempDir::new().unwrap();
+        let db = RocksDBConnection::new(temp_dir.path().to_str().unwrap()).unwrap();
+
+        let epoch = 1;
+        db.set_epoch(&epoch).unwrap();
+
+        let result = db.get_epoch().unwrap();
+        assert_eq!(result, epoch);
+    }
+}