refactor(aggregation_mode): migrate merkle tree to lambdaworks (#1911)

Co-authored-by: JuArce <[email protected]>

Author: MarcosNicolau

aggregation_mode/Cargo.lock

@@ -16,6 +16,7 @@ tokio = { version = "1", features = ["time"]}
 sha3 = "0.10.8"
 reqwest = { version = "0.12" }
 ciborium = "=0.2.2"
+lambdaworks-crypto = { git = "https://github.com/lambdaclass/lambdaworks.git", rev = "5f8f2cfcc8a1a22f77e8dff2d581f1166eefb80b", features = ["serde"]}
 # Necessary for the VerificationData type
 aligned-sdk = { path = "../batcher/aligned-sdk/" }
 # zkvms

aggregation_mode/Cargo.toml

@@ -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 = { git = "https://github.com/lambdaclass/lambdaworks.git", rev = "5f8f2cfcc8a1a22f77e8dff2d581f1166eefb80b", features = ["serde"]}
 
 [lib]
 path = "./src/lib.rs"

aggregation_mode/aggregation_programs/Cargo.lock

@@ -1,3 +1,4 @@
+use lambdaworks_crypto::merkle_tree::traits::IsMerkleTreeBackend;
 use serde::{Deserialize, Serialize};
 use tiny_keccak::{Hasher, Keccak};
 
@@ -21,6 +22,35 @@ impl Risc0ImageIdAndPubInputs {
     }
 }
 
+// Note: this MerkleTreeBackend is defined in three locations
+// - aggregation_mode/src/aggregators/mod.rs
+// - aggregation_mode/src/aggregators/risc0_aggregator.rs
+// - aggregation_mode/src/aggregators/sp1_aggregator.rs
+// All 3 implementations should match
+// The definition on aggregator/mod.rs supports taking proofs from both Risc0 and SP1,
+// Additionally, a version that takes the leaves as already hashed data is defined on:
+// - batcher/aligned-sdk/src/sdk/aggregation.rs
+// This one is used in the SDK since,
+// the user may not have access to the proofs that he didn't submit
+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/Cargo.toml

@@ -1,46 +1,11 @@
 #![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 +17,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/lib.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 = { git = "https://github.com/lambdaclass/lambdaworks.git", rev = "5f8f2cfcc8a1a22f77e8dff2d581f1166eefb80b", features = ["serde"]}
 
 [lib]
 path = "./src/lib.rs"

aggregation_mode/aggregation_programs/risc0/src/main.rs

@@ -1,3 +1,4 @@
+use lambdaworks_crypto::merkle_tree::traits::IsMerkleTreeBackend;
 use serde::{Deserialize, Serialize};
 use sha3::{Digest, Keccak256};
 
@@ -18,6 +19,32 @@ impl SP1VkAndPubInputs {
     }
 }
 
+// Note: this MerkleTreeBackend is defined in three locations
+// - aggregation_mode/src/aggregators/mod.rs
+// - aggregation_mode/src/aggregators/risc0_aggregator.rs and
+// - aggregation_mode/src/aggregators/sp1_aggregator.rs
+// All 3 implementations should match
+// The definition on aggregator/mod.rs supports taking proofs from both Risc0 and SP1,
+// Additionally, a version that takes the leaves as already hashed data is defined on:
+// - batcher/aligned-sdk/src/sdk/aggregation.rs
+// This one is used in the SDK since
+// the user may not have access to the proofs that they didn't submit
+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>,