fix: aggregated proofs commitment for the second layer

Author: MarcosNicolau

aggregation_mode/aggregation_programs/risc0/src/lib.rs

@@ -8,20 +8,25 @@ pub struct Risc0ImageIdAndPubInputs {
 }
 
 impl Risc0ImageIdAndPubInputs {
-    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);
+    pub fn commitment(&self, is_aggregated_chunk: bool) -> [u8; 32] {
+        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
+            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,
 }

aggregation_mode/aggregation_programs/risc0/src/main.rs

@@ -17,12 +17,18 @@ fn combine_hashes(hash_a: &[u8; 32], hash_b: &[u8; 32]) -> [u8; 32] {
 }
 
 /// Computes the merkle root for the given proofs
-fn compute_merkle_root(proofs: &[Risc0ImageIdAndPubInputs]) -> [u8; 32] {
+fn compute_merkle_root(proofs: &[Risc0ImageIdAndPubInputs], is_aggregated_chunk: bool) -> [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()),
+            [a, b] => combine_hashes(
+                &a.commitment(is_aggregated_chunk),
+                &b.commitment(is_aggregated_chunk),
+            ),
+            [a] => combine_hashes(
+                &a.commitment(is_aggregated_chunk),
+                &a.commitment(is_aggregated_chunk),
+            ),
             _ => panic!("Unexpected chunk leaves"),
         })
         .collect();
@@ -52,7 +58,10 @@ 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_root = compute_merkle_root(
+        &input.proofs_image_id_and_pub_inputs,
+        input.is_aggregated_chunk,
+    );
 
     env::commit_slice(&merkle_root);
 }

aggregation_mode/aggregation_programs/sp1/src/lib.rs

@@ -8,17 +8,22 @@ pub struct SP1VkAndPubInputs {
 }
 
 impl SP1VkAndPubInputs {
-    pub fn hash(&self) -> [u8; 32] {
-        let mut hasher = Keccak256::new();
-        for &word in &self.vk {
-            hasher.update(word.to_be_bytes());
+    pub fn commitment(&self, is_aggregated_chunk: bool) -> [u8; 32] {
+        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()
         }
-        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,
 }

aggregation_mode/aggregation_programs/sp1/src/main.rs

@@ -13,12 +13,18 @@ fn combine_hashes(hash_a: &[u8; 32], hash_b: &[u8; 32]) -> [u8; 32] {
 }
 
 /// Computes the merkle root for the given proofs using the vk
-fn compute_merkle_root(proofs: &[SP1VkAndPubInputs]) -> [u8; 32] {
+fn compute_merkle_root(proofs: &[SP1VkAndPubInputs], is_aggregated_chunk: bool) -> [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()),
+            [a, b] => combine_hashes(
+                &a.commitment(is_aggregated_chunk),
+                &b.commitment(is_aggregated_chunk),
+            ),
+            [a] => combine_hashes(
+                &a.commitment(is_aggregated_chunk),
+                &a.commitment(is_aggregated_chunk),
+            ),
             _ => panic!("Unexpected chunk leaves"),
         })
         .collect();
@@ -48,7 +54,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_root =
+        compute_merkle_root(&input.proofs_vk_and_pub_inputs, input.is_aggregated_chunk);
 
     sp1_zkvm::io::commit_slice(&merkle_root);
 }

aggregation_mode/src/aggregators/mod.rs

@@ -5,12 +5,13 @@ use std::fmt::Display;
 
 use risc0_aggregator::{
     AlignedRisc0VerificationError, Risc0AggregationError, Risc0ProofReceiptAndImageId,
-    Risc0ProofType,
 };
 use sp1_aggregator::{
-    AlignedSP1VerificationError, SP1AggregationError, SP1ProofType, SP1ProofWithPubValuesAndElf,
+    AlignedSP1VerificationError, SP1AggregationError, SP1ProofWithPubValuesAndElf,
 };
 
+const AGG_PROOF_CHUNKS: usize = 512;
+
 #[derive(Clone, Debug)]
 pub enum ZKVMEngine {
     SP1,
@@ -68,29 +69,21 @@ impl ZKVMEngine {
                     })
                     .collect();
 
-                // we run the aggregator in chunks of 512 proofs
-                let chunks = proofs.chunks(512);
-                let mut agg_proofs: Vec<SP1ProofWithPubValuesAndElf> = vec![];
-
-                let agg_chunks_type = if chunks.len() == 1 {
-                    SP1ProofType::Groth16
-                } else {
-                    SP1ProofType::Compressed
-                };
-
-                for chunk in chunks {
-                    let agg_proof =
-                        sp1_aggregator::aggregate_proofs(chunk, agg_chunks_type.clone())
+                let mut agg_proof = if proofs.len() > AGG_PROOF_CHUNKS {
+                    let chunks = proofs.chunks(AGG_PROOF_CHUNKS);
+                    let mut agg_proofs: Vec<SP1ProofWithPubValuesAndElf> = vec![];
+                    for chunk in chunks {
+                        let agg_proof = sp1_aggregator::aggregate_proofs(chunk, false, false)
                             .map_err(ProofAggregationError::SP1Aggregation)?;
 
-                    agg_proofs.push(agg_proof);
-                }
+                        agg_proofs.push(agg_proof);
+                    }
 
-                let mut agg_proof = if agg_proofs.len() > 1 {
-                    sp1_aggregator::aggregate_proofs(&agg_proofs, SP1ProofType::Groth16)
+                    sp1_aggregator::aggregate_proofs(&agg_proofs, true, true)
                         .map_err(ProofAggregationError::SP1Aggregation)?
                 } else {
-                    agg_proofs.pop().unwrap()
+                    sp1_aggregator::aggregate_proofs(&proofs, false, true)
+                        .map_err(ProofAggregationError::SP1Aggregation)?
                 };
 
                 let merkle_root: [u8; 32] = agg_proof
@@ -109,28 +102,20 @@ impl ZKVMEngine {
                     })
                     .collect();
 
-                let chunks = proofs.chunks(512);
-                let mut agg_proofs: Vec<Risc0ProofReceiptAndImageId> = vec![];
-
-                let agg_chunks_type = if chunks.len() == 1 {
-                    Risc0ProofType::Groth16
-                } else {
-                    Risc0ProofType::Composite
-                };
-
-                for chunk in chunks {
-                    let agg_proof =
-                        risc0_aggregator::aggregate_proofs(chunk, agg_chunks_type.clone())
+                let agg_proof = if proofs.len() > AGG_PROOF_CHUNKS {
+                    let chunks = proofs.chunks(AGG_PROOF_CHUNKS);
+                    let mut agg_proofs: Vec<Risc0ProofReceiptAndImageId> = vec![];
+                    for chunk in chunks {
+                        let agg_proof = risc0_aggregator::aggregate_proofs(chunk, false, false)
                             .map_err(ProofAggregationError::Risc0Aggregation)?;
+                        agg_proofs.push(agg_proof);
+                    }
 
-                    agg_proofs.push(agg_proof);
-                }
-
-                let agg_proof = if agg_proofs.len() > 1 {
-                    risc0_aggregator::aggregate_proofs(&agg_proofs, Risc0ProofType::Groth16)
+                    risc0_aggregator::aggregate_proofs(&agg_proofs, true, true)
                         .map_err(ProofAggregationError::Risc0Aggregation)?
                 } else {
-                    agg_proofs.pop().unwrap()
+                    risc0_aggregator::aggregate_proofs(&proofs, false, true)
+                        .map_err(ProofAggregationError::Risc0Aggregation)?
                 };
 
                 // Note: journal.decode() won't work here as risc0 deserializer works under u32 words

aggregation_mode/src/aggregators/risc0_aggregator.rs

@@ -46,16 +46,10 @@ pub enum Risc0AggregationError {
     Verification(String),
 }
 
-#[derive(Debug, Clone)]
-pub enum Risc0ProofType {
-    Groth16,
-    Composite,
-    Succinct,
-}
-
 pub(crate) fn aggregate_proofs(
     proofs: &[Risc0ProofReceiptAndImageId],
-    to_proof_type: Risc0ProofType,
+    is_aggregated_chunk: bool,
+    should_wrap_to_groth16: bool,
 ) -> Result<Risc0ProofReceiptAndImageId, Risc0AggregationError> {
     let mut env_builder = ExecutorEnv::builder();
 
@@ -72,6 +66,7 @@ pub(crate) fn aggregate_proofs(
     // write input data
     let input = risc0_aggregation_program::Input {
         proofs_image_id_and_pub_inputs,
+        is_aggregated_chunk,
     };
     env_builder
         .write(&input)
@@ -83,10 +78,10 @@ pub(crate) fn aggregate_proofs(
 
     let prover = default_prover();
 
-    let opts = match to_proof_type {
-        Risc0ProofType::Groth16 => ProverOpts::groth16(),
-        Risc0ProofType::Composite => ProverOpts::composite(),
-        Risc0ProofType::Succinct => ProverOpts::succinct(),
+    let opts = if should_wrap_to_groth16 {
+        ProverOpts::groth16()
+    } else {
+        ProverOpts::composite()
     };
 
     let receipt = prover

aggregation_mode/src/aggregators/sp1_aggregator.rs

@@ -38,21 +38,16 @@ pub enum SP1AggregationError {
     UnsupportedProof,
 }
 
-#[derive(Debug, Clone)]
-pub enum SP1ProofType {
-    Groth16,
-    Compressed,
-    Core,
-}
-
 pub(crate) fn aggregate_proofs(
     proofs: &[SP1ProofWithPubValuesAndElf],
-    to_proof_type: SP1ProofType,
+    is_aggregated_chunk: bool,
+    should_wrap_to_groth16: bool,
 ) -> Result<SP1ProofWithPubValuesAndElf, SP1AggregationError> {
     let mut stdin = SP1Stdin::new();
 
     let mut program_input = sp1_aggregation_program::Input {
         proofs_vk_and_pub_inputs: vec![],
+        is_aggregated_chunk,
     };
 
     // write vk + public inputs
@@ -85,10 +80,10 @@ pub(crate) fn aggregate_proofs(
 
     let (pk, vk) = client.setup(PROGRAM_ELF);
     let proof_builder = client.prove(&pk, &stdin);
-    let proof_builder = match to_proof_type {
-        SP1ProofType::Groth16 => proof_builder.groth16(),
-        SP1ProofType::Compressed => proof_builder.compressed(),
-        SP1ProofType::Core => proof_builder.core(),
+    let proof_builder = if should_wrap_to_groth16 {
+        proof_builder.groth16()
+    } else {
+        proof_builder.compressed()
     };
 
     let proof = proof_builder