refactor(agg-mode): reduce types, abstractions and double merkle root in PI (#1889)

Co-authored-by: MauroFab <[email protected]>

Author: MarcosNicolau

aggregation_mode/aggregation_programs/risc0/src/lib.rs

@@ -24,5 +24,4 @@ impl Risc0ImageIdAndPubInputs {
 #[derive(Serialize, Deserialize)]
 pub struct Input {
     pub proofs_image_id_and_pub_inputs: Vec<Risc0ImageIdAndPubInputs>,
-    pub merkle_root: [u8; 32],
 }

aggregation_mode/aggregation_programs/risc0/src/main.rs

@@ -54,7 +54,5 @@ fn main() {
 
     let merkle_root = compute_merkle_root(&input.proofs_image_id_and_pub_inputs);
 
-    assert_eq!(merkle_root, input.merkle_root);
-
     env::commit_slice(&merkle_root);
 }

aggregation_mode/aggregation_programs/sp1/src/lib.rs

@@ -18,21 +18,7 @@ impl SP1VkAndPubInputs {
     }
 }
 
-#[derive(Serialize, Deserialize)]
-pub enum ProofVkAndPubInputs {
-    SP1Compressed(SP1VkAndPubInputs),
-}
-
-impl ProofVkAndPubInputs {
-    pub fn hash(&self) -> [u8; 32] {
-        match self {
-            ProofVkAndPubInputs::SP1Compressed(proof_data) => proof_data.hash(),
-        }
-    }
-}
-
 #[derive(Serialize, Deserialize)]
 pub struct Input {
-    pub proofs_vk_and_pub_inputs: Vec<ProofVkAndPubInputs>,
-    pub merkle_root: [u8; 32],
+    pub proofs_vk_and_pub_inputs: Vec<SP1VkAndPubInputs>,
 }

aggregation_mode/aggregation_programs/sp1/src/main.rs

@@ -3,7 +3,7 @@ sp1_zkvm::entrypoint!(main);
 
 use sha2::{Digest, Sha256};
 use sha3::Keccak256;
-use sp1_aggregation_program::{Input, ProofVkAndPubInputs};
+use sp1_aggregation_program::{Input, SP1VkAndPubInputs};
 
 fn combine_hashes(hash_a: &[u8; 32], hash_b: &[u8; 32]) -> [u8; 32] {
     let mut hasher = Keccak256::new();
@@ -13,7 +13,7 @@ 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: &[ProofVkAndPubInputs]) -> [u8; 32] {
+fn compute_merkle_root(proofs: &[SP1VkAndPubInputs]) -> [u8; 32] {
     let mut leaves: Vec<[u8; 32]> = proofs
         .chunks(2)
         .map(|chunk| match chunk {
@@ -42,19 +42,13 @@ pub fn main() {
 
     // Verify the proofs.
     for proof in input.proofs_vk_and_pub_inputs.iter() {
-        match proof {
-            ProofVkAndPubInputs::SP1Compressed(proof) => {
-                let vkey = proof.vk;
-                let public_values = &proof.public_inputs;
-                let public_values_digest = Sha256::digest(public_values);
-                sp1_zkvm::lib::verify::verify_sp1_proof(&vkey, &public_values_digest.into());
-            }
-        }
+        let vkey = proof.vk;
+        let public_values = &proof.public_inputs;
+        let public_values_digest = Sha256::digest(public_values);
+        sp1_zkvm::lib::verify::verify_sp1_proof(&vkey, &public_values_digest.into());
     }
 
     let merkle_root = compute_merkle_root(&input.proofs_vk_and_pub_inputs);
 
-    assert_eq!(merkle_root, input.merkle_root);
-
     sp1_zkvm::io::commit_slice(&merkle_root);
 }

aggregation_mode/src/aggregators/lib.rs

@@ -1,11 +1,14 @@
-pub mod lib;
 pub mod risc0_aggregator;
 pub mod sp1_aggregator;
 
 use std::fmt::Display;
 
-use risc0_aggregator::{AlignedRisc0VerificationError, Risc0ProofReceiptAndImageId};
-use sp1_aggregator::{AlignedSP1VerificationError, SP1ProofWithPubValuesAndElf};
+use risc0_aggregator::{
+    AlignedRisc0VerificationError, Risc0AggregationError, Risc0ProofReceiptAndImageId,
+};
+use sp1_aggregator::{
+    AlignedSP1VerificationError, SP1AggregationError, SP1ProofWithPubValuesAndElf,
+};
 
 #[derive(Clone, Debug)]
 pub enum ZKVMEngine {
@@ -22,6 +25,13 @@ impl Display for ZKVMEngine {
     }
 }
 
+#[derive(Debug)]
+pub enum ProofAggregationError {
+    SP1Aggregation(SP1AggregationError),
+    Risc0Aggregation(Risc0AggregationError),
+    PublicInputsDeserialization,
+}
+
 impl ZKVMEngine {
     pub fn from_env() -> Option<Self> {
         let key = "AGGREGATOR";
@@ -34,6 +44,69 @@ impl ZKVMEngine {
 
         Some(engine)
     }
+
+    /// Aggregates a list of [`AlignedProof`]s into a single [`AlignedProof`].
+    ///
+    /// Returns a tuple containing:
+    /// - The aggregated [`AlignedProof`], representing the combined proof
+    /// - The Merkle root computed within the ZKVM, exposed as a public input
+    ///
+    /// This function performs proof aggregation and ensures the resulting Merkle root
+    /// can be independently verified by external systems.
+    pub fn aggregate_proofs(
+        &self,
+        proofs: Vec<AlignedProof>,
+    ) -> Result<(AlignedProof, [u8; 32]), ProofAggregationError> {
+        let res = match self {
+            ZKVMEngine::SP1 => {
+                let proofs = proofs
+                    .into_iter()
+                    // Fetcher already filtered for SP1
+                    // We do this for type casting, as to avoid using generics
+                    // or macros in this function
+                    .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)?;
+
+                let merkle_root: [u8; 32] = agg_proof
+                    .proof_with_pub_values
+                    .public_values
+                    .read::<[u8; 32]>();
+
+                (AlignedProof::SP1(agg_proof.into()), merkle_root)
+            }
+            ZKVMEngine::RISC0 => {
+                let proofs = proofs
+                    .into_iter()
+                    // Fetcher already filtered for Risc0
+                    // We do this for type casting, as to avoid using generics
+                    // or macros in this function
+                    .filter_map(|proof| match proof {
+                        AlignedProof::Risc0(proof) => Some(*proof),
+                        _ => None,
+                    })
+                    .collect();
+
+                let agg_proof = risc0_aggregator::aggregate_proofs(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();
+                let merkle_root: [u8; 32] = public_input_bytes
+                    .try_into()
+                    .map_err(|_| ProofAggregationError::PublicInputsDeserialization)?;
+
+                (AlignedProof::Risc0(agg_proof.into()), merkle_root)
+            }
+        };
+
+        Ok(res)
+    }
 }
 
 pub enum AlignedProof {
@@ -42,7 +115,7 @@ pub enum AlignedProof {
 }
 
 impl AlignedProof {
-    pub fn hash(&self) -> [u8; 32] {
+    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(),

aggregation_mode/src/aggregators/mod.rs

@@ -3,8 +3,6 @@ include!(concat!(env!("OUT_DIR"), "/methods.rs"));
 use risc0_zkvm::{default_prover, ExecutorEnv, ProverOpts, Receipt};
 use sha3::{Digest, Keccak256};
 
-use super::lib::{AggregatedProof, ProgramOutput, ProofAggregationError};
-
 /// Byte representation of the aggregator image_id, converted from `[u32; 8]` to `[u8; 32]`.
 const RISC0_AGGREGATOR_PROGRAM_ID_BYTES: [u8; 32] = {
     let mut res = [0u8; 32];
@@ -40,19 +38,22 @@ impl Risc0ProofReceiptAndImageId {
     }
 }
 
-pub struct Risc0AggregationInput {
-    pub proofs: Vec<Risc0ProofReceiptAndImageId>,
-    pub merkle_root: [u8; 32],
+#[derive(Debug)]
+pub enum Risc0AggregationError {
+    WriteInput(String),
+    BuildExecutor(String),
+    Prove(String),
+    Verification(String),
 }
 
 pub(crate) fn aggregate_proofs(
-    input: Risc0AggregationInput,
-) -> Result<ProgramOutput, ProofAggregationError> {
+    proofs: Vec<Risc0ProofReceiptAndImageId>,
+) -> Result<Risc0ProofReceiptAndImageId, Risc0AggregationError> {
     let mut env_builder = ExecutorEnv::builder();
 
     // write assumptions and proof image id + pub inputs
     let mut proofs_image_id_and_pub_inputs = vec![];
-    for proof in input.proofs {
+    for proof in proofs {
         proofs_image_id_and_pub_inputs.push(risc0_aggregation_program::Risc0ImageIdAndPubInputs {
             image_id: proof.image_id,
             public_inputs: proof.receipt.journal.bytes.clone(),
@@ -62,34 +63,33 @@ pub(crate) fn aggregate_proofs(
 
     // write input data
     let input = risc0_aggregation_program::Input {
-        merkle_root: input.merkle_root,
         proofs_image_id_and_pub_inputs,
     };
     env_builder
         .write(&input)
-        .map_err(|e| ProofAggregationError::Risc0Proving(e.to_string()))?;
+        .map_err(|e| Risc0AggregationError::WriteInput(e.to_string()))?;
 
     let env = env_builder
         .build()
-        .map_err(|e| ProofAggregationError::Risc0Proving(e.to_string()))?;
+        .map_err(|e| Risc0AggregationError::BuildExecutor(e.to_string()))?;
 
     let prover = default_prover();
 
     let receipt = prover
         .prove_with_opts(env, RISC0_AGGREGATOR_PROGRAM_ELF, &ProverOpts::groth16())
-        .map_err(|e| ProofAggregationError::Risc0Proving(e.to_string()))?
+        .map_err(|e| Risc0AggregationError::Prove(e.to_string()))?
         .receipt;
 
     receipt
         .verify(RISC0_AGGREGATOR_PROGRAM_ID)
-        .map_err(|e| ProofAggregationError::Risc0Proving(e.to_string()))?;
+        .map_err(|e| Risc0AggregationError::Verification(e.to_string()))?;
 
-    let output = Risc0ProofReceiptAndImageId {
+    let proof = Risc0ProofReceiptAndImageId {
         image_id: RISC0_AGGREGATOR_PROGRAM_ID_BYTES,
         receipt,
     };
 
-    Ok(ProgramOutput::new(AggregatedProof::Risc0(output.into())))
+    Ok(proof)
 }
 
 #[derive(Debug)]

aggregation_mode/src/aggregators/risc0_aggregator.rs

@@ -1,14 +1,12 @@
 use std::sync::LazyLock;
 
 use alloy::primitives::Keccak256;
-use sp1_aggregation_program::{ProofVkAndPubInputs, SP1VkAndPubInputs};
+use sp1_aggregation_program::SP1VkAndPubInputs;
 use sp1_sdk::{
     EnvProver, HashableKey, Prover, ProverClient, SP1ProofWithPublicValues, SP1Stdin,
     SP1VerifyingKey,
 };
 
-use super::lib::{AggregatedProof, ProgramOutput, ProofAggregationError};
-
 const PROGRAM_ELF: &[u8] =
     include_bytes!("../../aggregation_programs/sp1/elf/sp1_aggregator_program");
 
@@ -33,38 +31,39 @@ impl SP1ProofWithPubValuesAndElf {
     }
 }
 
-pub struct SP1AggregationInput {
-    pub proofs: Vec<SP1ProofWithPubValuesAndElf>,
-    pub merkle_root: [u8; 32],
+#[derive(Debug)]
+pub enum SP1AggregationError {
+    Verification(sp1_sdk::SP1VerificationError),
+    Prove(String),
+    UnsupportedProof,
 }
 
 pub(crate) fn aggregate_proofs(
-    input: SP1AggregationInput,
-) -> Result<ProgramOutput, ProofAggregationError> {
+    proofs: Vec<SP1ProofWithPubValuesAndElf>,
+) -> Result<SP1ProofWithPubValuesAndElf, SP1AggregationError> {
     let mut stdin = SP1Stdin::new();
 
     let mut program_input = sp1_aggregation_program::Input {
         proofs_vk_and_pub_inputs: vec![],
-        merkle_root: input.merkle_root,
     };
 
     // write vk + public inputs
-    for proof in input.proofs.iter() {
+    for proof in proofs.iter() {
         program_input
             .proofs_vk_and_pub_inputs
-            .push(ProofVkAndPubInputs::SP1Compressed(SP1VkAndPubInputs {
+            .push(SP1VkAndPubInputs {
                 public_inputs: proof.proof_with_pub_values.public_values.to_vec(),
                 vk: proof.vk().hash_u32(),
-            }));
+            });
     }
     stdin.write(&program_input);
 
     // write proofs
-    for input_proof in input.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 {
-            return Err(ProofAggregationError::UnsupportedProof);
+            return Err(SP1AggregationError::UnsupportedProof);
         };
         stdin.write_proof(*proof, vk);
     }
@@ -80,21 +79,19 @@ pub(crate) fn aggregate_proofs(
         .prove(&pk, &stdin)
         .groth16()
         .run()
-        .map_err(|_| ProofAggregationError::SP1Proving)?;
+        .map_err(|e| SP1AggregationError::Prove(e.to_string()))?;
 
     // a sanity check, vm already performs it
     client
         .verify(&proof, &vk)
-        .map_err(ProofAggregationError::SP1Verification)?;
+        .map_err(SP1AggregationError::Verification)?;
 
     let proof_and_elf = SP1ProofWithPubValuesAndElf {
         proof_with_pub_values: proof,
         elf: PROGRAM_ELF.to_vec(),
     };
 
-    let output = ProgramOutput::new(AggregatedProof::SP1(proof_and_elf.into()));
-
-    Ok(output)
+    Ok(proof_and_elf)
 }
 
 #[derive(Debug)]