feat: migrate to lambdaworks merkle tree in zkvms

Author: MarcosNicolau

aggregation_mode/aggregation_programs/Cargo.lock

@@ -10,6 +10,7 @@ serde_json = "1.0.117"
 # Unstable feature is necessary for tiny-keccak patch, see: https://dev.risczero.com/api/zkvm/precompiles#stability
 risc0-zkvm = { git = "https://github.com/risc0/risc0", tag="v2.0.0", default-features = false, features = ["unstable", "std"] }
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
+lambdaworks-crypto = { version = "0.12.0", features = ["serde"]}
 
 [lib]
 path = "./src/lib.rs"

aggregation_mode/aggregation_programs/risc0/Cargo.toml

@@ -1,7 +1,8 @@
+use lambdaworks_crypto::merkle_tree::traits::IsMerkleTreeBackend;
 use serde::{Deserialize, Serialize};
 use tiny_keccak::{Hasher, Keccak};
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Default)]
 pub struct Risc0ImageIdAndPubInputs {
     pub image_id: [u8; 32],
     pub public_inputs: Vec<u8>,
@@ -21,6 +22,25 @@ impl Risc0ImageIdAndPubInputs {
     }
 }
 
+impl IsMerkleTreeBackend for Risc0ImageIdAndPubInputs {
+    type Data = Risc0ImageIdAndPubInputs;
+    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 = Keccak::v256();
+        hasher.update(child_1);
+        hasher.update(child_2);
+
+        let mut hash = [0u8; 32];
+        hasher.finalize(&mut hash);
+        hash
+    }
+}
+
 #[derive(Serialize, Deserialize)]
 pub struct Input {
     pub proofs_image_id_and_pub_inputs: Vec<Risc0ImageIdAndPubInputs>,

aggregation_mode/aggregation_programs/risc0/src/lib.rs

@@ -1,46 +1,12 @@
 #![no_main]
 
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
 use risc0_aggregation_program::{Input, Risc0ImageIdAndPubInputs};
 use risc0_zkvm::guest::env;
 use tiny_keccak::{Hasher, Keccak};
 
 risc0_zkvm::guest::entry!(main);
 
-fn combine_hashes(hash_a: &[u8; 32], hash_b: &[u8; 32]) -> [u8; 32] {
-    let mut hasher = Keccak::v256();
-    hasher.update(hash_a);
-    hasher.update(hash_b);
-
-    let mut hash = [0u8; 32];
-    hasher.finalize(&mut hash);
-    hash
-}
-
-/// Computes the merkle root for the given proofs
-fn compute_merkle_root(proofs: &[Risc0ImageIdAndPubInputs]) -> [u8; 32] {
-    let mut leaves: Vec<[u8; 32]> = proofs
-        .chunks(2)
-        .map(|chunk| match chunk {
-            [a, b] => combine_hashes(&a.commitment(), &b.commitment()),
-            [a] => combine_hashes(&a.commitment(), &a.commitment()),
-            _ => panic!("Unexpected chunk leaves"),
-        })
-        .collect();
-
-    while leaves.len() > 1 {
-        leaves = leaves
-            .chunks(2)
-            .map(|chunk| match chunk {
-                [a, b] => combine_hashes(&a, &b),
-                [a] => combine_hashes(&a, &a),
-                _ => panic!("Unexpected chunk size in leaves"),
-            })
-            .collect()
-    }
-
-    leaves[0]
-}
-
 fn main() {
     let input = env::read::<Input>();
 
@@ -52,7 +18,8 @@ fn main() {
         env::verify(image_id.clone(), &public_inputs).expect("proof to be verified correctly");
     }
 
-    let merkle_root = compute_merkle_root(&input.proofs_image_id_and_pub_inputs);
+    let merkle_tree: MerkleTree<Risc0ImageIdAndPubInputs> =
+        MerkleTree::build(&input.proofs_image_id_and_pub_inputs).unwrap();
 
-    env::commit_slice(&merkle_root);
+    env::commit_slice(&merkle_tree.root);
 }

aggregation_mode/aggregation_programs/risc0/src/main.rs

@@ -10,6 +10,7 @@ sha2 = { git = "https://github.com/sp1-patches/RustCrypto-hashes", tag = "sha2-v
 sha3 = { git = "https://github.com/sp1-patches/RustCrypto-hashes", tag = "sha3-v0.10.8-patch-v1" }
 serde = { version = "1.0.203", features = ["derive"] }
 serde_json = "1.0.117"
+lambdaworks-crypto = { version = "0.12.0", features = ["serde"]}
 
 [lib]
 path = "./src/lib.rs"

aggregation_mode/aggregation_programs/sp1/Cargo.toml

@@ -1,7 +1,8 @@
+use lambdaworks_crypto::merkle_tree::traits::IsMerkleTreeBackend;
 use serde::{Deserialize, Serialize};
 use sha3::{Digest, Keccak256};
 
-#[derive(Serialize, Deserialize)]
+#[derive(Serialize, Deserialize, Default)]
 pub struct SP1VkAndPubInputs {
     pub vk: [u32; 8],
     pub public_inputs: Vec<u8>,
@@ -18,6 +19,22 @@ impl SP1VkAndPubInputs {
     }
 }
 
+impl IsMerkleTreeBackend for SP1VkAndPubInputs {
+    type Data = SP1VkAndPubInputs;
+    type Node = [u8; 32];
+
+    fn hash_data(leaf: &Self::Data) -> Self::Node {
+        leaf.hash()
+    }
+
+    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(Serialize, Deserialize)]
 pub struct Input {
     pub proofs_vk_and_pub_inputs: Vec<SP1VkAndPubInputs>,

aggregation_mode/aggregation_programs/sp1/src/lib.rs

@@ -1,42 +1,11 @@
 #![no_main]
 sp1_zkvm::entrypoint!(main);
 
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
 use sha2::{Digest, Sha256};
 use sha3::Keccak256;
 use sp1_aggregation_program::{Input, SP1VkAndPubInputs};
 
-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()
-}
-
-/// Computes the merkle root for the given proofs using the vk
-fn compute_merkle_root(proofs: &[SP1VkAndPubInputs]) -> [u8; 32] {
-    let mut leaves: Vec<[u8; 32]> = proofs
-        .chunks(2)
-        .map(|chunk| match chunk {
-            [a, b] => combine_hashes(&a.hash(), &b.hash()),
-            [a] => combine_hashes(&a.hash(), &a.hash()),
-            _ => panic!("Unexpected chunk leaves"),
-        })
-        .collect();
-
-    while leaves.len() > 1 {
-        leaves = leaves
-            .chunks(2)
-            .map(|chunk| match chunk {
-                [a, b] => combine_hashes(&a, &b),
-                [a] => combine_hashes(&a, &a),
-                _ => panic!("Unexpected chunk size in leaves"),
-            })
-            .collect()
-    }
-
-    leaves[0]
-}
-
 pub fn main() {
     let input = sp1_zkvm::io::read::<Input>();
 
@@ -48,7 +17,8 @@ pub fn main() {
         sp1_zkvm::lib::verify::verify_sp1_proof(&vkey, &public_values_digest.into());
     }
 
-    let merkle_root = compute_merkle_root(&input.proofs_vk_and_pub_inputs);
+    let merkle_tree: MerkleTree<SP1VkAndPubInputs> =
+        MerkleTree::build(&input.proofs_vk_and_pub_inputs).unwrap();
 
-    sp1_zkvm::io::commit_slice(&merkle_root);
+    sp1_zkvm::io::commit_slice(&merkle_tree.root);
 }