feat: update risc0 programs to correctly compute the merkle tree

Author: MarcosNicolau

aggregation_mode/aggregation_programs/risc0/Cargo.toml

@@ -17,8 +17,8 @@ path = "./src/lib.rs"
 
 [[bin]]
 name = "risc0_chunk_aggregator_program"
-path = "./bin/chunk_aggregator.rs"
+path = "./src/chunk_aggregator_main.rs"
 
 [[bin]]
 name = "risc0_root_aggregator_program"
-path = "./bin/root_aggregator.rs"
+path = "./src/root_aggregator_main.rs"

aggregation_mode/aggregation_programs/risc0/bin/chunk_aggregator.rs

@@ -0,0 +1,22 @@
+#![no_main]
+
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
+use risc0_aggregation_program::{ChunkAggregatorInput, Risc0ImageIdAndPubInputs};
+use risc0_zkvm::guest::env;
+
+risc0_zkvm::guest::entry!(main);
+
+fn main() {
+    let input = env::read::<ChunkAggregatorInput>();
+
+    for proof in &input.proofs_image_id_and_pub_inputs {
+        env::verify(proof.image_id.clone(), &proof.public_inputs)
+            .expect("proof to be verified correctly");
+    }
+
+    let merkle_tree =
+        MerkleTree::<Risc0ImageIdAndPubInputs>::build(&input.proofs_image_id_and_pub_inputs)
+            .unwrap();
+
+    env::commit_slice(&merkle_tree.root);
+}

aggregation_mode/aggregation_programs/risc0/bin/root_aggregator.rs

@@ -52,32 +52,32 @@ impl IsMerkleTreeBackend for Risc0ImageIdAndPubInputs {
 }
 
 #[derive(Serialize, Deserialize)]
-pub struct Input {
-    pub proofs_image_id_and_pub_inputs: Vec<Risc0ImageIdAndPubInputs>,
-}
+pub struct Hash32(pub [u8; 32]);
 
-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);
+impl IsMerkleTreeBackend for Hash32 {
+    type Data = Hash32;
+    type Node = [u8; 32];
 
-    let mut hash = [0u8; 32];
-    hasher.finalize(&mut hash);
-    hash
-}
+    fn hash_data(leaf: &Self::Data) -> Self::Node {
+        leaf.0
+    }
 
-/// Computes the merkle root for the given proofs
-pub fn compute_merkle_root(mut leaves: Vec<[u8; 32]>) -> [u8; 32] {
-    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()
+    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
     }
+}
 
-    leaves[0]
+#[derive(Serialize, Deserialize)]
+pub struct ChunkAggregatorInput {
+    pub proofs_image_id_and_pub_inputs: Vec<Risc0ImageIdAndPubInputs>,
+}
+
+#[derive(Serialize, Deserialize)]
+pub struct RootAggregatorInput {
+    pub proofs_and_leaves_commitment: Vec<(Risc0ImageIdAndPubInputs, Vec<[u8; 32]>)>,
 }

aggregation_mode/aggregation_programs/risc0/src/chunk_aggregator_main.rs

@@ -0,0 +1,46 @@
+#![no_main]
+
+use lambdaworks_crypto::merkle_tree::merkle::MerkleTree;
+use risc0_aggregation_program::{Hash32, RootAggregatorInput};
+use risc0_zkvm::guest::env;
+
+risc0_zkvm::guest::entry!(main);
+
+pub const CHUNK_AGGREGATOR_PROGRAM_IMAGE_ID: [u8; 32] = [0u8; 32];
+
+fn main() {
+    let input = env::read::<RootAggregatorInput>();
+
+    let mut leaves: Vec<Hash32> = vec![];
+
+    for (proof, leaves_commitment) in input.proofs_and_leaves_commitment {
+        let image_id = proof.image_id;
+
+        // 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.
+        assert!(image_id == CHUNK_AGGREGATOR_PROGRAM_IMAGE_ID);
+
+        // Ensure the committed root matches the root of the provided leaves
+        let merkle_root: [u8; 32] = proof
+            .public_inputs
+            .clone()
+            .try_into()
+            .expect("Public input to be the chunk merkle root");
+
+        let leaves_commitment: Vec<Hash32> =
+            leaves_commitment.into_iter().map(|el| Hash32(el)).collect();
+        let merkle_tree = MerkleTree::<Hash32>::build(&leaves_commitment).unwrap();
+        assert!(merkle_root == merkle_tree.root);
+
+        leaves.extend(leaves_commitment);
+
+        // finally verify the proof
+        env::verify(image_id, &proof.public_inputs).expect("proof to be verified correctly");
+    }
+
+    let merkle_tree = MerkleTree::<Hash32>::build(&leaves).unwrap();
+
+    env::commit_slice(&merkle_tree.root);
+}