chore: some more test fixes + clippy

Author: stephenh-axiom-xyz

crates/continuations/src/tests/e2e.rs

@@ -48,7 +48,7 @@ use super::{
 };
 use crate::{
     circuit::{
-        deferral::{verify::DeferralMerkleProofs, DeferralCircuitPvs, DEF_HOOK_PVS_AIR_ID},
+        deferral::{DeferralCircuitPvs, DeferralMerkleProofs, DEF_HOOK_PVS_AIR_ID},
         inner::ProofsType,
     },
     prover::{
@@ -74,7 +74,7 @@ const INPUT_RAW_0: [u8; 8] = [1, 2, 3, 4, 5, 6, 7, 8];
 const INPUT_RAW_1: [u8; 8] = [8, 7, 6, 5, 4, 3, 2, 1];
 const INPUT_RAW_2: [u8; 8] = [9, 9, 9, 9, 9, 9, 9, 9];
 
-fn make_deferral_extension() -> DeferralExtension {
+fn make_deferral_extension(commits: Vec<[u8; 32]>) -> DeferralExtension {
     let fns: Vec<_> = (0..NUM_DEF_CIRCUITS)
         .map(|idx| {
             Arc::new(DeferralFn::new(move |input_raw| {
@@ -89,7 +89,7 @@ fn make_deferral_extension() -> DeferralExtension {
             }))
         })
         .collect();
-    DeferralExtension::new(fns)
+    DeferralExtension::new(fns, commits)
 }
 
 fn deferral_fn_output(idx: u16, input_raw: &[u8]) -> Vec<u8> {
@@ -230,17 +230,14 @@ fn test_deferral_e2e() -> Result<()> {
     let (_, def_circuit_vk) = gpu_engine.keygen(&[empty_air]);
     let def_circuit_vk = Arc::new(def_circuit_vk);
 
-    let def_leaf_prover = DeferralInnerProver::<2>::new::<GpuEngine>(
-        def_circuit_vk.clone(),
-        leaf_system_params(),
-        false,
-    );
-    let def_i0_prover = DeferralInnerProver::<2>::new::<GpuEngine>(
+    let def_leaf_prover =
+        DeferralInnerProver::new::<GpuEngine>(def_circuit_vk.clone(), leaf_system_params(), false);
+    let def_i0_prover = DeferralInnerProver::new::<GpuEngine>(
         def_leaf_prover.get_vk(),
         internal_system_params(),
         false,
     );
-    let def_i1_prover = DeferralInnerProver::<2>::new::<GpuEngine>(
+    let def_i1_prover = DeferralInnerProver::new::<GpuEngine>(
         def_i0_prover.get_vk(),
         internal_system_params(),
         true,
@@ -288,7 +285,7 @@ fn test_deferral_e2e() -> Result<()> {
         rv32i: Rv32I,
         rv32m: Rv32M::default(),
         io: Rv32Io,
-        deferral: make_deferral_extension(),
+        deferral: make_deferral_extension(transpiler_commits.clone()),
     };
 
     let elf = Elf::decode(
@@ -408,7 +405,7 @@ fn test_deferral_e2e() -> Result<()> {
         proofs: Vec<Proof<SC>>,
         leaf_children: Vec<([F; DIGEST_SIZE], [F; DIGEST_SIZE])>,
     ) -> Result<Proof<SC>> {
-        let leaf_prover = DeferralInnerProver::<2>::new::<GpuEngine>(
+        let leaf_prover = DeferralInnerProver::new::<GpuEngine>(
             def_circuit_vk.clone(),
             leaf_system_params(),
             false,
@@ -431,7 +428,7 @@ fn test_deferral_e2e() -> Result<()> {
         current_proofs = next;
         child_merkle_depth += 1;
 
-        let i4l_prover = DeferralInnerProver::<2>::new::<GpuEngine>(
+        let i4l_prover = DeferralInnerProver::new::<GpuEngine>(
             leaf_prover.get_vk(),
             internal_system_params(),
             false,
@@ -453,7 +450,7 @@ fn test_deferral_e2e() -> Result<()> {
         current_proofs = next;
         child_merkle_depth += 1;
 
-        let ir_prover = DeferralInnerProver::<2>::new::<GpuEngine>(
+        let ir_prover = DeferralInnerProver::new::<GpuEngine>(
             i4l_prover.get_vk(),
             internal_system_params(),
             true,

crates/continuations/src/tests/mod.rs

@@ -360,16 +360,9 @@ fn test_root_prover_trace_heights() -> Result<()> {
 // DEFERRAL TESTS
 ///////////////////////////////////////////////////////////////////////////////
 
-#[cfg(feature = "cuda")]
-fn generate_single_def_proof(
-    _child_extra_recursive_layers: usize,
-) -> Result<(Arc<MultiStarkVerifyingKey<SC>>, Proof<SC>)> {
-    dummy::generate_single_dummy_def_proof()
-}
-
 #[cfg(feature = "cuda")]
 fn aggregate_deferral_layer(
-    prover: &DeferralInnerProver<2>,
+    prover: &DeferralInnerProver,
     proofs: &[Proof<SC>],
     child_is_def: bool,
     child_merkle_depth: usize,
@@ -412,17 +405,13 @@ pub(in crate::tests) fn generate_deferral_internal_recursive_proof_from_copies(
     let mut current_proofs = vec![def_proof.clone(); num_copies];
     let mut child_merkle_depth = 0usize;
 
-    let leaf_prover =
-        DeferralInnerProver::<2>::new::<Engine>(deferral_vk, leaf_system_params(), false);
+    let leaf_prover = DeferralInnerProver::new::<Engine>(deferral_vk, leaf_system_params(), false);
     current_proofs =
         aggregate_deferral_layer(&leaf_prover, &current_proofs, false, child_merkle_depth)?;
     child_merkle_depth += 1;
 
-    let internal_for_leaf_prover = DeferralInnerProver::<2>::new::<Engine>(
-        leaf_prover.get_vk(),
-        internal_system_params(),
-        false,
-    );
+    let internal_for_leaf_prover =
+        DeferralInnerProver::new::<Engine>(leaf_prover.get_vk(), internal_system_params(), false);
     current_proofs = aggregate_deferral_layer(
         &internal_for_leaf_prover,
         &current_proofs,
@@ -433,7 +422,7 @@ pub(in crate::tests) fn generate_deferral_internal_recursive_proof_from_copies(
 
     let child_vk = internal_for_leaf_prover.get_vk();
     let internal_recursive_prover =
-        DeferralInnerProver::<2>::new::<Engine>(child_vk, internal_system_params(), true);
+        DeferralInnerProver::new::<Engine>(child_vk, internal_system_params(), true);
     loop {
         current_proofs = aggregate_deferral_layer(
             &internal_recursive_prover,
@@ -506,19 +495,14 @@ fn expected_deferral_inner_merkle_commit_from_copies(
 }
 
 #[cfg(feature = "cuda")]
-#[test_case(0, 1 ; "internal_recursive_dag_commit not set + one child")]
-#[test_case(1, 1 ; "internal_recursive_dag_commit set + one child")]
-#[test_case(0, 2 ; "internal_recursive_dag_commit not set + two children")]
-#[test_case(1, 2 ; "internal_recursive_dag_commit set + two children")]
-fn test_deferral_leaf_prover(
-    child_extra_recursive_layers: usize,
-    num_children: usize,
-) -> Result<()> {
+#[test_case(1)]
+#[test_case(2)]
+fn test_deferral_leaf_prover(num_children: usize) -> Result<()> {
     setup_tracing_with_log_level(Level::INFO);
-    let (deferral_vk, def_proof) = generate_single_def_proof(child_extra_recursive_layers)?;
+    let (deferral_vk, def_proof) = dummy::generate_single_dummy_def_proof()?;
 
     let deferral_inner_prover =
-        DeferralInnerProver::<2>::new::<Engine>(deferral_vk, leaf_system_params(), false);
+        DeferralInnerProver::new::<Engine>(deferral_vk, leaf_system_params(), false);
     let wrapped_proof = deferral_inner_prover.agg_prove::<Engine>(
         &vec![def_proof.clone(); num_children],
         DeferralChildVkKind::DeferralCircuit,
@@ -549,7 +533,7 @@ fn test_deferral_leaf_prover(
 #[test_case(5 ; "partially empty aggregation tree")]
 fn test_deferral_aggregation(num_children: usize) -> Result<()> {
     setup_tracing_with_log_level(Level::INFO);
-    let (deferral_vk, def_proof) = generate_single_def_proof(0)?;
+    let (deferral_vk, def_proof) = dummy::generate_single_dummy_def_proof()?;
     let (vk, final_proof) = generate_deferral_internal_recursive_proof_from_copies(
         deferral_vk,
         def_proof.clone(),
@@ -575,21 +559,17 @@ fn test_deferral_aggregation(num_children: usize) -> Result<()> {
 #[test]
 fn test_deferral_internal_recursive_vk_stabilization() -> Result<()> {
     setup_tracing_with_log_level(Level::INFO);
-    let (deferral_vk, _) = generate_single_def_proof(0)?;
+    let (deferral_vk, _) = dummy::generate_single_dummy_def_proof()?;
 
-    let leaf_prover =
-        DeferralInnerProver::<2>::new::<Engine>(deferral_vk, leaf_system_params(), false);
-    let internal_0_prover = DeferralInnerProver::<2>::new::<Engine>(
-        leaf_prover.get_vk(),
-        internal_system_params(),
-        false,
-    );
-    let internal_1_prover = DeferralInnerProver::<2>::new::<Engine>(
+    let leaf_prover = DeferralInnerProver::new::<Engine>(deferral_vk, leaf_system_params(), false);
+    let internal_0_prover =
+        DeferralInnerProver::new::<Engine>(leaf_prover.get_vk(), internal_system_params(), false);
+    let internal_1_prover = DeferralInnerProver::new::<Engine>(
         internal_0_prover.get_vk(),
         internal_system_params(),
         false,
     );
-    let test_prover = DeferralInnerProver::<2>::new::<Engine>(
+    let test_prover = DeferralInnerProver::new::<Engine>(
         internal_1_prover.get_vk(),
         internal_system_params(),
         true,
@@ -610,7 +590,7 @@ fn test_deferral_hook_prover(num_children: usize) -> Result<()> {
     use crate::circuit::utils::vk_commit_components;
 
     setup_tracing_with_log_level(Level::INFO);
-    let (deferral_vk, def_proof) = generate_single_def_proof(0)?;
+    let (deferral_vk, def_proof) = dummy::generate_single_dummy_def_proof()?;
     let (deferral_internal_recursive_vk, final_inner_proof) =
         generate_deferral_internal_recursive_proof_from_copies(
             deferral_vk,

crates/sdk/src/prover/deferral/circuit.rs

@@ -31,6 +31,11 @@ pub struct SingleDefCircuitProver {
     pub internal_for_leaf_prover: DeferralInnerProver,
 }
 
+pub struct SingleDefCircuitResult {
+    pub internal_for_leaf_proofs: Vec<Proof<SC>>,
+    pub leaf_io_commits: Vec<DeferralIoCommit<F>>,
+}
+
 impl SingleDefCircuitProver {
     pub fn new<DP: DeferralCircuitProver<SC> + 'static>(
         def_circuit_prover: DP,
@@ -64,19 +69,16 @@ impl SingleDefCircuitProver {
         }
     }
 
-    pub fn prove(
-        &self,
-        inputs: &DeferralInput,
-    ) -> Result<(Vec<Proof<SC>>, Vec<DeferralIoCommit<F>>)> {
+    pub fn prove(&self, inputs: &DeferralInput) -> Result<SingleDefCircuitResult> {
         // Generate deferral circuit proofs
         let def_proofs = inputs
             .byte_vec
             .iter()
-            .map(|input| self.def_circuit_prover.prove(&input))
+            .map(|input| self.def_circuit_prover.prove(input))
             .collect_vec();
 
         // Extract leaf IO commits from the deferral circuit proofs
-        let leaf_children = def_proofs
+        let leaf_io_commits = def_proofs
             .iter()
             .map(|proof| {
                 let pvs: &DeferralCircuitPvs<F> = proof.public_values[DEF_CIRCUIT_PVS_AIR_ID]
@@ -116,8 +118,8 @@ impl SingleDefCircuitProver {
         // Verify leaf-layer proofs and generate internal-for-leaf-layer proofs
         let mut internal_node_idx = 0u32;
         let child_merkle_depth = (leaf_proofs.len() != 1).then_some(1);
-        let internal_proofs =
-            info_span!("agg_layer", group = "internal_for_leaf").in_scope(|| {
+        let internal_for_leaf_proofs = info_span!("agg_layer", group = "internal_for_leaf")
+            .in_scope(|| {
                 leaf_proofs
                     .chunks(2)
                     .map(|proofs| {
@@ -135,7 +137,10 @@ impl SingleDefCircuitProver {
                     .collect::<Result<Vec<_>>>()
             })?;
 
-        Ok((internal_proofs, leaf_children))
+        Ok(SingleDefCircuitResult {
+            internal_for_leaf_proofs,
+            leaf_io_commits,
+        })
     }
 
     pub fn vk_commit(&self, internal_for_leaf_dag_commit: DagCommit<F>) -> Digest {

crates/sdk/src/prover/deferral/mod.rs

@@ -3,7 +3,6 @@ use std::sync::Arc;
 use eyre::Result;
 use openvm_continuations::{
     bn254::CommitBytes,
-    circuit::deferral::hook::DeferralIoCommit,
     prover::{DeferralChildVkKind, DeferralCircuitProver},
     SC,
 };
@@ -40,6 +39,7 @@ pub use merkle::*;
 
 pub type DefAggProvingKey = AggProvingKey;
 
+#[allow(clippy::large_enum_variant)]
 pub enum DeferralProof {
     Present(Proof<SC>),
     Absent(DeferralPvs<F>),
@@ -51,6 +51,9 @@ pub struct DeferralProver {
     pub def_hook_prover: DeferralHookProver,
 }
 
+unsafe impl Send for DeferralProver {}
+unsafe impl Sync for DeferralProver {}
+
 impl DeferralProver {
     pub fn new<DP: DeferralCircuitProver<SC> + 'static>(
         def_circuit_prover: DP,
@@ -123,7 +126,7 @@ impl DeferralProver {
 
     pub fn prove(&self, inputs: &[DeferralInput]) -> Result<Vec<DeferralProof>> {
         // Generate internal-for-leaf proofs and leaf IO commits per circuit
-        let per_circuit: Vec<(Vec<Proof<SC>>, Vec<DeferralIoCommit<F>>)> = self
+        let per_circuit = self
             .single_circuit_provers
             .iter()
             .zip(inputs)
@@ -134,7 +137,8 @@ impl DeferralProver {
         per_circuit
             .into_iter()
             .enumerate()
-            .map(|(circuit_idx, (mut proofs, leaf_children))| {
+            .map(|(circuit_idx, res)| {
+                let mut proofs = res.internal_for_leaf_proofs;
                 if proofs.is_empty() {
                     let def_vk_commit = self.single_circuit_provers[circuit_idx]
                         .vk_commit(self.internal_recursive_prover.get_dag_commit(false));
@@ -199,7 +203,7 @@ impl DeferralProver {
                     // Generate the deferral hook proof
                     Ok(DeferralProof::Present(
                         self.def_hook_prover
-                            .prove::<E>(proofs.pop().unwrap(), leaf_children)?,
+                            .prove::<E>(proofs.pop().unwrap(), res.leaf_io_commits)?,
                     ))
                 }
             })

crates/sdk/src/stdin.rs

@@ -80,7 +80,7 @@ impl DeferralInput {
     pub fn into_inputs<I: Decode>(self) -> Vec<I> {
         self.byte_vec
             .iter()
-            .map(|input| I::decode_from_bytes(&input).unwrap())
+            .map(|input| I::decode_from_bytes(input).unwrap())
             .collect_vec()
     }
 

crates/sdk/src/tests.rs

@@ -10,9 +10,8 @@ use openvm_stark_sdk::config::{
     root_params_with_100_bits_security,
 };
 use openvm_transpiler::elf::Elf;
-use openvm_verify_stark_circuit::{
-    extension::{get_deferral_state, get_raw_deferral_results, verify_stark_deferral_fn},
-    prover::{DefaultEngine, DeferredVerifyDefaultCircuitProver, DeferredVerifyDefaultProver},
+use openvm_verify_stark_circuit::extension::{
+    get_deferral_state, get_raw_deferral_results, verify_stark_deferral_fn,
 };
 use openvm_verify_stark_host::vk::NonRootStarkVerifyingKey;
 
@@ -22,7 +21,19 @@ use crate::{
     CpuSdk, DeferralInput, Sdk, StdIn,
 };
 
-type E = openvm_stark_sdk::config::baby_bear_bn254_poseidon2::BabyBearBn254Poseidon2CpuEngine;
+cfg_if::cfg_if! {
+    if #[cfg(feature = "cuda")] {
+        use openvm_verify_stark_circuit::prover::DeferredVerifyGpuProver as VerifyProver;
+        use openvm_verify_stark_circuit::prover::DeferredVerifyGpuCircuitProver as VerifyCircuitProver;
+        type E = openvm_cuda_backend::BabyBearPoseidon2GpuEngine;
+        type RootE = openvm_stark_sdk::config::baby_bear_bn254_poseidon2::BabyBearBn254Poseidon2CpuEngine;
+    } else {
+        use openvm_verify_stark_circuit::prover::DeferredVerifyCpuProver as VerifyProver;
+        use openvm_verify_stark_circuit::prover::DeferredVerifyCpuCircuitProver as VerifyCircuitProver;
+        type E = openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2CpuEngine;
+        type RootE = openvm_stark_sdk::config::baby_bear_bn254_poseidon2::BabyBearBn254Poseidon2CpuEngine;
+    }
+}
 
 #[test]
 fn test_root_prover_trace_heights() -> Result<()> {
@@ -45,7 +56,7 @@ fn test_root_prover_trace_heights() -> Result<()> {
 
     let proof = evm_prover.prove(stdin)?;
     let vk = evm_prover.root_prover.0.get_vk();
-    let engine = E::new(vk.inner.params.clone());
+    let engine = RootE::new(vk.inner.params.clone());
     engine.verify(&vk, &proof)?;
 
     Ok(())
@@ -85,15 +96,15 @@ fn test_verify_stark_deferral() -> Result<()> {
     let num_user_pvs = fib_system_config.num_public_values;
 
     let def_circuit_params = internal_params_with_100_bits_security();
-    let deferred_verify_prover = DeferredVerifyDefaultProver::new::<DefaultEngine>(
+    let deferred_verify_prover = VerifyProver::new::<E>(
         ir_vk,
         ir_pcs_data,
         def_circuit_params,
         memory_dimensions,
         num_user_pvs,
         None,
     );
-    let verify_stark_prover = DeferredVerifyDefaultCircuitProver::new(deferred_verify_prover);
+    let verify_stark_prover = VerifyCircuitProver::new(deferred_verify_prover);
 
     // ---- Step 3: Create DeferralProver ----
     let hook_params = root_params_with_100_bits_security();