feat: split aggregation in two programs

Author: MarcosNicolau

aggregation_mode/aggregation_programs/risc0/Cargo.toml

@@ -15,5 +15,9 @@ tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 path = "./src/lib.rs"
 
 [[bin]]
-name = "risc0_aggregator_program"
-path = "./src/main.rs"
+name = "risc0_chunk_aggregator_program"
+path = "./bin/chunk_aggregator.rs"
+
+[[bin]]
+name = "risc0_root_aggregator_program"
+path = "./bin/root_aggregator.rs"

aggregation_mode/aggregation_programs/risc0/bin/chunk_aggregator.rs

@@ -0,0 +1,22 @@
+#![no_main]
+
+use risc0_aggregation_program::{compute_merkle_root, Input, Risc0ImageIdAndPubInputs};
+use risc0_zkvm::guest::env;
+
+risc0_zkvm::guest::entry!(main);
+
+fn main() {
+    let input = env::read::<Input>();
+
+    let mut proofs_hash: Vec<[u8; 32]> = vec![];
+
+    for proof in &input.proofs_image_id_and_pub_inputs {
+        proofs_hash.push(proof.commitment());
+        env::verify(proof.image_id.clone(), &proof.public_inputs)
+            .expect("proof to be verified correctly");
+    }
+
+    let merkle_root = compute_merkle_root(proofs_hash);
+
+    env::commit_slice(&merkle_root);
+}

aggregation_mode/aggregation_programs/risc0/bin/root_aggregator.rs

@@ -0,0 +1,29 @@
+#![no_main]
+
+use risc0_aggregation_program::{compute_merkle_root, Input, Risc0ImageIdAndPubInputs};
+use risc0_zkvm::guest::env;
+
+risc0_zkvm::guest::entry!(main);
+
+fn main() {
+    let input = env::read::<Input>();
+
+    let mut chunks_merkle_root: Vec<[u8; 32]> = vec![];
+
+    for Risc0ImageIdAndPubInputs {
+        image_id,
+        public_inputs,
+    } in &input.proofs_image_id_and_pub_inputs
+    {
+        let merkle_root: [u8; 32] = public_inputs
+            .clone()
+            .try_into()
+            .expect("Public input to be the chunk merkle root");
+        chunks_merkle_root.push(merkle_root);
+        env::verify(image_id.clone(), &public_inputs).expect("proof to be verified correctly");
+    }
+
+    let merkle_root = compute_merkle_root(chunks_merkle_root);
+
+    env::commit_slice(&merkle_root);
+}

aggregation_mode/aggregation_programs/risc0/src/lib.rs

@@ -8,27 +8,46 @@ pub struct Risc0ImageIdAndPubInputs {
 }
 
 impl Risc0ImageIdAndPubInputs {
-    pub fn commitment(&self, is_aggregated_chunk: bool) -> [u8; 32] {
-        // If the proof is the aggregation of a chunk
-        // the commitment is simply the public input (the merkle root of the chunk)
-        if is_aggregated_chunk {
-            self.public_inputs.clone().try_into().unwrap()
-        } else {
-            let mut hasher = Keccak::v256();
-            for &word in &self.image_id {
-                hasher.update(&word.to_be_bytes());
-            }
-            hasher.update(&self.public_inputs);
-
-            let mut hash = [0u8; 32];
-            hasher.finalize(&mut hash);
-            hash
+    pub fn commitment(&self) -> [u8; 32] {
+        let mut hasher = Keccak::v256();
+        for &word in &self.image_id {
+            hasher.update(&word.to_be_bytes());
         }
+        hasher.update(&self.public_inputs);
+
+        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>,
-    pub is_aggregated_chunk: bool,
+}
+
+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
+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()
+    }
+
+    leaves[0]
 }

aggregation_mode/aggregation_programs/risc0/src/main.rs

@@ -15,5 +15,9 @@ serde_json = "1.0.117"
 path = "./src/lib.rs"
 
 [[bin]]
-name = "sp1_aggregator_program"
-path = "./src/main.rs"
+name = "sp1_chunk_aggregator_program"
+path = "./bin/chunk_aggregator.rs"
+
+[[bin]]
+name = "sp1_root_aggregator_program"
+path = "./bin/root_aggregator.rs"

aggregation_mode/aggregation_programs/sp1/Cargo.toml

@@ -0,0 +1,26 @@
+#![no_main]
+sp1_zkvm::entrypoint!(main);
+
+use sha2::{Digest, Sha256};
+use sp1_aggregation_program::{compute_merkle_root, Input};
+
+pub fn main() {
+    let input = sp1_zkvm::io::read::<Input>();
+
+    let mut proofs_commitment: Vec<[u8; 32]> = vec![];
+
+    // 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);
+
+    sp1_zkvm::io::commit_slice(&merkle_root);
+}

aggregation_mode/aggregation_programs/sp1/bin/chunk_aggregator.rs

@@ -0,0 +1,38 @@
+#![no_main]
+sp1_zkvm::entrypoint!(main);
+
+use sha2::{Digest, Sha256};
+use sp1_aggregation_program::{compute_merkle_root, Input};
+
+pub const LEAVES_AGG_PROGRAM_VK_HASH: [u32; 8] = [0, 0, 0, 0, 0, 0, 0, 0];
+
+pub fn main() {
+    let input = sp1_zkvm::io::read::<Input>();
+
+    let mut proofs_hash: Vec<[u8; 32]> = vec![];
+
+    // Verify the proofs.
+    for proof in input.proofs_vk_and_pub_inputs.iter() {
+        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.
+        assert!(proof.vk == LEAVES_AGG_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);
+
+        sp1_zkvm::lib::verify::verify_sp1_proof(&vkey, &public_values_digest.into());
+    }
+
+    let merkle_root = compute_merkle_root(proofs_hash);
+
+    sp1_zkvm::io::commit_slice(&merkle_root);
+}

aggregation_mode/aggregation_programs/sp1/bin/root_aggregator.rs

@@ -8,24 +8,40 @@ pub struct SP1VkAndPubInputs {
 }
 
 impl SP1VkAndPubInputs {
-    pub fn commitment(&self, is_aggregated_chunk: bool) -> [u8; 32] {
-        // If the proof is the aggregation of a chunk
-        // the commitment is simply the public input (the merkle root of the chunk)
-        if is_aggregated_chunk {
-            self.public_inputs.clone().try_into().unwrap()
-        } else {
-            let mut hasher = Keccak256::new();
-            for &word in &self.vk {
-                hasher.update(word.to_be_bytes());
-            }
-            hasher.update(&self.public_inputs);
-            hasher.finalize().into()
+    pub fn commitment(&self) -> [u8; 32] {
+        let mut hasher = Keccak256::new();
+        for &word in &self.vk {
+            hasher.update(word.to_be_bytes());
         }
+        hasher.update(&self.public_inputs);
+        hasher.finalize().into()
     }
 }
 
 #[derive(Serialize, Deserialize)]
 pub struct Input {
     pub proofs_vk_and_pub_inputs: Vec<SP1VkAndPubInputs>,
-    pub is_aggregated_chunk: bool,
+}
+
+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
+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()
+    }
+
+    leaves_hash[0]
 }