feat: migrate to lambdaworks merkle tree in proof aggregator

Author: MarcosNicolau

aggregation_mode/Cargo.lock

@@ -17,6 +17,7 @@ tokio = { version = "1", features = ["time"]}
 sha3 = "0.10.8"
 reqwest = { version = "0.12" }
 ciborium = "=0.2.2"
+lambdaworks-crypto = { version = "0.12.0", features = ["serde"]}
 # Necessary for the VerificationData type
 aligned-sdk = { path = "../batcher/aligned-sdk/" }
 # zkvms

aggregation_mode/Cargo.toml

@@ -3,9 +3,11 @@ pub mod sp1_aggregator;
 
 use std::fmt::Display;
 
+use lambdaworks_crypto::merkle_tree::traits::IsMerkleTreeBackend;
 use risc0_aggregator::{
     AlignedRisc0VerificationError, Risc0AggregationError, Risc0ProofReceiptAndImageId,
 };
+use sha3::{Digest, Keccak256};
 use sp1_aggregator::{
     AlignedSP1VerificationError, SP1AggregationError, SP1ProofWithPubValuesAndElf,
 };
@@ -114,6 +116,12 @@ pub enum AlignedProof {
     Risc0(Box<Risc0ProofReceiptAndImageId>),
 }
 
+impl Default for AlignedProof {
+    fn default() -> Self {
+        todo!()
+    }
+}
+
 impl AlignedProof {
     pub fn commitment(&self) -> [u8; 32] {
         match self {
@@ -123,6 +131,22 @@ impl AlignedProof {
     }
 }
 
+impl IsMerkleTreeBackend for AlignedProof {
+    type Data = AlignedProof;
+    type Node = [u8; 32];
+
+    fn hash_data(leaf: &Self::Data) -> Self::Node {
+        leaf.commitment()
+    }
+
+    fn hash_new_parent(child_1: &Self::Node, child_2: &Self::Node) -> Self::Node {
+        let mut hasher = Keccak256::new();
+        hasher.update(child_1);
+        hasher.update(child_2);
+        hasher.finalize().into()
+    }
+}
+
 #[derive(Debug)]
 pub enum AlignedVerificationError {
     Sp1(AlignedSP1VerificationError),

aggregation_mode/src/aggregators/mod.rs

@@ -1,29 +1,12 @@
 use crate::aggregators::AlignedProof;
-use sha3::{Digest, Keccak256};
-
-pub fn combine_hashes(hash_a: &[u8; 32], hash_b: &[u8; 32]) -> [u8; 32] {
-    let mut hasher = Keccak256::new();
-    hasher.update(hash_a);
-    hasher.update(hash_b);
-    hasher.finalize().into()
-}
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
 
 /// Returns (merkle_root, leaves)
-pub fn compute_proofs_merkle_root(proofs: &[AlignedProof]) -> ([u8; 32], Vec<[u8; 32]>) {
+pub fn compute_proofs_merkle_root(
+    proofs: &[AlignedProof],
+) -> Option<(MerkleTree<AlignedProof>, Vec<[u8; 32]>)> {
+    let merkle_tree: MerkleTree<AlignedProof> = MerkleTree::build(proofs)?;
     let leaves: Vec<[u8; 32]> = proofs.iter().map(|proof| proof.commitment()).collect();
 
-    let mut root = leaves.clone();
-
-    while root.len() > 1 {
-        root = root
-            .chunks(2)
-            .map(|chunk| match chunk {
-                [a, b] => combine_hashes(a, b),
-                [a] => combine_hashes(a, a),
-                _ => panic!("Unexpected chunk size in leaves"),
-            })
-            .collect()
-    }
-
-    (root[0], leaves)
+    Some((merkle_tree, leaves))
 }

aggregation_mode/src/backend/merkle_tree.rs

@@ -35,6 +35,7 @@ pub enum AggregatedProofSubmissionError {
     ReceiptError(PendingTransactionError),
     FetchingProofs(ProofsFetcherError),
     ZKVMAggregation(ProofAggregationError),
+    BuildingMerkleRoot,
     MerkleRootMisMatch,
 }
 
@@ -105,7 +106,9 @@ impl ProofAggregator {
         }
 
         info!("Proofs fetched, constructing merkle root...");
-        let (merkle_root, leaves) = compute_proofs_merkle_root(&proofs);
+        let (merkle_tree, leaves) = compute_proofs_merkle_root(&proofs)
+            .ok_or(AggregatedProofSubmissionError::BuildingMerkleRoot)?;
+        let merkle_root = merkle_tree.root;
         info!("Merkle root constructed: 0x{}", hex::encode(merkle_root));
 
         info!("Starting proof aggregation program...");