feat: store merkle tree on program [skip-line-limit]

.github/workflows/ci.yml

@@ -680,6 +680,7 @@ jobs:
       - name: Test Template
         run: |
           cd templates/default
+          pnpm compile
           chmod 755 .enclave/generated/contracts tests
           chmod 644 .enclave/generated/contracts/ImageID.sol
           pnpm test:integration

crates/compute-provider/Cargo.toml

@@ -17,3 +17,4 @@ light-poseidon = "=0.2.0"
 ark-ff = "=0.4.2"
 ark-bn254 = "=0.4.0"
 rayon = "=1.10.0"
+zk-kit-imt = "0.0.7"

crates/compute-provider/src/compute_input.rs

@@ -5,7 +5,7 @@
 // or FITNESS FOR A PARTICULAR PURPOSE.
 
 use crate::ciphertext_output::ComputeResult;
-use crate::merkle_tree::MerkleTree;
+use crate::merkle_tree_builder::MerkleTreeBuilder;
 use sha3::{Digest, Keccak256};
 
 pub type FHEProcessor = fn(&FHEInputs) -> Vec<u8>;
@@ -34,12 +34,11 @@ impl ComputeInput {
             "Ciphertext hash mismatch"
         );
 
-        let merkle_root = MerkleTree {
-            leaf_hashes: self.leaf_hashes.clone(),
-        }
-        .build_tree()
-        .root()
-        .unwrap();
+        let merkle_root = MerkleTreeBuilder::new()
+            .with_leaf_hashes(self.leaf_hashes.clone())
+            .build_tree()
+            .root()
+            .unwrap();
 
         ComputeResult {
             ciphertext_hash: processed_hash,

crates/compute-provider/src/compute_manager.rs

@@ -6,7 +6,7 @@
 
 use crate::ciphertext_output::ComputeProvider;
 use crate::compute_input::{ComputeInput, FHEInputs};
-use crate::merkle_tree::MerkleTree;
+use crate::merkle_tree_builder::MerkleTreeBuilder;
 use crate::FHEProcessor;
 use rayon::prelude::*;
 use sha3::{Digest, Keccak256};
@@ -56,9 +56,10 @@ where
     }
 
     fn start_sequential(&mut self) -> (P::Output, Vec<u8>) {
-        let mut tree_handler = MerkleTree::new();
-        tree_handler.compute_leaf_hashes(&self.input.fhe_inputs.ciphertexts);
-        self.input.leaf_hashes = tree_handler.leaf_hashes.clone();
+        let mut tree_builder = MerkleTreeBuilder::new();
+
+        tree_builder.compute_leaf_hashes(&self.input.fhe_inputs.ciphertexts);
+        self.input.leaf_hashes = tree_builder.leaf_hashes.clone();
 
         // Compute the ciphertext
         let ciphertext = (self.processor)(&self.input.fhe_inputs);
@@ -83,9 +84,10 @@ where
         let tally_results: Vec<(P::Output, Vec<u8>, String)> = chunks
             .into_par_iter()
             .map(|chunk| {
-                let mut tree_handler = MerkleTree::new();
-                tree_handler.compute_leaf_hashes(&chunk);
-                let merkle_root = tree_handler.build_tree().root().unwrap();
+                let mut tree_builder = MerkleTreeBuilder::new();
+
+                tree_builder.compute_leaf_hashes(&chunk);
+                let merkle_root = tree_builder.build_tree().root().unwrap();
 
                 let fhe_inputs = FHEInputs {
                     ciphertexts: chunk.clone(),
@@ -98,7 +100,7 @@ where
                 let input = ComputeInput {
                     fhe_inputs,
                     ciphertext_hash,
-                    leaf_hashes: tree_handler.leaf_hashes.clone(),
+                    leaf_hashes: tree_builder.leaf_hashes.clone(),
                 };
 
                 (self.provider.prove(&input), ciphertext, merkle_root)

crates/compute-provider/src/lib.rs

@@ -7,7 +7,7 @@
 mod ciphertext_output;
 mod compute_input;
 mod compute_manager;
-mod merkle_tree;
+mod merkle_tree_builder;
 
 pub use ciphertext_output::*;
 pub use compute_input::*;

crates/compute-provider/src/merkle_tree.rs → crates/compute-provider/src/merkle_tree_builder.rs

@@ -6,24 +6,35 @@
 
 use ark_bn254::Fr;
 use ark_ff::{BigInt, BigInteger};
-use lean_imt::LeanIMT;
 use light_poseidon::{Poseidon, PoseidonHasher};
 use num_bigint::BigUint;
 use num_traits::Num;
 use sha3::{Digest, Keccak256};
 use std::str::FromStr;
+use zk_kit_imt::imt::IMT;
 
-pub struct MerkleTree {
+pub struct MerkleTreeBuilder {
     pub leaf_hashes: Vec<String>,
+    pub arity: usize,
+    pub zero_value: String,
+    pub depth: usize,
 }
 
-impl MerkleTree {
+impl MerkleTreeBuilder {
     pub fn new() -> Self {
         Self {
             leaf_hashes: Vec::new(),
+            arity: 2,
+            zero_value: "0".to_string(),
+            depth: 20,
         }
     }
 
+    pub fn with_leaf_hashes(mut self, leaf_hashes: Vec<String>) -> Self {
+        self.leaf_hashes = leaf_hashes;
+        self
+    }
+
     pub fn compute_leaf_hashes(&mut self, data: &[(Vec<u8>, u64)]) {
         for item in data {
             let hex_output = hex::encode(Keccak256::digest(&item.0));
@@ -60,9 +71,20 @@ impl MerkleTree {
         hex::encode(result_hash.to_bytes_be())
     }
 
-    pub fn build_tree(&self) -> LeanIMT {
-        let mut tree = LeanIMT::new(Self::poseidon_hash);
-        tree.insert_many(self.leaf_hashes.clone()).unwrap();
+    pub fn build_tree(&self) -> IMT {
+        let mut tree = IMT::new(
+            Self::poseidon_hash,
+            self.depth,
+            self.zero_value.clone(),
+            self.arity,
+            vec![],
+        )
+        .unwrap();
+
+        for leaf in &self.leaf_hashes {
+            tree.insert(leaf.clone()).unwrap();
+        }
+
         tree
     }
 }

crates/evm-helpers/src/events.rs

@@ -43,9 +43,6 @@ sol! {
         bytes plaintextOutput;
     }
 
-    #[derive(Debug)]
-    event InputPublished(uint256 indexed e3Id, bytes data, uint256 inputHash, uint256 index);
-
     #[derive(Debug)]
     event CiphertextOutputPublished(uint256 indexed e3Id, bytes ciphertextOutput);
 

crates/indexer/src/indexer.rs

@@ -14,7 +14,7 @@ use alloy::sol_types::SolEvent;
 use async_trait::async_trait;
 use e3_evm_helpers::{
     contracts::{EnclaveContract, EnclaveContractFactory, EnclaveRead, ReadOnly},
-    events::{CiphertextOutputPublished, E3Activated, InputPublished, PlaintextOutputPublished},
+    events::{CiphertextOutputPublished, E3Activated, PlaintextOutputPublished},
     listener::EventListener,
 };
 use eyre::eyre;
@@ -276,30 +276,6 @@ impl<S: DataStore> EnclaveIndexer<S> {
         Ok(())
     }
 
-    async fn register_input_published(&mut self) -> Result<()> {
-        let store = self.store.clone();
-        self.listener
-            .add_event_handler(move |e: InputPublished| {
-                let store = SharedStore::new(store.clone());
-                async move {
-                    println!(
-                        "InputPublished: e3_id={}, index={}, data=0x{}...",
-                        e.e3Id,
-                        e.index,
-                        hex::encode(&e.data[..8.min(e.data.len())])
-                    );
-                    let e3_id = u64_try_from(e.e3Id)?;
-
-                    let mut repo = E3Repository::new(store, e3_id);
-                    repo.insert_ciphertext_input(e.data.to_vec(), e.index.to::<u64>())
-                        .await?;
-                    Ok(())
-                }
-            })
-            .await;
-        Ok(())
-    }
-
     async fn register_ciphertext_output_published(&mut self) -> Result<()> {
         let store = self.store.clone();
         self.listener
@@ -349,7 +325,6 @@ impl<S: DataStore> EnclaveIndexer<S> {
 
     async fn setup_listeners(&mut self) -> Result<()> {
         self.register_e3_activated().await?;
-        self.register_input_published().await?;
         self.register_ciphertext_output_published().await?;
         self.register_plaintext_output_published().await?;
         Ok(())

crates/indexer/tests/integration.rs

@@ -61,56 +61,6 @@ async fn test_indexer() -> Result<()> {
         .watch()
         .await?;
 
-    // InputPublished
-    let data = "Random data that wont actually be a string".to_string();
-    contract
-        .emitInputPublished(
-            Uint::from(e3_id),
-            Bytes::from(data.clone().into_bytes()),
-            Uint::from(1111),
-            Uint::from(1),
-        )
-        .send()
-        .await?
-        .watch()
-        .await?;
-
-    contract
-        .emitInputPublished(
-            Uint::from(e3_id),
-            Bytes::from(data.clone().into_bytes()),
-            Uint::from(2222),
-            Uint::from(2),
-        )
-        .send()
-        .await?
-        .watch()
-        .await?;
-
-    contract
-        .emitInputPublished(
-            Uint::from(e3_id),
-            Bytes::from(data.clone().into_bytes()),
-            Uint::from(3333),
-            Uint::from(3),
-        )
-        .send()
-        .await?
-        .watch()
-        .await?;
-
-    sleep(Duration::from_millis(10)).await;
-
-    assert_eq!(indexer.get_e3(e3_id).await?.ciphertext_inputs.len(), 3);
-    assert_eq!(
-        indexer.get_e3(e3_id).await?.ciphertext_inputs,
-        vec![
-            (Bytes::from(data.clone().into_bytes()).to_vec(), 1),
-            (Bytes::from(data.clone().into_bytes()).to_vec(), 2),
-            (Bytes::from(data.clone().into_bytes()).to_vec(), 3),
-        ]
-    );
-
     let ciphertext_output = vec![9, 8, 7, 6, 5, 4, 3, 2, 1];
     contract
         .emitCiphertextOutputPublished(Uint::from(e3_id), Bytes::from(ciphertext_output.clone()))
@@ -125,8 +75,5 @@ async fn test_indexer() -> Result<()> {
 
     assert_eq!(e3.ciphertext_output, ciphertext_output);
 
-    let store = indexer.get_store();
-    let val = store.get::<u64>("input_count").await?.unwrap();
-    assert_eq!(val, 3);
     Ok(())
 }