feat: Add e2e stark proof support

crates/cli/src/commands/prove.rs

@@ -2,18 +2,18 @@ use std::{path::PathBuf, sync::Arc};
 
 use clap::Parser;
 use eyre::Result;
+#[cfg(feature = "evm-prove")]
+use openvm_sdk::fs::write_evm_proof_to_file;
 use openvm_sdk::{
     commit::AppExecutionCommit,
-    config::SdkVmConfig,
-    fs::{read_app_pk_from_file, read_exe_from_file, write_app_proof_to_file},
+    config::{AggregationTreeConfig, SdkVmConfig},
+    fs::{
+        encode_to_file, read_agg_pk_from_file, read_app_pk_from_file, read_exe_from_file,
+        write_app_proof_to_file,
+    },
     keygen::AppProvingKey,
     NonRootCommittedExe, Sdk, StdIn,
 };
-#[cfg(feature = "evm-prove")]
-use openvm_sdk::{
-    config::AggregationTreeConfig,
-    fs::{read_agg_pk_from_file, write_evm_proof_to_file},
-};
 
 use crate::{
     default::*,
@@ -42,6 +42,22 @@ enum ProveSubCommand {
         #[arg(long, action, help = "Path to output proof", default_value = DEFAULT_APP_PROOF_PATH)]
         output: PathBuf,
     },
+    Stark {
+        #[arg(long, action, help = "Path to app proving key", default_value = DEFAULT_APP_PK_PATH)]
+        app_pk: PathBuf,
+
+        #[arg(long, action, help = "Path to OpenVM executable", default_value = DEFAULT_APP_EXE_PATH)]
+        exe: PathBuf,
+
+        #[arg(long, value_parser, help = "Input to OpenVM program")]
+        input: Option<Input>,
+
+        #[arg(long, action, help = "Path to output proof", default_value = DEFAULT_APP_PROOF_PATH)]
+        output: PathBuf,
+
+        #[command(flatten)]
+        agg_tree_config: AggregationTreeConfig,
+    },
     #[cfg(feature = "evm-prove")]
     Evm {
         #[arg(long, action, help = "Path to app proving key", default_value = DEFAULT_APP_PK_PATH)]
@@ -76,6 +92,28 @@ impl ProveCmd {
                 let app_proof = sdk.generate_app_proof(app_pk, committed_exe, input)?;
                 write_app_proof_to_file(app_proof, output)?;
             }
+            ProveSubCommand::Stark {
+                app_pk,
+                exe,
+                input,
+                output,
+                agg_tree_config,
+            } => {
+                let mut sdk = sdk;
+                sdk.set_agg_tree_config(*agg_tree_config);
+                let (app_pk, committed_exe, input) =
+                    Self::prepare_execution(&sdk, app_pk, exe, input)?;
+                let agg_pk = read_agg_pk_from_file(DEFAULT_AGG_PK_PATH).map_err(|e| {
+                    eyre::eyre!("Failed to read aggregation proving key: {}\nPlease run 'cargo openvm setup' first", e)
+                })?;
+                let stark_proof = sdk.generate_e2e_stark_proof(
+                    app_pk,
+                    committed_exe,
+                    agg_pk.agg_stark_pk,
+                    input,
+                )?;
+                encode_to_file(output, stark_proof)?;
+            }
             #[cfg(feature = "evm-prove")]
             ProveSubCommand::Evm {
                 app_pk,

crates/sdk/src/codec.rs

@@ -16,7 +16,8 @@ use openvm_stark_backend::{
 };
 use p3_fri::CommitPhaseProofStep;
 
-use super::{F, SC}; // BabyBearPoseidon2Config
+use super::{F, SC};
+use crate::types::E2eStarkProof; // BabyBearPoseidon2Config
 
 type Challenge = BinomialExtensionField<F, 4>;
 
@@ -59,6 +60,13 @@ impl Encode for ContinuationVmProof<SC> {
     }
 }
 
+impl Encode for E2eStarkProof {
+    fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
+        self.proof.encode(writer)?;
+        encode_slice(&self.user_public_values, writer)
+    }
+}
+
 impl Encode for UserPublicValuesProof<DIGEST_SIZE, F> {
     fn encode<W: Write>(&self, writer: &mut W) -> Result<()> {
         encode_slice(&self.proof, writer)?;
@@ -323,6 +331,17 @@ impl Decode for ContinuationVmProof<SC> {
     }
 }
 
+impl Decode for E2eStarkProof {
+    fn decode<R: Read>(reader: &mut R) -> Result<Self> {
+        let proof = Proof::decode(reader)?;
+        let user_public_values = decode_vec(reader)?;
+        Ok(Self {
+            proof,
+            user_public_values,
+        })
+    }
+}
+
 impl Decode for UserPublicValuesProof<DIGEST_SIZE, F> {
     fn decode<R: Read>(reader: &mut R) -> Result<Self> {
         let proof = decode_vec(reader)?;

crates/sdk/src/lib.rs

@@ -58,6 +58,8 @@ pub mod prover;
 mod stdin;
 pub use stdin::*;
 
+use crate::types::E2eStarkProof;
+
 pub mod fs;
 pub mod types;
 
@@ -268,6 +270,23 @@ impl<E: StarkFriEngine<SC>> GenericSdk<E> {
         Ok(proof)
     }
 
+    pub fn generate_e2e_stark_proof<VC: VmConfig<F>>(
+        &self,
+        app_pk: Arc<AppProvingKey<VC>>,
+        app_exe: Arc<NonRootCommittedExe>,
+        agg_stark_pk: AggStarkProvingKey,
+        inputs: StdIn,
+    ) -> Result<E2eStarkProof>
+    where
+        VC::Executor: Chip<SC>,
+        VC::Periphery: Chip<SC>,
+    {
+        let stark_prover =
+            StarkProver::<VC, E>::new(app_pk, app_exe, agg_stark_pk, self.agg_tree_config);
+        let proof = stark_prover.generate_e2e_stark_proof(inputs);
+        Ok(proof)
+    }
+
     #[cfg(feature = "evm-prove")]
     pub fn generate_evm_proof<VC: VmConfig<F>>(
         &self,

crates/sdk/src/prover/agg.rs

@@ -6,6 +6,7 @@ use openvm_continuations::verifier::{
     root::types::RootVmVerifierInput,
 };
 use openvm_native_circuit::NativeConfig;
+use openvm_native_compiler::ir::DIGEST_SIZE;
 use openvm_native_recursion::hints::Hintable;
 use openvm_stark_sdk::{engine::StarkFriEngine, openvm_stark_backend::proof::Proof};
 use tracing::info_span;
@@ -17,6 +18,7 @@ use crate::{
         vm::{local::VmLocalProver, SingleSegmentVmProver},
         RootVerifierLocalProver,
     },
+    types::E2eStarkProof,
     NonRootCommittedExe, RootSC, F, SC,
 };
 
@@ -83,50 +85,30 @@ impl<E: StarkFriEngine<SC>> AggStarkProver<E> {
         self.generate_root_proof_impl(root_verifier_input)
     }
 
+    pub fn generate_leaf_proofs(&self, app_proofs: &ContinuationVmProof<SC>) -> Vec<Proof<SC>> {
+        self.leaf_controller
+            .generate_proof(&self.leaf_prover, app_proofs)
+    }
+
     pub fn generate_root_verifier_input(
         &self,
         app_proofs: ContinuationVmProof<SC>,
     ) -> RootVmVerifierInput<SC> {
-        let leaf_proofs = self
-            .leaf_controller
-            .generate_proof(&self.leaf_prover, &app_proofs);
+        let leaf_proofs = self.generate_leaf_proofs(&app_proofs);
         let public_values = app_proofs.user_public_values.public_values;
-        let internal_proof = self.generate_internal_proof_impl(leaf_proofs, &public_values);
-        RootVmVerifierInput {
-            proofs: vec![internal_proof],
-            public_values,
-        }
+        let e2e_stark_proof = self.generate_e2e_stark_proof(leaf_proofs, public_values);
+        self.wrap_e2e_stark_proof(e2e_stark_proof)
     }
 
-    fn generate_internal_proof_impl(
+    pub fn generate_e2e_stark_proof(
         &self,
         leaf_proofs: Vec<Proof<SC>>,
-        public_values: &[F],
-    ) -> Proof<SC> {
+        public_values: Vec<F>,
+    ) -> E2eStarkProof {
         let mut internal_node_idx = -1;
         let mut internal_node_height = 0;
         let mut proofs = leaf_proofs;
-        let mut wrapper_layers = 0;
-        loop {
-            if proofs.len() == 1 {
-                let actual_air_heights =
-                    self.root_prover
-                        .execute_for_air_heights(RootVmVerifierInput {
-                            proofs: vec![proofs[0].clone()],
-                            public_values: public_values.to_vec(),
-                        });
-                // Root verifier can handle the internal proof. We can stop here.
-                if heights_le(
-                    &actual_air_heights,
-                    &self.root_prover.root_verifier_pk.air_heights,
-                ) {
-                    break;
-                }
-                if wrapper_layers >= self.max_internal_wrapper_layers {
-                    panic!("The heights of the root verifier still exceed the required heights after {} wrapper layers", self.max_internal_wrapper_layers);
-                }
-                wrapper_layers += 1;
-            }
+        while proofs.len() > 1 {
             let internal_inputs = InternalVmVerifierInput::chunk_leaf_or_internal_proofs(
                 self.internal_prover
                     .committed_exe
@@ -158,7 +140,26 @@ impl<E: StarkFriEngine<SC>> AggStarkProver<E> {
             });
             internal_node_height += 1;
         }
-        proofs.pop().unwrap()
+        E2eStarkProof {
+            proof: proofs.pop().unwrap(),
+            user_public_values: public_values,
+        }
+    }
+
+    /// Wrap the e2e stark proof until its heights meet the requirements of the root verifier.
+    pub fn wrap_e2e_stark_proof(&self, e2e_stark_proof: E2eStarkProof) -> RootVmVerifierInput<SC> {
+        let internal_commit = self
+            .internal_prover
+            .committed_exe
+            .get_program_commit()
+            .into();
+        wrap_e2e_stark_proof(
+            &self.internal_prover,
+            &self.root_prover,
+            internal_commit,
+            self.max_internal_wrapper_layers,
+            e2e_stark_proof,
+        )
     }
 
     fn generate_root_proof_impl(&self, root_input: RootVmVerifierInput<SC>) -> Proof<RootSC> {
@@ -204,6 +205,58 @@ impl LeafProvingController {
     }
 }
 
+/// Wrap the e2e stark proof until its heights meet the requirements of the root verifier.
+pub fn wrap_e2e_stark_proof<E: StarkFriEngine<SC>>(
+    internal_prover: &VmLocalProver<SC, NativeConfig, E>,
+    root_prover: &RootVerifierLocalProver,
+    internal_commit: [F; DIGEST_SIZE],
+    max_internal_wrapper_layers: usize,
+    e2e_stark_proof: E2eStarkProof,
+) -> RootVmVerifierInput<SC> {
+    let E2eStarkProof {
+        mut proof,
+        user_public_values,
+    } = e2e_stark_proof;
+    let mut wrapper_layers = 0;
+    loop {
+        let actual_air_heights = root_prover.execute_for_air_heights(RootVmVerifierInput {
+            proofs: vec![proof.clone()],
+            public_values: user_public_values.clone(),
+        });
+        // Root verifier can handle the internal proof. We can stop here.
+        if heights_le(
+            &actual_air_heights,
+            &root_prover.root_verifier_pk.air_heights,
+        ) {
+            break;
+        }
+        if wrapper_layers >= max_internal_wrapper_layers {
+            panic!("The heights of the root verifier still exceed the required heights after {} wrapper layers", max_internal_wrapper_layers);
+        }
+        wrapper_layers += 1;
+        let input = InternalVmVerifierInput {
+            self_program_commit: internal_commit,
+            proofs: vec![proof.clone()],
+        };
+        proof = info_span!(
+            "wrapper_layer",
+            group = format!("internal_wrapper.{wrapper_layers}")
+        )
+        .in_scope(|| {
+            #[cfg(feature = "bench-metrics")]
+            {
+                metrics::counter!("fri.log_blowup")
+                    .absolute(internal_prover.fri_params().log_blowup as u64);
+            }
+            SingleSegmentVmProver::prove(internal_prover, input.write())
+        });
+    }
+    RootVmVerifierInput {
+        proofs: vec![proof],
+        public_values: user_public_values,
+    }
+}
+
 fn heights_le(a: &[usize], b: &[usize]) -> bool {
     assert_eq!(a.len(), b.len());
     a.iter().zip(b.iter()).all(|(a, b)| a <= b)

crates/sdk/src/prover/stark.rs

@@ -9,6 +9,7 @@ use crate::{
     config::AggregationTreeConfig,
     keygen::{AggStarkProvingKey, AppProvingKey},
     prover::{agg::AggStarkProver, app::AppProver},
+    types::E2eStarkProof,
     NonRootCommittedExe, RootSC, StdIn, F, SC,
 };
 
@@ -68,4 +69,16 @@ impl<VC, E: StarkFriEngine<SC>> StarkProver<VC, E> {
         let app_proof = self.app_prover.generate_app_proof(input);
         self.agg_prover.generate_root_verifier_input(app_proof)
     }
+
+    pub fn generate_e2e_stark_proof(&self, input: StdIn) -> E2eStarkProof
+    where
+        VC: VmConfig<F>,
+        VC::Executor: Chip<SC>,
+        VC::Periphery: Chip<SC>,
+    {
+        let app_proof = self.app_prover.generate_app_proof(input);
+        let leaf_proofs = self.agg_prover.generate_leaf_proofs(&app_proof);
+        self.agg_prover
+            .generate_e2e_stark_proof(leaf_proofs, app_proof.user_public_values.public_values)
+    }
 }

crates/sdk/src/types.rs

@@ -1,7 +1,9 @@
 use std::iter::{once, repeat};
 
 use itertools::Itertools;
+use openvm_continuations::{F, SC};
 use openvm_native_recursion::halo2::{wrapper::EvmVerifierByteCode, Fr, RawEvmProof};
+use openvm_stark_backend::proof::Proof;
 use serde::{Deserialize, Serialize};
 use serde_with::serde_as;
 use thiserror::Error;
@@ -49,6 +51,12 @@ pub struct EvmProof {
     pub proof_data: ProofData,
 }
 
+#[derive(Clone, Deserialize, Serialize)]
+pub struct E2eStarkProof {
+    pub proof: Proof<SC>,
+    pub user_public_values: Vec<F>,
+}
+
 #[derive(Debug, Error)]
 pub enum EvmProofConversionError {
     #[error("Invalid length of proof")]