feat: aggregate proofs in chunks of 512 proofs

Author: MarcosNicolau

aggregation_mode/src/aggregators/mod.rs

@@ -5,9 +5,10 @@ use std::fmt::Display;
 
 use risc0_aggregator::{
     AlignedRisc0VerificationError, Risc0AggregationError, Risc0ProofReceiptAndImageId,
+    Risc0ProofType,
 };
 use sp1_aggregator::{
-    AlignedSP1VerificationError, SP1AggregationError, SP1ProofWithPubValuesAndElf,
+    AlignedSP1VerificationError, SP1AggregationError, SP1ProofType, SP1ProofWithPubValuesAndElf,
 };
 
 #[derive(Clone, Debug)]
@@ -59,16 +60,38 @@ impl ZKVMEngine {
     ) -> Result<(AlignedProof, [u8; 32]), ProofAggregationError> {
         let res = match self {
             ZKVMEngine::SP1 => {
-                let proofs = proofs
+                let proofs: Vec<SP1ProofWithPubValuesAndElf> = proofs
                     .into_iter()
                     .filter_map(|proof| match proof {
                         AlignedProof::SP1(proof) => Some(*proof),
                         _ => None,
                     })
                     .collect();
 
-                let mut agg_proof = sp1_aggregator::aggregate_proofs(proofs)
-                    .map_err(ProofAggregationError::SP1Aggregation)?;
+                // 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())
+                            .map_err(ProofAggregationError::SP1Aggregation)?;
+
+                    agg_proofs.push(agg_proof);
+                }
+
+                let mut agg_proof = if agg_proofs.len() > 1 {
+                    sp1_aggregator::aggregate_proofs(&agg_proofs, SP1ProofType::Groth16)
+                        .map_err(ProofAggregationError::SP1Aggregation)?
+                } else {
+                    agg_proofs.pop().unwrap()
+                };
 
                 let merkle_root: [u8; 32] = agg_proof
                     .proof_with_pub_values
@@ -78,16 +101,37 @@ impl ZKVMEngine {
                 (AlignedProof::SP1(agg_proof.into()), merkle_root)
             }
             ZKVMEngine::RISC0 => {
-                let proofs = proofs
+                let proofs: Vec<Risc0ProofReceiptAndImageId> = proofs
                     .into_iter()
                     .filter_map(|proof| match proof {
                         AlignedProof::Risc0(proof) => Some(*proof),
                         _ => None,
                     })
                     .collect();
 
-                let agg_proof = risc0_aggregator::aggregate_proofs(proofs)
-                    .map_err(ProofAggregationError::Risc0Aggregation)?;
+                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())
+                            .map_err(ProofAggregationError::Risc0Aggregation)?;
+
+                    agg_proofs.push(agg_proof);
+                }
+
+                let agg_proof = if agg_proofs.len() > 1 {
+                    risc0_aggregator::aggregate_proofs(&agg_proofs, Risc0ProofType::Groth16)
+                        .map_err(ProofAggregationError::Risc0Aggregation)?
+                } else {
+                    agg_proofs.pop().unwrap()
+                };
 
                 // Note: journal.decode() won't work here as risc0 deserializer works under u32 words
                 let public_input_bytes = agg_proof.receipt.journal.as_ref();

aggregation_mode/src/aggregators/risc0_aggregator.rs

@@ -46,8 +46,16 @@ pub enum Risc0AggregationError {
     Verification(String),
 }
 
+#[derive(Debug, Clone)]
+pub enum Risc0ProofType {
+    Groth16,
+    Composite,
+    Succinct,
+}
+
 pub(crate) fn aggregate_proofs(
-    proofs: Vec<Risc0ProofReceiptAndImageId>,
+    proofs: &[Risc0ProofReceiptAndImageId],
+    to_proof_type: Risc0ProofType,
 ) -> Result<Risc0ProofReceiptAndImageId, Risc0AggregationError> {
     let mut env_builder = ExecutorEnv::builder();
 
@@ -58,7 +66,7 @@ pub(crate) fn aggregate_proofs(
             image_id: proof.image_id,
             public_inputs: proof.receipt.journal.bytes.clone(),
         });
-        env_builder.add_assumption(proof.receipt);
+        env_builder.add_assumption(proof.receipt.clone());
     }
 
     // write input data
@@ -75,8 +83,14 @@ 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 receipt = prover
-        .prove_with_opts(env, RISC0_AGGREGATOR_PROGRAM_ELF, &ProverOpts::groth16())
+        .prove_with_opts(env, RISC0_AGGREGATOR_PROGRAM_ELF, &opts)
         .map_err(|e| Risc0AggregationError::Prove(e.to_string()))?
         .receipt;
 

aggregation_mode/src/aggregators/sp1_aggregator.rs

@@ -38,8 +38,16 @@ pub enum SP1AggregationError {
     UnsupportedProof,
 }
 
+#[derive(Debug, Clone)]
+pub enum SP1ProofType {
+    Groth16,
+    Compressed,
+    Core,
+}
+
 pub(crate) fn aggregate_proofs(
-    proofs: Vec<SP1ProofWithPubValuesAndElf>,
+    proofs: &[SP1ProofWithPubValuesAndElf],
+    to_proof_type: SP1ProofType,
 ) -> Result<SP1ProofWithPubValuesAndElf, SP1AggregationError> {
     let mut stdin = SP1Stdin::new();
 
@@ -62,7 +70,8 @@ pub(crate) fn aggregate_proofs(
     for input_proof in proofs {
         let vk = input_proof.vk().vk;
         // we only support sp1 Compressed proofs for now
-        let sp1_sdk::SP1Proof::Compressed(proof) = input_proof.proof_with_pub_values.proof else {
+        let sp1_sdk::SP1Proof::Compressed(proof) = input_proof.proof_with_pub_values.proof.clone()
+        else {
             return Err(SP1AggregationError::UnsupportedProof);
         };
         stdin.write_proof(*proof, vk);
@@ -75,9 +84,14 @@ pub(crate) fn aggregate_proofs(
     let client = ProverClient::builder().mock().build();
 
     let (pk, vk) = client.setup(PROGRAM_ELF);
-    let proof = client
-        .prove(&pk, &stdin)
-        .groth16()
+    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 = proof_builder
         .run()
         .map_err(|e| SP1AggregationError::Prove(e.to_string()))?;