feat: update sp1 programs to correctly compute the merkle tree

Author: MarcosNicolau

aggregation_mode/aggregation_programs/Cargo.toml

@@ -8,4 +8,3 @@ members = ["sp1", "risc0"]
 [patch.crates-io]
 # Adding RISC Zero keccak precompile support
 tiny-keccak = { git = "https://github.com/risc0/tiny-keccak", tag = "tiny-keccak/v2.0.2-risczero.0" }
-

aggregation_mode/aggregation_programs/sp1/src/chunk_aggregator_main.rs

@@ -1,26 +1,24 @@
 #![no_main]
 sp1_zkvm::entrypoint!(main);
 
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
 use sha2::{Digest, Sha256};
-use sp1_aggregation_program::{compute_merkle_root, Input};
+use sp1_aggregation_program::{ChunkAggregatorInput, SP1VkAndPubInputs};
 
 pub fn main() {
-    let input = sp1_zkvm::io::read::<Input>();
-
-    let mut proofs_commitment: Vec<[u8; 32]> = vec![];
+    let input = sp1_zkvm::io::read::<ChunkAggregatorInput>();
 
     // Verify the proofs.
     for proof in input.proofs_vk_and_pub_inputs.iter() {
         let vkey = proof.vk;
         let public_values = &proof.public_inputs;
         let public_values_digest = Sha256::digest(public_values);
 
-        proofs_commitment.push(proof.commitment());
-
         sp1_zkvm::lib::verify::verify_sp1_proof(&vkey, &public_values_digest.into());
     }
 
-    let merkle_root = compute_merkle_root(proofs_commitment);
+    let merkle_tree =
+        MerkleTree::<SP1VkAndPubInputs>::build(&input.proofs_vk_and_pub_inputs).unwrap();
 
-    sp1_zkvm::io::commit_slice(&merkle_root);
+    sp1_zkvm::io::commit_slice(&merkle_tree.root);
 }

aggregation_mode/aggregation_programs/sp1/src/lib.rs

@@ -46,29 +46,30 @@ impl IsMerkleTreeBackend for SP1VkAndPubInputs {
 }
 
 #[derive(Serialize, Deserialize)]
-pub struct Input {
-    pub proofs_vk_and_pub_inputs: Vec<SP1VkAndPubInputs>,
-}
+pub struct Hash32(pub [u8; 32]);
 
-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()
-}
+impl IsMerkleTreeBackend for Hash32 {
+    type Data = Hash32;
+    type Node = [u8; 32];
 
-/// Computes the merkle root for the given proofs
-pub fn compute_merkle_root(mut leaves_hash: Vec<[u8; 32]>) -> [u8; 32] {
-    while leaves_hash.len() > 1 {
-        leaves_hash = leaves_hash
-            .chunks(2)
-            .map(|chunk| match chunk {
-                [a, b] => combine_hashes(&a, &b),
-                [a] => combine_hashes(&a, &a),
-                _ => panic!("Unexpected chunk size in leaves"),
-            })
-            .collect()
+    fn hash_data(leaf: &Self::Data) -> Self::Node {
+        leaf.0
+    }
+
+    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()
     }
+}
 
-    leaves_hash[0]
+#[derive(Serialize, Deserialize)]
+pub struct ChunkAggregatorInput {
+    pub proofs_vk_and_pub_inputs: Vec<SP1VkAndPubInputs>,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct RootAggregatorInput {
+    pub proofs_and_leaves_commitment: Vec<(SP1VkAndPubInputs, Vec<[u8; 32]>)>,
 }

aggregation_mode/aggregation_programs/sp1/src/root_aggregator_main.rs

@@ -1,51 +1,47 @@
 #![no_main]
 sp1_zkvm::entrypoint!(main);
 
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
 use sha2::{Digest, Sha256};
-use sp1_aggregation_program::{compute_merkle_root, Input};
+use sp1_aggregation_program::{Hash32, RootAggregatorInput};
 
-pub const LEAVES_AGG_PROGRAM_VK_HASH: [u32; 8] = [0, 0, 0, 0, 0, 0, 0, 0];
+pub const CHUNK_AGGREGATOR_PROGRAM_VK_HASH: [u32; 8] = [0, 0, 0, 0, 0, 0, 0, 0];
 
 pub fn main() {
-    let input = sp1_zkvm::io::read::<Input>();
-    
-    println!("Number of proofs in input: {}", input.proofs_vk_and_pub_inputs.len());
+    let input = sp1_zkvm::io::read::<RootAggregatorInput>();
 
-    let mut proofs_hash: Vec<[u8; 32]> = vec![];
-    println!("Initial proofs_hash length: {}", proofs_hash.len());
+    let mut leaves = vec![];
 
     // Verify the proofs.
-    for (i, proof) in input.proofs_vk_and_pub_inputs.iter().enumerate() {
-        println!("Processing proof {}, public_inputs length: {}", i, proof.public_inputs.len());
-        
+    for (proof, leaves_commitment) in input.proofs_and_leaves_commitment {
         let vkey = proof.vk;
         let public_values_digest = Sha256::digest(&proof.public_inputs);
 
         // Ensure the aggregated chunk originates from the L1 aggregation program.
         // This validation step guarantees that the proof was genuinely verified
         // by this program. Without this check, a different program using the
         // same public inputs could bypass verification.
-        
-        // TODO: Add the assert here
-        //assert!(proof.vk == LEAVES_AGG_PROGRAM_VK_HASH);
+        assert!(proof.vk == CHUNK_AGGREGATOR_PROGRAM_VK_HASH);
 
         let merkle_root: [u8; 32] = proof
             .public_inputs
             .clone()
             .try_into()
             .expect("Public input to be the hash of the chunk tree");
 
-        proofs_hash.push(merkle_root);
-        println!("proofs_hash length after push: {}", proofs_hash.len());
+        // Reconstruct the merkle tree and verify that the roots match
+        let leaves_commitment: Vec<Hash32> =
+            leaves_commitment.into_iter().map(|el| Hash32(el)).collect();
+        let merkle_tree: MerkleTree<Hash32> = MerkleTree::build(&leaves_commitment).unwrap();
+        assert!(merkle_tree.root == merkle_root);
 
-        println!("vkey (debug): {:?}", vkey);
-        println!("public_values_digest (debug): {:?}", public_values_digest);
+        leaves.extend(leaves_commitment);
 
         sp1_zkvm::lib::verify::verify_sp1_proof(&vkey, &public_values_digest.into());
     }
 
-    println!("Final proofs_hash length before compute_merkle_root: {}", proofs_hash.len());
-    let merkle_root = compute_merkle_root(proofs_hash);
+    // Finally, compute the final merkle root with all the leaves
+    let merkle_tree: MerkleTree<Hash32> = MerkleTree::build(&leaves).unwrap();
 
-    sp1_zkvm::io::commit_slice(&merkle_root);
+    sp1_zkvm::io::commit_slice(&merkle_tree.root);
 }