feat: make provers generic over engine (#1553)

Golovanov399 · web-flow · commit 40927aa95fd1 · 2025-04-02T16:35:28.000-07:00
This adds the ability to use different engines for proving.
diff --git a/benchmarks/src/bin/fib_e2e.rs b/benchmarks/src/bin/fib_e2e.rs
@@ -12,7 +12,9 @@ use openvm_rv32im_transpiler::{
 use openvm_sdk::{
     commit::commit_app_exe, prover::ContinuationProver, DefaultStaticVerifierPvHandler, Sdk, StdIn,
 };
-use openvm_stark_sdk::bench::run_with_metric_collection;
+use openvm_stark_sdk::{
+    bench::run_with_metric_collection, config::baby_bear_poseidon2::BabyBearPoseidon2Engine,
+};
 use openvm_transpiler::{transpiler::Transpiler, FromElf};
 
 const NUM_PUBLIC_VALUES: usize = DEFAULT_MAX_NUM_PUBLIC_VALUES;
@@ -56,8 +58,12 @@ async fn main() -> Result<()> {
     let mut stdin = StdIn::default();
     stdin.write(&n);
     run_with_metric_collection("OUTPUT_PATH", || {
-        let mut e2e_prover =
-            ContinuationProver::new(&halo2_params_reader, app_pk, app_committed_exe, full_agg_pk);
+        let mut e2e_prover = ContinuationProver::<_, BabyBearPoseidon2Engine>::new(
+            &halo2_params_reader,
+            app_pk,
+            app_committed_exe,
+            full_agg_pk,
+        );
         e2e_prover.set_program_name("fib_e2e");
         let _proof = e2e_prover.generate_proof_for_evm(stdin);
     });
diff --git a/benchmarks/src/bin/kitchen_sink.rs b/benchmarks/src/bin/kitchen_sink.rs
@@ -13,7 +13,9 @@ use openvm_sdk::{
     commit::commit_app_exe, config::SdkVmConfig, prover::ContinuationProver,
     DefaultStaticVerifierPvHandler, Sdk, StdIn,
 };
-use openvm_stark_sdk::bench::run_with_metric_collection;
+use openvm_stark_sdk::{
+    bench::run_with_metric_collection, config::baby_bear_poseidon2::BabyBearPoseidon2Engine,
+};
 use openvm_transpiler::FromElf;
 
 fn main() -> Result<()> {
@@ -78,8 +80,12 @@ fn main() -> Result<()> {
     )?;
 
     run_with_metric_collection("OUTPUT_PATH", || -> Result<()> {
-        let mut prover =
-            ContinuationProver::new(&halo2_params_reader, app_pk, app_committed_exe, full_agg_pk);
+        let mut prover = ContinuationProver::<_, BabyBearPoseidon2Engine>::new(
+            &halo2_params_reader,
+            app_pk,
+            app_committed_exe,
+            full_agg_pk,
+        );
         prover.set_program_name("kitchen_sink");
         let stdin = StdIn::default();
         let _proof = prover.generate_proof_for_evm(stdin);
diff --git a/benchmarks/src/bin/verify_fibair.rs b/benchmarks/src/bin/verify_fibair.rs
@@ -59,8 +59,8 @@ fn main() -> Result<()> {
         let app_pk = sdk.app_keygen(app_config)?;
         let app_vk = app_pk.get_app_vk();
         let committed_exe = sdk.commit_app_exe(app_fri_params, program.into())?;
-        let prover =
-            AppProver::new(app_pk.app_vm_pk, committed_exe).with_program_name("verify_fibair");
+        let prover = AppProver::<_, BabyBearPoseidon2Engine>::new(app_pk.app_vm_pk, committed_exe)
+            .with_program_name("verify_fibair");
         let proof = prover.generate_app_proof_without_continuations(input_stream.into());
         sdk.verify_app_proof_without_continuations(&app_vk, &proof)?;
         Ok(())
diff --git a/benchmarks/src/utils.rs b/benchmarks/src/utils.rs
@@ -24,6 +24,7 @@ use openvm_stark_sdk::{
         baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine},
         FriParameters,
     },
+    engine::StarkFriEngine,
     openvm_stark_backend::Chip,
     p3_baby_bear::BabyBear,
 };
@@ -157,7 +158,7 @@ impl BenchmarkCli {
         VC::Periphery: Chip<SC>,
     {
         let app_config = self.app_config(vm_config);
-        bench_from_exe(
+        bench_from_exe::<VC, BabyBearPoseidon2Engine>(
             bench_name,
             app_config,
             exe,
@@ -202,7 +203,7 @@ pub fn build_bench(manifest_dir: PathBuf, profile: impl ToString) -> Result<Elf>
 /// 6. Verify STARK proofs.
 ///
 /// Returns the data necessary for proof aggregation.
-pub fn bench_from_exe<VC>(
+pub fn bench_from_exe<VC, E: StarkFriEngine<SC>>(
     bench_name: impl ToString,
     app_config: AppConfig<VC>,
     exe: impl Into<VmExe<F>>,
@@ -231,18 +232,17 @@ where
     // 4. Generate trace
     // 5. Generate STARK proofs for each segment (segmentation is determined by `config`), with timer.
     let app_vk = app_pk.get_app_vk();
-    let prover = AppProver::new(app_pk.app_vm_pk, committed_exe).with_program_name(bench_name);
+    let prover =
+        AppProver::<VC, E>::new(app_pk.app_vm_pk, committed_exe).with_program_name(bench_name);
     let app_proof = prover.generate_app_proof(input_stream);
     // 6. Verify STARK proofs, including boundary conditions.
     let sdk = Sdk::new();
     sdk.verify_app_proof(&app_vk, &app_proof)
         .expect("Verification failed");
     if let Some(leaf_vm_config) = leaf_vm_config {
         let leaf_vm_pk = leaf_keygen(app_config.leaf_fri_params.fri_params, leaf_vm_config);
-        let leaf_prover = VmLocalProver::<SC, NativeConfig, BabyBearPoseidon2Engine>::new(
-            leaf_vm_pk,
-            app_pk.leaf_committed_exe,
-        );
+        let leaf_prover =
+            VmLocalProver::<SC, NativeConfig, E>::new(leaf_vm_pk, app_pk.leaf_committed_exe);
         let leaf_controller = LeafProvingController {
             num_children: DEFAULT_NUM_CHILDREN_LEAF,
         };
diff --git a/crates/sdk/examples/sdk_app.rs b/crates/sdk/examples/sdk_app.rs
@@ -9,7 +9,7 @@ use openvm_sdk::{
     prover::AppProver,
     Sdk, StdIn,
 };
-use openvm_stark_sdk::config::FriParameters;
+use openvm_stark_sdk::config::{baby_bear_poseidon2::BabyBearPoseidon2Engine, FriParameters};
 use openvm_transpiler::elf::Elf;
 use serde::{Deserialize, Serialize};
 
@@ -90,8 +90,11 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
     // 9a. Generate a proof
     let proof = sdk.generate_app_proof(app_pk.clone(), app_committed_exe.clone(), stdin.clone())?;
     // 9b. Generate a proof with an AppProver with custom fields
-    let app_prover = AppProver::new(app_pk.app_vm_pk.clone(), app_committed_exe.clone())
-        .with_program_name("test_program");
+    let app_prover = AppProver::<_, BabyBearPoseidon2Engine>::new(
+        app_pk.app_vm_pk.clone(),
+        app_committed_exe.clone(),
+    )
+    .with_program_name("test_program");
     let proof = app_prover.generate_app_proof(stdin.clone());
     // ANCHOR_END: proof_generation
 
diff --git a/crates/sdk/src/lib.rs b/crates/sdk/src/lib.rs
@@ -176,7 +176,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
-        let app_prover = AppProver::new(app_pk.app_vm_pk.clone(), app_committed_exe);
+        let app_prover = AppProver::<VC, E>::new(app_pk.app_vm_pk.clone(), app_committed_exe);
         let proof = app_prover.generate_app_proof(inputs);
         Ok(proof)
     }
@@ -241,7 +241,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
-        let stark_prover = StarkProver::new(app_pk, app_exe, agg_stark_pk);
+        let stark_prover = StarkProver::<VC, E>::new(app_pk, app_exe, agg_stark_pk);
         let proof = stark_prover.generate_root_verifier_input(inputs);
         Ok(proof)
     }
@@ -258,7 +258,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
-        let e2e_prover = ContinuationProver::new(reader, app_pk, app_exe, agg_pk);
+        let e2e_prover = ContinuationProver::<VC, E>::new(reader, app_pk, app_exe, agg_pk);
         let proof = e2e_prover.generate_proof_for_evm(inputs);
         Ok(proof)
     }
diff --git a/crates/sdk/src/prover/agg.rs b/crates/sdk/src/prover/agg.rs
@@ -7,9 +7,7 @@ use openvm_continuations::verifier::{
 };
 use openvm_native_circuit::NativeConfig;
 use openvm_native_recursion::hints::Hintable;
-use openvm_stark_sdk::{
-    config::baby_bear_poseidon2::BabyBearPoseidon2Engine, openvm_stark_backend::proof::Proof,
-};
+use openvm_stark_sdk::{engine::StarkFriEngine, openvm_stark_backend::proof::Proof};
 use tracing::info_span;
 
 use crate::{
@@ -25,11 +23,11 @@ pub const DEFAULT_NUM_CHILDREN_LEAF: usize = 1;
 const DEFAULT_NUM_CHILDREN_INTERNAL: usize = 2;
 const DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS: usize = 4;
 
-pub struct AggStarkProver {
-    leaf_prover: VmLocalProver<SC, NativeConfig, BabyBearPoseidon2Engine>,
+pub struct AggStarkProver<E: StarkFriEngine<SC>> {
+    leaf_prover: VmLocalProver<SC, NativeConfig, E>,
     leaf_controller: LeafProvingController,
 
-    internal_prover: VmLocalProver<SC, NativeConfig, BabyBearPoseidon2Engine>,
+    internal_prover: VmLocalProver<SC, NativeConfig, E>,
     root_prover: RootVerifierLocalProver,
 
     pub num_children_internal: usize,
@@ -41,19 +39,17 @@ pub struct LeafProvingController {
     pub num_children: usize,
 }
 
-impl AggStarkProver {
+impl<E: StarkFriEngine<SC>> AggStarkProver<E> {
     pub fn new(
         agg_stark_pk: AggStarkProvingKey,
         leaf_committed_exe: Arc<NonRootCommittedExe>,
     ) -> Self {
-        let leaf_prover = VmLocalProver::<SC, NativeConfig, BabyBearPoseidon2Engine>::new(
-            agg_stark_pk.leaf_vm_pk,
-            leaf_committed_exe,
-        );
+        let leaf_prover =
+            VmLocalProver::<SC, NativeConfig, E>::new(agg_stark_pk.leaf_vm_pk, leaf_committed_exe);
         let leaf_controller = LeafProvingController {
             num_children: DEFAULT_NUM_CHILDREN_LEAF,
         };
-        let internal_prover = VmLocalProver::<SC, NativeConfig, BabyBearPoseidon2Engine>::new(
+        let internal_prover = VmLocalProver::<SC, NativeConfig, E>::new(
             agg_stark_pk.internal_vm_pk,
             agg_stark_pk.internal_committed_exe,
         );
@@ -184,9 +180,9 @@ impl LeafProvingController {
         self
     }
 
-    pub fn generate_proof(
+    pub fn generate_proof<E: StarkFriEngine<SC>>(
         &self,
-        prover: &VmLocalProver<SC, NativeConfig, BabyBearPoseidon2Engine>,
+        prover: &VmLocalProver<SC, NativeConfig, E>,
         app_proofs: &ContinuationVmProof<SC>,
     ) -> Vec<Proof<SC>> {
         info_span!("agg_layer", group = "leaf").in_scope(|| {
diff --git a/crates/sdk/src/prover/app.rs b/crates/sdk/src/prover/app.rs
@@ -3,7 +3,7 @@ use std::sync::Arc;
 use getset::Getters;
 use openvm_circuit::arch::{ContinuationVmProof, VmConfig};
 use openvm_stark_backend::{proof::Proof, Chip};
-use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine;
+use openvm_stark_sdk::engine::StarkFriEngine;
 use tracing::info_span;
 
 use super::vm::SingleSegmentVmProver;
@@ -13,13 +13,13 @@ use crate::{
 };
 
 #[derive(Getters)]
-pub struct AppProver<VC> {
+pub struct AppProver<VC, E: StarkFriEngine<SC>> {
     pub program_name: Option<String>,
     #[getset(get = "pub")]
-    app_prover: VmLocalProver<SC, VC, BabyBearPoseidon2Engine>,
+    app_prover: VmLocalProver<SC, VC, E>,
 }
 
-impl<VC> AppProver<VC> {
+impl<VC, E: StarkFriEngine<SC>> AppProver<VC, E> {
     pub fn new(
         app_vm_pk: Arc<VmProvingKey<SC, VC>>,
         app_committed_exe: Arc<NonRootCommittedExe>,
@@ -29,10 +29,7 @@ impl<VC> AppProver<VC> {
     {
         Self {
             program_name: None,
-            app_prover: VmLocalProver::<SC, VC, BabyBearPoseidon2Engine>::new(
-                app_vm_pk,
-                app_committed_exe,
-            ),
+            app_prover: VmLocalProver::<SC, VC, E>::new(app_vm_pk, app_committed_exe),
         }
     }
     pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
diff --git a/crates/sdk/src/prover/mod.rs b/crates/sdk/src/prover/mod.rs
@@ -1,7 +1,7 @@
 use std::sync::Arc;
 
 use openvm_circuit::arch::VmConfig;
-use openvm_stark_sdk::openvm_stark_backend::Chip;
+use openvm_stark_sdk::{engine::StarkFriEngine, openvm_stark_backend::Chip};
 
 use crate::{keygen::AppProvingKey, stdin::StdIn, NonRootCommittedExe, F, SC};
 
@@ -24,12 +24,12 @@ pub use stark::*;
 
 use crate::{keygen::AggProvingKey, prover::halo2::Halo2Prover, types::EvmProof};
 
-pub struct ContinuationProver<VC> {
-    stark_prover: StarkProver<VC>,
+pub struct ContinuationProver<VC, E: StarkFriEngine<SC>> {
+    stark_prover: StarkProver<VC, E>,
     halo2_prover: Halo2Prover,
 }
 
-impl<VC> ContinuationProver<VC> {
+impl<VC, E: StarkFriEngine<SC>> ContinuationProver<VC, E> {
     pub fn new(
         reader: &impl Halo2ParamsReader,
         app_pk: Arc<AppProvingKey<VC>>,
diff --git a/crates/sdk/src/prover/stark.rs b/crates/sdk/src/prover/stark.rs
@@ -3,18 +3,19 @@ use std::sync::Arc;
 use openvm_circuit::arch::VmConfig;
 use openvm_continuations::verifier::root::types::RootVmVerifierInput;
 use openvm_stark_backend::{proof::Proof, Chip};
+use openvm_stark_sdk::engine::StarkFriEngine;
 
 use crate::{
     keygen::{AggStarkProvingKey, AppProvingKey},
     prover::{agg::AggStarkProver, app::AppProver},
     NonRootCommittedExe, RootSC, StdIn, F, SC,
 };
 
-pub struct StarkProver<VC> {
-    app_prover: AppProver<VC>,
-    agg_prover: AggStarkProver,
+pub struct StarkProver<VC, E: StarkFriEngine<SC>> {
+    app_prover: AppProver<VC, E>,
+    agg_prover: AggStarkProver<E>,
 }
-impl<VC> StarkProver<VC> {
+impl<VC, E: StarkFriEngine<SC>> StarkProver<VC, E> {
     pub fn new(
         app_pk: Arc<AppProvingKey<VC>>,
         app_committed_exe: Arc<NonRootCommittedExe>,

Original file line number	Diff line number	Diff line change
`@@ -176,7 +176,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {`
`176`	`176`	`VC::Executor: Chip<SC>,`
`177`	`177`	`VC::Periphery: Chip<SC>,`
`178`	`178`	`{`
`179`		`- let app_prover = AppProver::new(app_pk.app_vm_pk.clone(), app_committed_exe);`
	`179`	`+ let app_prover = AppProver::<VC, E>::new(app_pk.app_vm_pk.clone(), app_committed_exe);`
`180`	`180`	`let proof = app_prover.generate_app_proof(inputs);`
`181`	`181`	`Ok(proof)`
`182`	`182`	`}`
`@@ -241,7 +241,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {`
`241`	`241`	`VC::Executor: Chip<SC>,`
`242`	`242`	`VC::Periphery: Chip<SC>,`
`243`	`243`	`{`
`244`		`- let stark_prover = StarkProver::new(app_pk, app_exe, agg_stark_pk);`
	`244`	`+ let stark_prover = StarkProver::<VC, E>::new(app_pk, app_exe, agg_stark_pk);`
`245`	`245`	`let proof = stark_prover.generate_root_verifier_input(inputs);`
`246`	`246`	`Ok(proof)`
`247`	`247`	`}`
`@@ -258,7 +258,7 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {`
`258`	`258`	`VC::Executor: Chip<SC>,`
`259`	`259`	`VC::Periphery: Chip<SC>,`
`260`	`260`	`{`
`261`		`- let e2e_prover = ContinuationProver::new(reader, app_pk, app_exe, agg_pk);`
	`261`	`+ let e2e_prover = ContinuationProver::<VC, E>::new(reader, app_pk, app_exe, agg_pk);`
`262`	`262`	`let proof = e2e_prover.generate_proof_for_evm(inputs);`
`263`	`263`	`Ok(proof)`
`264`	`264`	`}`