refactor: redis as JMT backend (temporarily)

Author: distractedm1nd

src/common.rs

@@ -1,11 +1,11 @@
 use anyhow::{bail, Result};
 use borsh::{BorshDeserialize, BorshSerialize};
-use indexed_merkle_tree::{sha256_mod, Hash};
+use indexed_merkle_tree::Hash;
 use jmt::KeyHash;
 use serde::{Deserialize, Serialize};
 use std::fmt::Display;
 
-use crate::tree::Hasher;
+use crate::tree::{hash, Digest, Hasher};
 
 #[derive(Clone, BorshDeserialize, BorshSerialize, Serialize, Deserialize, Debug, PartialEq)]
 // An [`Operation`] represents a state transition in the system.
@@ -74,7 +74,20 @@ impl Hashchain {
         }
     }
 
-    pub fn push(&mut self, operation: Operation) -> Result<Hash> {
+    pub fn create_account(&mut self, value: String, source: AccountSource) -> Result<Digest> {
+        let operation = Operation::CreateAccount {
+            id: self.id.clone(),
+            value,
+            source,
+        };
+        self.push(operation)
+    }
+
+    pub fn get(&self, idx: usize) -> &HashchainEntry {
+        &self.entries[idx]
+    }
+
+    pub fn push(&mut self, operation: Operation) -> Result<Digest> {
         if let Operation::CreateAccount { .. } = operation {
             bail!("Cannot CreateAccount on an already existing hashchain");
         }
@@ -85,7 +98,7 @@ impl Hashchain {
         let previous_hash = self
             .entries
             .last()
-            .map_or(Hash::new([0u8; 32]), |entry| entry.hash);
+            .map_or(Digest::new([0u8; 32]), |entry| entry.hash);
 
         let entry = HashchainEntry::new(operation, previous_hash);
         self.entries.push(entry.clone());
@@ -94,15 +107,15 @@ impl Hashchain {
     }
 
     // TODO: Obviously, this needs to be authenticated by an existing key.
-    pub fn add(&mut self, value: String) -> Result<Hash> {
+    pub fn add(&mut self, value: String) -> Result<Digest> {
         let operation = Operation::Add {
             id: self.id.clone(),
             value,
         };
         self.push(operation)
     }
 
-    pub fn revoke(&mut self, value: String) -> Result<Hash> {
+    pub fn revoke(&mut self, value: String) -> Result<Digest> {
         let operation = Operation::Revoke {
             id: self.id.clone(),
             value,
@@ -123,19 +136,19 @@ impl Hashchain {
 // A [`HashchainEntry`] represents a single entry in an account's hashchain.
 // The value in the leaf of the corresponding account's node in the IMT is the hash of the last node in the hashchain.
 pub struct HashchainEntry {
-    pub hash: Hash,
-    pub previous_hash: Hash,
+    pub hash: Digest,
+    pub previous_hash: Digest,
     pub operation: Operation,
 }
 
 impl HashchainEntry {
-    pub fn new(operation: Operation, previous_hash: Hash) -> Self {
+    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());
             // TODO: replace with sha256 after JMT complete
-            sha256_mod(&data)
+            hash(&data)
         };
         Self {
             hash,

src/da/mod.rs

@@ -2,6 +2,7 @@ use crate::{
     circuits::{Bls12Proof, VerifyingKey},
     common::Operation,
     error::GeneralError,
+    tree::Digest,
     utils::SignedContent,
 };
 use anyhow::Result;
@@ -18,10 +19,10 @@ pub mod memory;
 #[derive(BorshSerialize, BorshDeserialize, Clone, Debug)]
 pub struct FinalizedEpoch {
     pub height: u64,
-    pub prev_commitment: Hash,
-    pub current_commitment: Hash,
-    pub proof: Bls12Proof,
-    pub verifying_key: VerifyingKey,
+    pub prev_commitment: Digest,
+    pub current_commitment: Digest,
+    // pub proof: Bls12Proof,
+    // pub verifying_key: VerifyingKey,
     pub signature: Option<String>,
 }
 

src/node_types/lightclient.rs

@@ -80,13 +80,13 @@ impl LightClient {
                                 let prev_commitment = &epoch_json.prev_commitment;
                                 let current_commitment = &epoch_json.current_commitment;
 
-                                let proof = match epoch_json.proof.clone().try_into() {
-                                    Ok(proof) => proof,
-                                    Err(e) => {
-                                        error!("failed to deserialize proof, skipping a blob at height {}: {:?}", i, e);
-                                        continue;
-                                    }
-                                };
+                                // let proof = match epoch_json.proof.clone().try_into() {
+                                //     Ok(proof) => proof,
+                                //     Err(e) => {
+                                //         error!("failed to deserialize proof, skipping a blob at height {}: {:?}", i, e);
+                                //         continue;
+                                //     }
+                                // };
 
                                 // TODO(@distractedm1nd): i don't know rust yet but this seems like non-idiomatic rust -
                                 // is there not a Trait that can satisfy these properties for us?

src/node_types/sequencer.rs

@@ -1,14 +1,9 @@
+use crate::tree::{hash, Digest, KeyDirectoryTree, Proof, RedisKDTree};
 use anyhow::{Context, Result};
 use async_trait::async_trait;
 use ed25519::Signature;
 use ed25519_dalek::{Signer, SigningKey};
-use jmt::{storage::{TreeReader, TreeWriter}, KeyHash};
-// use indexed_merkle_tree::{
-//     node::Node,
-//     tree::{IndexedMerkleTree, Proof},
-//     Hash,
-// };
-use crate::tree::{hash, Digest, KeyDirectoryTree, Proof};
+use jmt::KeyHash;
 use std::{self, str::FromStr, sync::Arc};
 use tokio::{
     sync::{
@@ -24,7 +19,6 @@ use crate::error::DataAvailabilityError;
 
 use crate::{
     cfg::Config,
-    circuits::BatchMerkleProofCircuit,
     common::{AccountSource, Hashchain, HashchainEntry, Operation},
     consts::{CHANNEL_BUFFER_SIZE, DA_RETRY_COUNT, DA_RETRY_INTERVAL},
     da::{DataAvailabilityLayer, FinalizedEpoch},
@@ -34,10 +28,7 @@ use crate::{
     webserver::{OperationInput, WebServer},
 };
 
-pub struct Sequencer<'a, S>
-where
-    S: 'a + TreeReader + TreeWriter,
-{
+pub struct Sequencer<'a> {
     pub db: Arc<dyn Database>,
     pub da: Arc<dyn DataAvailabilityLayer>,
     pub ws: WebServer,
@@ -52,17 +43,14 @@ where
     // [`pending_operations`] is a buffer for operations that have not yet been
     // posted to the DA layer.
     pending_operations: Arc<Mutex<Vec<Operation>>>,
-    tree: Arc<Mutex<KeyDirectoryTree<'a, S>>>,
+    tree: Arc<Mutex<RedisKDTree<'a>>>,
 
     epoch_buffer_tx: Arc<Sender<FinalizedEpoch>>,
     epoch_buffer_rx: Arc<Mutex<Receiver<FinalizedEpoch>>>,
 }
 
 #[async_trait]
-impl<'a, S> NodeType for Sequencer<'a, S>
-where
-    S: 'a + TreeReader + TreeWriter,
-{
+impl<'a> NodeType for Sequencer<'a> {
     async fn start(self: Arc<Self>) -> Result<()> {
         self.da.start().await.context("Failed to start DA layer")?;
 
@@ -95,13 +83,15 @@ impl Sequencer {
 
         let start_height = cfg.celestia_config.unwrap_or_default().start_height;
 
+        let tree = KeyDirectoryTree::new(&db.clone());
+
         Ok(Sequencer {
             db,
             da,
             ws: WebServer::new(ws),
             key,
             start_height,
-            tree: Arc::new(Mutex::new(IndexedMerkleTree::new_with_size(1024).unwrap())),
+            tree,
             pending_operations: Arc::new(Mutex::new(Vec::new())),
             epoch_buffer_tx: Arc::new(tx),
             epoch_buffer_rx: Arc::new(Mutex::new(rx)),
@@ -309,22 +299,22 @@ impl Sequencer {
             .context("Failed to set new epoch")?;
         // add the commitment for the operations ran since the last epoch
         self.db
-            .add_commitment(&epoch, &current_commitment)
+            .set_commitment(&epoch, &current_commitment)
             .context("Failed to add commitment for new epoch")?;
 
-        let batch_circuit =
-            BatchMerkleProofCircuit::new(&prev_commitment, &current_commitment, proofs)
-                .context("Failed to create BatchMerkleProofCircuit")?;
-        let (proof, verifying_key) = batch_circuit
-            .create_and_verify_snark()
-            .context("Failed to create and verify snark")?;
+        // let batch_circuit =
+        //     BatchMerkleProofCircuit::new(&prev_commitment, &current_commitment, proofs)
+        //         .context("Failed to create BatchMerkleProofCircuit")?;
+        // let (proof, verifying_key) = batch_circuit
+        //     .create_and_verify_snark()
+        //     .context("Failed to create and verify snark")?;
 
         let epoch_json = FinalizedEpoch {
             height: epoch,
             prev_commitment,
             current_commitment,
-            proof: proof.into(),
-            verifying_key: verifying_key.into(),
+            // proof: proof.into(),
+            // verifying_key: verifying_key.into(),
             signature: None,
         };
 
@@ -377,25 +367,23 @@ impl Sequencer {
                 let new_chain_entry = HashchainEntry::new(operation.clone(), previous_hash);
                 current_chain.push(new_chain_entry.clone());
 
-                let updated_node = Node::new_leaf(
-                    node.is_left_sibling(),
-                    hashed_id,
-                    new_chain_entry.hash,
-                    node.get_next(),
-                );
+                // let updated_node = Node::new_leaf(
+                //     node.is_left_sibling(),
+                //     hashed_id,
+                //     new_chain_entry.hash,
+                //     node.get_next(),
+                // );
 
                 debug!("updating hashchain for user id {}", id.clone());
-                self.tree.insert(KeyHash::with(hashed_id), )
+                let proof = self.tree.update(KeyHash::with(hashed_id), current_chain)?;
                 self.db
                     .update_hashchain(operation, &current_chain)
                     .context(format!(
                         "Failed to update hashchain for operation {:?}",
                         operation
                     ))?;
 
-                tree.update_node(index, updated_node)
-                    .map(Proof::Update)
-                    .context("Failed to update node in tree")
+                proof
             }
             Operation::CreateAccount { id, value, source } => {
                 // validation of account source
@@ -420,10 +408,11 @@ impl Sequencer {
                 }
 
                 debug!("creating new hashchain for user id {}", id.clone());
-                let new_chain = vec![HashchainEntry::new(operation.clone(), Node::HEAD)];
+                let chain = Hashchain::new(id.clone());
+                chain.create_account(value.into(), *source);
 
                 self.db
-                    .update_hashchain(operation, &new_chain)
+                    .update_hashchain(operation, &chain)
                     .context(format!(
                         "Failed to create hashchain for operation {:?}",
                         operation
@@ -432,11 +421,7 @@ impl Sequencer {
                 let mut tree = self.tree.lock().await;
                 let hashed_id = hash(id.as_bytes());
 
-                let mut node =
-                    Node::new_leaf(true, hashed_id, new_chain.first().unwrap().hash, Node::TAIL);
-                tree.insert_node(&mut node)
-                    .map(Proof::Insert)
-                    .context("Failed to insert node into tree")
+                tree.insert(KeyHash::with(hashed_id), chain)
             }
         }
     }
@@ -479,7 +464,7 @@ mod tests {
     }
 
     // Helper function to create a test Sequencer instance
-    async fn create_test_sequencer() -> Arc<Sequencer> {
+    async fn create_test_sequencer() -> Arc<Sequencer<'static>> {
         let (da_layer, _rx, _brx) = InMemoryDataAvailabilityLayer::new(1);
         let da_layer = Arc::new(da_layer);
         let db = Arc::new(setup_db());
@@ -578,7 +563,7 @@ mod tests {
         assert_ne!(prev_commitment, new_commitment);
 
         let hashchain = sequencer.db.get_hashchain(id.as_str());
-        let value = hashchain.unwrap().first().unwrap().operation.value();
+        let value = hashchain.unwrap().get(0).operation.value();
         assert_eq!(value, "test");
 
         teardown_db(&db);
@@ -696,7 +681,7 @@ mod tests {
 
         let hashchain = sequencer.db.get_hashchain("user@example.com").unwrap();
         assert_eq!(hashchain.len(), 2);
-        assert_eq!(hashchain[1].operation.value(), "new_value");
+        assert_eq!(hashchain.get(1).operation.value(), "new_value");
     }
 
     #[tokio::test]
@@ -724,7 +709,10 @@ mod tests {
 
         let hashchain = sequencer.db.get_hashchain("user@example.com").unwrap();
         assert_eq!(hashchain.len(), 2);
-        assert!(matches!(hashchain[1].operation, Operation::Revoke { .. }));
+        assert!(matches!(
+            hashchain.get(1).operation,
+            Operation::Revoke { .. }
+        ));
     }
 
     #[tokio::test]