feat: move verification to filter to compute vk only once and in parallel

Author: MarcosNicolau

aggregation_mode/src/aggregators/mod.rs

@@ -219,24 +219,3 @@ impl IsMerkleTreeBackend for AlignedProof {
         hasher.finalize().into()
     }
 }
-
-#[derive(Debug)]
-pub enum AlignedVerificationError {
-    Sp1(AlignedSP1VerificationError),
-    Risc0(AlignedRisc0VerificationError),
-}
-
-impl AlignedProof {
-    pub fn verify(&self) -> Result<(), AlignedVerificationError> {
-        match self {
-            AlignedProof::SP1(proof) => sp1_aggregator::verify(proof).map_err(
-                |arg0: sp1_aggregator::AlignedSP1VerificationError| {
-                    AlignedVerificationError::Sp1(arg0)
-                },
-            ),
-            AlignedProof::Risc0(proof) => {
-                risc0_aggregator::verify(proof).map_err(AlignedVerificationError::Risc0)
-            }
-        }
-    }
-}

aggregation_mode/src/aggregators/risc0_aggregator.rs

@@ -8,7 +8,36 @@ pub struct Risc0ProofReceiptAndImageId {
     pub receipt: Receipt,
 }
 
+#[derive(Debug)]
+pub enum AlignedRisc0VerificationError {
+    Verification(String),
+    UnsupportedProof,
+}
+
 impl Risc0ProofReceiptAndImageId {
+    pub fn new(image_id: [u8; 32], receipt: Receipt) -> Self {
+        Self { image_id, receipt }
+    }
+
+    pub fn new_with_verification(
+        image_id: [u8; 32],
+        receipt: Receipt,
+    ) -> Result<Self, AlignedRisc0VerificationError> {
+        let is_supported_proof =
+            proof.receipt.inner.composite().is_ok() || proof.receipt.inner.succinct().is_ok();
+
+        if is_supported_proof {
+            proof
+                .receipt
+                .verify(proof.image_id)
+                .map_err(|e| AlignedRisc0VerificationError::Verification(e.to_string()))?;
+        } else {
+            return Err(AlignedRisc0VerificationError::UnsupportedProof);
+        }
+
+        Ok(Self { image_id, receipt })
+    }
+
     pub fn public_inputs(&self) -> &Vec<u8> {
         &self.receipt.journal.bytes
     }
@@ -22,12 +51,6 @@ pub enum Risc0AggregationError {
     Verification(String),
 }
 
-#[derive(Debug)]
-pub enum AlignedRisc0VerificationError {
-    Verification(String),
-    UnsupportedProof,
-}
-
 /// Byte representation of the user proofs aggregator image_id, converted from `[u32; 8]` to `[u8; 32]`.
 pub const RISC0_USER_PROOFS_AGGREGATOR_PROGRAM_ID_BYTES: [u8; 32] = {
     let mut res = [0u8; 32];
@@ -169,17 +192,3 @@ pub(crate) fn run_chunk_aggregator(
 
     Ok(proof)
 }
-
-pub(crate) fn verify(
-    proof: &Risc0ProofReceiptAndImageId,
-) -> Result<(), AlignedRisc0VerificationError> {
-    // only stark proofs are supported for recursion
-    if proof.receipt.inner.composite().is_ok() || proof.receipt.inner.succinct().is_ok() {
-        proof
-            .receipt
-            .verify(proof.image_id)
-            .map_err(|e| AlignedRisc0VerificationError::Verification(e.to_string()))
-    } else {
-        Err(AlignedRisc0VerificationError::UnsupportedProof)
-    }
-}

aggregation_mode/src/aggregators/sp1_aggregator.rs

@@ -26,6 +26,12 @@ pub struct SP1ProofWithPubValuesAndElf {
     pub vk: SP1VerifyingKey,
 }
 
+#[derive(Debug)]
+pub enum AlignedSP1VerificationError {
+    Verification(sp1_sdk::SP1VerificationError),
+    UnsupportedProof,
+}
+
 impl SP1ProofWithPubValuesAndElf {
     pub fn new(proof_with_pub_values: SP1ProofWithPublicValues, elf: Vec<u8>) -> Self {
         let vk = vk_from_elf(&elf);
@@ -37,6 +43,35 @@ impl SP1ProofWithPubValuesAndElf {
         }
     }
 
+    pub fn new_with_verification(
+        proof_with_pub_values: SP1ProofWithPublicValues,
+        elf: Vec<u8>,
+    ) -> Result<Self, AlignedSP1VerificationError> {
+        let client = &*SP1_PROVER_CLIENT_CPU;
+
+        let (_pk, vk) = client.setup(&sp1_proof_with_pub_values_and_elf.elf);
+
+        // only sp1 compressed proofs are supported for aggregation now
+        match sp1_proof_with_pub_values_and_elf
+            .proof_with_pub_values
+            .proof
+        {
+            sp1_sdk::SP1Proof::Compressed(_) => client
+                .verify(
+                    &sp1_proof_with_pub_values_and_elf.proof_with_pub_values,
+                    &vk,
+                )
+                .map_err(AlignedSP1VerificationError::Verification),
+            _ => Err(AlignedSP1VerificationError::UnsupportedProof),
+        }?;
+
+        Ok(Self {
+            proof_with_pub_values,
+            elf,
+            vk,
+        })
+    }
+
     pub fn hash_vk_and_pub_inputs(&self) -> [u8; 32] {
         let mut hasher = Keccak256::new();
         let vk_bytes = &self.vk.hash_bytes();
@@ -187,34 +222,6 @@ pub(crate) fn run_chunk_aggregator(
     Ok(proof_and_elf)
 }
 
-#[derive(Debug)]
-pub enum AlignedSP1VerificationError {
-    Verification(sp1_sdk::SP1VerificationError),
-    UnsupportedProof,
-}
-
-pub(crate) fn verify(
-    sp1_proof_with_pub_values_and_elf: &SP1ProofWithPubValuesAndElf,
-) -> Result<(), AlignedSP1VerificationError> {
-    let client = &*SP1_PROVER_CLIENT_CPU;
-
-    let (_pk, vk) = client.setup(&sp1_proof_with_pub_values_and_elf.elf);
-
-    // only sp1 compressed proofs are supported for aggregation now
-    match sp1_proof_with_pub_values_and_elf
-        .proof_with_pub_values
-        .proof
-    {
-        sp1_sdk::SP1Proof::Compressed(_) => client
-            .verify(
-                &sp1_proof_with_pub_values_and_elf.proof_with_pub_values,
-                &vk,
-            )
-            .map_err(AlignedSP1VerificationError::Verification),
-        _ => Err(AlignedSP1VerificationError::UnsupportedProof),
-    }
-}
-
 pub fn vk_from_elf(elf: &[u8]) -> SP1VerifyingKey {
     let prover = &*SP1_PROVER_CLIENT_CPU;
     let (_, vk) = prover.setup(elf);

aggregation_mode/src/backend/fetcher.rs

@@ -112,38 +112,61 @@ impl ProofsFetcher {
             info!("Data downloaded from S3, number of proofs {}", data.len());
 
             // Filter compatible proofs to be aggregated and push to queue
-            let mut proofs_to_add: Vec<AlignedProof> = match engine {
+            let proofs_to_add: Vec<AlignedProof> = match engine {
                 ZKVMEngine::SP1 => data
-                    .into_iter()
-                    .filter_map(|p| match p.proving_system {
-                        ProvingSystemId::SP1 => {
-                            let elf = p.vm_program_code?;
-                            let proof_with_pub_values = bincode::deserialize(&p.proof).ok()?;
-                            let sp1_proof =
-                                SP1ProofWithPubValuesAndElf::new(proof_with_pub_values, elf);
-
-                            Some(AlignedProof::SP1(sp1_proof.into()))
+                    .into_par_iter()
+                    .filter_map(|p| {
+                        if p.proving_system != ProvingSystemId::SP1 {
+                            return None;
+                        };
+
+                        let elf = p.vm_program_code?;
+                        let proof_with_pub_values = bincode::deserialize(&p.proof).ok()?;
+                        let sp1_proof = if self.pre_verification_enabled {
+                            SP1ProofWithPubValuesAndElf::new_with_verification(
+                                proof_with_pub_values,
+                                elf,
+                            )
+                        } else {
+                            Ok(SP1ProofWithPubValuesAndElf::new(proof_with_pub_values, elf))
+                        };
+
+                        match sp1_proof {
+                            Ok(proof) => Some(AlignedProof::SP1(sp1_proof)),
+                            Err(err) => {
+                                error!("Could not add proof, verification failed: {:?}", err);
+                                None
+                            }
                         }
-
-                        _ => None,
                     })
                     .collect(),
                 ZKVMEngine::RISC0 => data
-                    .into_iter()
-                    .filter_map(|p| match p.proving_system {
-                        ProvingSystemId::Risc0 => {
-                            let mut image_id = [0u8; 32];
-                            image_id.copy_from_slice(p.vm_program_code?.as_slice());
-                            let public_inputs = p.pub_input?;
-                            let inner_receipt: risc0_zkvm::InnerReceipt =
-                                bincode::deserialize(&p.proof).ok()?;
-
-                            let receipt = Receipt::new(inner_receipt, public_inputs);
-                            let risc0_proof = Risc0ProofReceiptAndImageId { image_id, receipt };
-
-                            Some(AlignedProof::Risc0(risc0_proof.into()))
+                    .into_par_iter()
+                    .filter_map(|p| {
+                        if p.proving_system != ProvingSystemId::Risc0 {
+                            return None;
+                        };
+
+                        let mut image_id = [0u8; 32];
+                        image_id.copy_from_slice(p.vm_program_code?.as_slice());
+                        let public_inputs = p.pub_input?;
+                        let inner_receipt: risc0_zkvm::InnerReceipt =
+                            bincode::deserialize(&p.proof).ok()?;
+
+                        let receipt = Receipt::new(inner_receipt, public_inputs);
+                        let risc0_proof = if self.pre_verification_enabled {
+                            Risc0ProofReceiptAndImageId::new_with_verification(image_id, receipt)
+                        } else {
+                            Ok(Risc0ProofReceiptAndImageId::new(image_id, receipt))
+                        };
+
+                        match risc0_proof {
+                            Ok(proof) => Some(AlignedProof::Risc0(risc0_proof)),
+                            Err(err) => {
+                                error!("Could not add proof, verification failed: {:?}", err);
+                                None
+                            }
                         }
-                        _ => None,
                     })
                     .collect(),
             };
@@ -154,22 +177,6 @@ impl ProofsFetcher {
                 proofs_to_add.len()
             );
 
-            if self.pre_verification_enabled {
-                // Try to add them to the queue
-                // We do this in parallel, as SP1 can take quite some time in verifying
-                // because of the overhead of setting up the prover
-                proofs_to_add = proofs_to_add
-                    .into_par_iter()
-                    .filter(|proof| match proof.verify() {
-                        Ok(_) => true,
-                        Err(err) => {
-                            error!("Could not add proof, verification failed: {:?}", err);
-                            return false;
-                        }
-                    })
-                    .collect();
-            }
-
             proofs.extend(proofs_to_add);
         }