feat: fetch proofs based on zkvm engine

Author: MarcosNicolau

aggregation_mode/Cargo.toml

@@ -22,7 +22,7 @@ aligned-sdk = { path = "../batcher/aligned-sdk/" }
 # zkvms
 sp1-sdk = "4.1.3"
 sp1_aggregation_program = { path = "./aggregation_programs/sp1/" }
-risc0-zkvm = { version = "2.0.1"}
+risc0-zkvm = { version = "2.0.1" }
 risc0_aggregation_program = { path = "./aggregation_programs/risc0/" }
 risc0-ethereum-contracts = { git = "https://github.com/risc0/risc0-ethereum/", tag = "v2.0.0" }
 
@@ -38,4 +38,7 @@ methods = ["./aggregation_programs/risc0"]
 opt-level = 3
 
 [features]
+default = []
 prove = []
+risc0 = []
+sp1 = []

aggregation_mode/src/aggregators/mod.rs

@@ -2,13 +2,42 @@ 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};
+
+#[derive(Clone, Debug)]
 pub enum ZKVMEngine {
     SP1,
     RISC0,
 }
 
+impl Display for ZKVMEngine {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        match self {
+            Self::SP1 => write!(f, "SP1"),
+            Self::RISC0 => write!(f, "Risc0"),
+        }
+    }
+}
+
+impl ZKVMEngine {
+    pub fn from_rust_features() -> Option<Self> {
+        #[cfg(feature = "sp1")]
+        {
+            return Some(ZKVMEngine::SP1);
+        }
+
+        #[cfg(feature = "risc0")]
+        {
+            return Some(ZKVMEngine::RISC0);
+        }
+
+        None
+    }
+}
+
 pub enum AlignedProof {
     SP1(SP1ProofWithPubValuesAndElf),
     Risc0(Risc0ProofReceiptAndImageId),

aggregation_mode/src/aggregators/risc0_aggregator.rs

@@ -40,7 +40,7 @@ impl Risc0ProofReceiptAndImageId {
 }
 
 pub struct Risc0AggregationInput {
-    pub receipts: Vec<Risc0ProofReceiptAndImageId>,
+    pub proofs: Vec<Risc0ProofReceiptAndImageId>,
     pub merkle_root: [u8; 32],
 }
 
@@ -51,7 +51,7 @@ pub(crate) fn aggregate_proofs(
 
     // write assumptions and proof image id + pub inputs
     let mut proofs_image_id_and_pub_inputs = vec![];
-    for proof in input.receipts {
+    for proof in input.proofs {
         proofs_image_id_and_pub_inputs.push(risc0_aggregation_program::Risc0ImageIdAndPubInputs {
             image_id: proof.image_id,
             public_inputs: proof.receipt.journal.bytes.clone(),

aggregation_mode/src/backend/fetcher.rs

@@ -7,7 +7,7 @@ use super::{
 use crate::{
     aggregators::{
         risc0_aggregator::Risc0ProofReceiptAndImageId, sp1_aggregator::SP1ProofWithPubValuesAndElf,
-        AlignedProof,
+        AlignedProof, ZKVMEngine,
     },
     backend::s3::get_aligned_batch_from_s3,
 };
@@ -50,7 +50,7 @@ impl ProofsFetcher {
         }
     }
 
-    pub async fn fetch(&self) -> Result<Vec<AlignedProof>, ProofsFetcherError> {
+    pub async fn fetch(&self, engine: ZKVMEngine) -> Result<Vec<AlignedProof>, ProofsFetcherError> {
         let from_block = self.get_block_number_to_fetch_from().await?;
         info!(
             "Fetching proofs from batch logs starting from block number {}",
@@ -87,37 +87,47 @@ impl ProofsFetcher {
             info!("Data downloaded from S3, number of proofs {}", data.len());
 
             // Filter compatible proofs to be aggregated and push to queue
-            let proofs_to_add: Vec<AlignedProof> = 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 {
-                            proof_with_pub_values,
-                            elf,
-                        };
-
-                        Some(AlignedProof::SP1(sp1_proof))
-                    }
-                    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))
-                    }
-                    _ => None,
-                })
-                .collect();
+            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 {
+                                proof_with_pub_values,
+                                elf,
+                            };
+
+                            Some(AlignedProof::SP1(sp1_proof))
+                        }
+
+                        _ => 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))
+                        }
+                        _ => None,
+                    })
+                    .collect(),
+            };
 
             info!(
-                "Proofs filtered, total proofs to add {}",
+                "{} Proofs filtered, compatible proofs found {}",
+                engine,
                 proofs_to_add.len()
             );
 

aggregation_mode/src/backend/mod.rs

@@ -63,10 +63,12 @@ impl ProofAggregator {
                 .expect("AlignedProofAggregationService address should be valid"),
             rpc_provider,
         );
+
+        let engine = ZKVMEngine::from_rust_features().expect("A feature defining zkvm engine");
         let fetcher = ProofsFetcher::new(config);
 
         Self {
-            engine: ZKVMEngine::SP1,
+            engine,
             proof_aggregation_service,
             fetcher,
         }
@@ -93,7 +95,7 @@ impl ProofAggregator {
     ) -> Result<(), AggregatedProofSubmissionError> {
         let proofs = self
             .fetcher
-            .fetch()
+            .fetch(self.engine.clone())
             .await
             .map_err(AggregatedProofSubmissionError::FetchingProofs)?;
 
@@ -127,7 +129,7 @@ impl ProofAggregator {
                     .map_err(AggregatedProofSubmissionError::Aggregation)?
             }
             ZKVMEngine::RISC0 => {
-                let receipts = proofs
+                let proofs = proofs
                     .into_iter()
                     .filter_map(|proof| match proof {
                         AlignedProof::Risc0(proof) => Some(proof),
@@ -136,7 +138,7 @@ impl ProofAggregator {
                     .collect();
 
                 let input = Risc0AggregationInput {
-                    receipts,
+                    proofs,
                     merkle_root,
                 };