feat: aggregator backends for new programs

MarcosNicolau · MarcosNicolau · commit 5335d2a219c2 · 2025-05-15T10:28:31.000-03:00
diff --git a/aggregation_mode/src/aggregators/mod.rs b/aggregation_mode/src/aggregators/mod.rs
@@ -85,26 +85,27 @@ impl ZKVMEngine {
                     .collect();
 
                 info!("Total proofs to aggregate {}", proofs.len());
-                let mut agg_proof = if proofs.len() > MAX_PROOFS_PER_AGGREGATION {
-                    let chunks = proofs.chunks(MAX_PROOFS_PER_AGGREGATION);
-                    let mut agg_proofs: Vec<SP1ProofWithPubValuesAndElf> = vec![];
-
-                    info!("Proofs length is higher than {}, aggregation will be performed in {} chunks", MAX_PROOFS_PER_AGGREGATION, chunks.len());
-                    for (i, chunk) in chunks.enumerate() {
-                        let agg_proof = sp1_aggregator::aggregate_proofs(chunk, false, false)
-                            .map_err(ProofAggregationError::SP1Aggregation)?;
-                        agg_proofs.push(agg_proof);
-
-                        info!("Chunk number {} has been aggregated", i);
-                    }
-
-                    info!("All chunks have been aggregated, performing last aggregation...");
-                    sp1_aggregator::aggregate_proofs(&agg_proofs, true, true)
-                        .map_err(ProofAggregationError::SP1Aggregation)?
-                } else {
-                    sp1_aggregator::aggregate_proofs(&proofs, false, true)
-                        .map_err(ProofAggregationError::SP1Aggregation)?
-                };
+                let chunks = proofs.chunks(MAX_PROOFS_PER_AGGREGATION);
+                let mut agg_proofs: Vec<(SP1ProofWithPubValuesAndElf, Vec<[u8; 32]>)> = vec![];
+
+                info!(
+                    "Proofs length is higher than {}, aggregation will be performed in {} chunks",
+                    MAX_PROOFS_PER_AGGREGATION,
+                    chunks.len()
+                );
+                for (i, chunk) in chunks.enumerate() {
+                    let leaves_commitment =
+                        chunk.iter().map(|e| e.hash_vk_and_pub_inputs()).collect();
+                    let agg_proof = sp1_aggregator::run_chunk_aggregator(chunk)
+                        .map_err(ProofAggregationError::SP1Aggregation)?;
+                    agg_proofs.push((agg_proof, leaves_commitment));
+
+                    info!("Chunk number {} has been aggregated", i);
+                }
+
+                info!("All chunks have been aggregated, performing last aggregation...");
+                let mut agg_proof = sp1_aggregator::run_root_aggregator(&agg_proofs)
+                    .map_err(ProofAggregationError::SP1Aggregation)?;
 
                 let merkle_root: [u8; 32] = agg_proof
                     .proof_with_pub_values
@@ -126,26 +127,29 @@ impl ZKVMEngine {
                     .collect();
 
                 info!("Total proofs to aggregate {}", proofs.len());
-                let agg_proof = if proofs.len() > MAX_PROOFS_PER_AGGREGATION {
-                    let chunks = proofs.chunks(MAX_PROOFS_PER_AGGREGATION);
-                    let mut agg_proofs: Vec<Risc0ProofReceiptAndImageId> = vec![];
-
-                    info!("Proofs length is higher than {}, aggregation will be performed in {} chunks", MAX_PROOFS_PER_AGGREGATION, chunks.len());
-                    for (i, chunk) in chunks.enumerate() {
-                        let agg_proof = risc0_aggregator::aggregate_proofs(chunk, false, false)
-                            .map_err(ProofAggregationError::Risc0Aggregation)?;
-                        agg_proofs.push(agg_proof);
-
-                        info!("Chunk number {} has been aggregated", i);
-                    }
-
-                    info!("All chunks have been aggregated, performing last aggregation...");
-                    risc0_aggregator::aggregate_proofs(&agg_proofs, true, true)
-                        .map_err(ProofAggregationError::Risc0Aggregation)?
-                } else {
-                    risc0_aggregator::aggregate_proofs(&proofs, false, true)
-                        .map_err(ProofAggregationError::Risc0Aggregation)?
-                };
+                let chunks = proofs.chunks(MAX_PROOFS_PER_AGGREGATION);
+                let mut agg_proofs: Vec<(Risc0ProofReceiptAndImageId, Vec<[u8; 32]>)> = vec![];
+
+                info!(
+                    "Proofs length is higher than {}, aggregation will be performed in {} chunks",
+                    MAX_PROOFS_PER_AGGREGATION,
+                    chunks.len()
+                );
+                for (i, chunk) in chunks.enumerate() {
+                    let leaves_commitment = chunk
+                        .iter()
+                        .map(|e| e.hash_image_id_and_public_inputs())
+                        .collect();
+                    let agg_proof = risc0_aggregator::run_chunk_aggregator(chunk)
+                        .map_err(ProofAggregationError::Risc0Aggregation)?;
+                    agg_proofs.push((agg_proof, leaves_commitment));
+
+                    info!("Chunk number {} has been aggregated", i);
+                }
+
+                info!("All chunks have been aggregated, performing last aggregation...");
+                let agg_proof = risc0_aggregator::run_root_aggregator(&agg_proofs)
+                    .map_err(ProofAggregationError::Risc0Aggregation)?;
 
                 // Note: journal.decode() won't work here as risc0 deserializer works under u32 words
                 let public_input_bytes = agg_proof.receipt.journal.as_ref();
@@ -170,9 +174,7 @@ impl AlignedProof {
     pub fn commitment(&self) -> [u8; 32] {
         match self {
             AlignedProof::SP1(proof) => proof.hash_vk_and_pub_inputs(),
-            //AlignedProof::Risc0(proof) => proof.hash_image_id_and_public_inputs(),
-            AlignedProof::Risc0(proof) => todo!(),
-
+            AlignedProof::Risc0(proof) => proof.hash_image_id_and_public_inputs(),
         }
     }
 }
@@ -225,8 +227,7 @@ impl AlignedProof {
                 },
             ),
             AlignedProof::Risc0(proof) => {
-                todo!()
-                //risc0_aggregator::verify(proof).map_err(AlignedVerificationError::Risc0)
+                risc0_aggregator::verify(proof).map_err(AlignedVerificationError::Risc0)
             }
         }
     }
diff --git a/aggregation_mode/src/aggregators/risc0_aggregator.rs b/aggregation_mode/src/aggregators/risc0_aggregator.rs
@@ -1,9 +1,8 @@
-//include!(concat!(env!("OUT_DIR"), "/methods.rs"));
+include!(concat!(env!("OUT_DIR"), "/methods.rs"));
 
 use risc0_zkvm::{default_prover, ExecutorEnv, ProverOpts, Receipt};
 use sha3::{Digest, Keccak256};
 
-
 pub struct Risc0ProofReceiptAndImageId {
     pub image_id: [u8; 32],
     pub receipt: Receipt,
@@ -29,13 +28,12 @@ pub enum AlignedRisc0VerificationError {
     UnsupportedProof,
 }
 
-/// Byte representation of the aggregator image_id, converted from `[u32; 8]` to `[u8; 32]`.
-/* 
-pub const RISC0_AGGREGATOR_PROGRAM_ID_BYTES: [u8; 32] = {
+/// Byte representation of the chunk aggregator image_id, converted from `[u32; 8]` to `[u8; 32]`.
+pub const RISC0_CHUNK_AGGREGATOR_PROGRAM_ID_BYTES: [u8; 32] = {
     let mut res = [0u8; 32];
     let mut i = 0;
     while i < 8 {
-        let bytes = RISC0_AGGREGATOR_PROGRAM_ID[i].to_le_bytes();
+        let bytes = RISC0_CHUNK_AGGREGATOR_PROGRAM_ID[i].to_le_bytes();
         res[i * 4] = bytes[0];
         res[i * 4 + 1] = bytes[1];
         res[i * 4 + 2] = bytes[2];
@@ -45,6 +43,20 @@ pub const RISC0_AGGREGATOR_PROGRAM_ID_BYTES: [u8; 32] = {
     res
 };
 
+/// Byte representation of the root aggregator image_id, converted from `[u32; 8]` to `[u8; 32]`.
+pub const RISC0_ROOT_AGGREGATOR_PROGRAM_ID_BYTES: [u8; 32] = {
+    let mut res = [0u8; 32];
+    let mut i = 0;
+    while i < 8 {
+        let bytes = RISC0_ROOT_AGGREGATOR_PROGRAM_ID[i].to_le_bytes();
+        res[i * 4] = bytes[0];
+        res[i * 4 + 1] = bytes[1];
+        res[i * 4 + 2] = bytes[2];
+        res[i * 4 + 3] = bytes[3];
+        i += 1;
+    }
+    res
+};
 
 impl Risc0ProofReceiptAndImageId {
     pub fn hash_image_id_and_public_inputs(&self) -> [u8; 32] {
@@ -55,11 +67,8 @@ impl Risc0ProofReceiptAndImageId {
     }
 }
 
-
-pub(crate) fn aggregate_proofs(
+pub(crate) fn run_chunk_aggregator(
     proofs: &[Risc0ProofReceiptAndImageId],
-    is_aggregated_chunk: bool,
-    should_wrap_to_groth16: bool,
 ) -> Result<Risc0ProofReceiptAndImageId, Risc0AggregationError> {
     let mut env_builder = ExecutorEnv::builder();
 
@@ -74,8 +83,8 @@ pub(crate) fn aggregate_proofs(
     }
 
     // write input data
-    let input = risc0_aggregation_program::Input {
-        proofs_image_id_and_pub_inputs
+    let input = risc0_aggregation_program::ChunkAggregatorInput {
+        proofs_image_id_and_pub_inputs,
     };
     env_builder
         .write(&input)
@@ -87,30 +96,79 @@ pub(crate) fn aggregate_proofs(
 
     let prover = default_prover();
 
-    let opts = if should_wrap_to_groth16 {
-        ProverOpts::groth16()
-    } else {
-        ProverOpts::composite()
-    };
-
     let receipt = prover
-        .prove_with_opts(env, RISC0_AGGREGATOR_PROGRAM_ELF, &opts)
+        .prove_with_opts(
+            env,
+            RISC0_CHUNK_AGGREGATOR_PROGRAM_ELF,
+            &ProverOpts::composite(),
+        )
         .map_err(|e| Risc0AggregationError::Prove(e.to_string()))?
         .receipt;
 
     receipt
-        .verify(RISC0_AGGREGATOR_PROGRAM_ID)
+        .verify(RISC0_CHUNK_AGGREGATOR_PROGRAM_ID)
         .map_err(|e| Risc0AggregationError::Verification(e.to_string()))?;
 
     let proof = Risc0ProofReceiptAndImageId {
-        image_id: RISC0_AGGREGATOR_PROGRAM_ID_BYTES,
+        image_id: RISC0_CHUNK_AGGREGATOR_PROGRAM_ID_BYTES,
         receipt,
     };
 
     Ok(proof)
 }
 
+pub(crate) fn run_root_aggregator(
+    proofs: &[(Risc0ProofReceiptAndImageId, Vec<[u8; 32]>)],
+) -> Result<Risc0ProofReceiptAndImageId, Risc0AggregationError> {
+    let mut env_builder = ExecutorEnv::builder();
+
+    // write assumptions and proof image id + pub inputs
+    let mut proofs_and_leaves_commitment = vec![];
+    for (proof, leaves_commitment) in proofs {
+        proofs_and_leaves_commitment.push((
+            risc0_aggregation_program::Risc0ImageIdAndPubInputs {
+                image_id: proof.image_id,
+                public_inputs: proof.receipt.journal.bytes.clone(),
+            },
+            leaves_commitment.clone(),
+        ));
+        env_builder.add_assumption(proof.receipt.clone());
+    }
 
+    // write input data
+    let input = risc0_aggregation_program::RootAggregatorInput {
+        proofs_and_leaves_commitment: proofs_and_leaves_commitment,
+    };
+    env_builder
+        .write(&input)
+        .map_err(|e| Risc0AggregationError::WriteInput(e.to_string()))?;
+
+    let env = env_builder
+        .build()
+        .map_err(|e| Risc0AggregationError::BuildExecutor(e.to_string()))?;
+
+    let prover = default_prover();
+
+    let receipt = prover
+        .prove_with_opts(
+            env,
+            RISC0_ROOT_AGGREGATOR_PROGRAM_ELF,
+            &ProverOpts::groth16(),
+        )
+        .map_err(|e| Risc0AggregationError::Prove(e.to_string()))?
+        .receipt;
+
+    receipt
+        .verify(RISC0_ROOT_AGGREGATOR_PROGRAM_ID)
+        .map_err(|e| Risc0AggregationError::Verification(e.to_string()))?;
+
+    let proof = Risc0ProofReceiptAndImageId {
+        image_id: RISC0_ROOT_AGGREGATOR_PROGRAM_ID_BYTES,
+        receipt,
+    };
+
+    Ok(proof)
+}
 
 pub(crate) fn verify(
     proof: &Risc0ProofReceiptAndImageId,
@@ -125,19 +183,3 @@ pub(crate) fn verify(
         Err(AlignedRisc0VerificationError::UnsupportedProof)
     }
 }
-*/
-
-pub(crate) fn aggregate_proofs(
-    proofs: &[Risc0ProofReceiptAndImageId],
-    is_aggregated_chunk: bool,
-    should_wrap_to_groth16: bool,
-) -> Result<Risc0ProofReceiptAndImageId, Risc0AggregationError> {
-
-    todo!()
-}
-
-pub(crate) fn verify(
-    proof: &Risc0ProofReceiptAndImageId,
-) -> Result<(), AlignedRisc0VerificationError> {
-    todo!();
-}
diff --git a/aggregation_mode/src/aggregators/sp1_aggregator.rs b/aggregation_mode/src/aggregators/sp1_aggregator.rs
