openvm_sdk/prover/
app.rs

1use std::sync::{Arc, OnceLock};
2
3#[cfg(feature = "async")]
4pub use async_prover::*;
5use getset::Getters;
6use itertools::Itertools;
7use openvm_circuit::{
8    arch::{
9        hasher::poseidon2::{vm_poseidon2_hasher, Poseidon2Hasher},
10        instructions::exe::VmExe,
11        verify_segments, ContinuationVmProof, ContinuationVmProver, Executor, MeteredExecutor,
12        PreflightExecutor, VerifiedExecutionPayload, VirtualMachine, VirtualMachineError,
13        VmBuilder, VmExecutionConfig, VmInstance, VmVerificationError,
14    },
15    system::memory::CHUNK,
16};
17use openvm_stark_backend::{
18    config::{Com, Val},
19    keygen::types::MultiStarkVerifyingKey,
20    p3_field::PrimeField32,
21};
22use openvm_stark_sdk::{
23    config::baby_bear_poseidon2::BabyBearPoseidon2Engine,
24    engine::{StarkEngine, StarkFriEngine},
25};
26use tracing::instrument;
27
28use crate::{
29    commit::{AppExecutionCommit, CommitBytes},
30    keygen::AppVerifyingKey,
31    prover::vm::{new_local_prover, types::VmProvingKey},
32    util::check_max_constraint_degrees,
33    StdIn, F, SC,
34};
35
36#[derive(Getters)]
37pub struct AppProver<E, VB>
38where
39    E: StarkEngine,
40    VB: VmBuilder<E>,
41{
42    pub program_name: Option<String>,
43    #[getset(get = "pub")]
44    instance: VmInstance<E, VB>,
45    #[getset(get = "pub")]
46    app_vm_vk: MultiStarkVerifyingKey<E::SC>,
47    #[getset(get = "pub")]
48    leaf_verifier_program_commit: Com<E::SC>,
49
50    app_execution_commit: OnceLock<AppExecutionCommit>,
51}
52
53impl<E, VB> AppProver<E, VB>
54where
55    E: StarkFriEngine,
56    VB: VmBuilder<E>,
57    Val<E::SC>: PrimeField32,
58    Com<E::SC>: AsRef<[Val<E::SC>; CHUNK]> + From<[Val<E::SC>; CHUNK]> + Into<[Val<E::SC>; CHUNK]>,
59{
60    /// Creates a new [AppProver] instance. This method will re-commit the `exe` program on device.
61    /// If a cached version of the program already exists on device, then directly use the
62    /// [`Self::new_from_instance`] constructor.
63    ///
64    /// The `leaf_verifier_program_commit` is the commitment to the program of the leaf verifier
65    /// that verifies the App VM circuit. It can be found in the `AppProvingKey`.
66    pub fn new(
67        vm_builder: VB,
68        app_vm_pk: &VmProvingKey<E::SC, VB::VmConfig>,
69        app_exe: Arc<VmExe<Val<E::SC>>>,
70        leaf_verifier_program_commit: Com<E::SC>,
71    ) -> Result<Self, VirtualMachineError> {
72        let instance = new_local_prover(vm_builder, app_vm_pk, app_exe)?;
73        let app_vm_vk = app_vm_pk.vm_pk.get_vk();
74
75        Ok(Self::new_from_instance(
76            instance,
77            app_vm_vk,
78            leaf_verifier_program_commit,
79        ))
80    }
81
82    pub fn new_from_instance(
83        instance: VmInstance<E, VB>,
84        app_vm_vk: MultiStarkVerifyingKey<E::SC>,
85        leaf_verifier_program_commit: Com<E::SC>,
86    ) -> Self {
87        Self {
88            program_name: None,
89            instance,
90            app_vm_vk,
91            leaf_verifier_program_commit,
92            app_execution_commit: OnceLock::new(),
93        }
94    }
95
96    pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
97        self.program_name = Some(program_name.as_ref().to_string());
98        self
99    }
100    pub fn with_program_name(mut self, program_name: impl AsRef<str>) -> Self {
101        self.set_program_name(program_name);
102        self
103    }
104
105    /// Returns [AppExecutionCommit], which is a commitment to **both** the App VM and the App
106    /// VmExe.
107    pub fn app_commit(&self) -> AppExecutionCommit {
108        *self.app_execution_commit.get_or_init(|| {
109            AppExecutionCommit::compute::<E::SC>(
110                &self.instance().vm.config().as_ref().memory_config,
111                self.instance().exe(),
112                self.instance().program_commitment().clone(),
113                self.leaf_verifier_program_commit.clone(),
114            )
115        })
116    }
117
118    pub fn app_program_commit(&self) -> Com<E::SC> {
119        self.instance().program_commitment().clone()
120    }
121
122    /// Generates proof for every continuation segment
123    #[instrument(
124        name = "app_prove",
125        skip_all,
126        fields(group = self.program_name.as_ref().unwrap_or(&"app_proof".to_string()))
127    )]
128    pub fn prove(
129        &mut self,
130        input: StdIn<Val<E::SC>>,
131    ) -> Result<ContinuationVmProof<E::SC>, VirtualMachineError>
132    where
133        <VB::VmConfig as VmExecutionConfig<Val<E::SC>>>::Executor: Executor<Val<E::SC>>
134            + MeteredExecutor<Val<E::SC>>
135            + PreflightExecutor<Val<E::SC>, VB::RecordArena>,
136    {
137        assert!(self.vm_config().as_ref().continuation_enabled);
138        check_max_constraint_degrees(
139            self.vm_config().as_ref(),
140            &self.instance.vm.engine.fri_params(),
141        );
142        #[cfg(feature = "metrics")]
143        metrics::counter!("fri.log_blowup")
144            .absolute(self.instance.vm.engine.fri_params().log_blowup as u64);
145        ContinuationVmProver::prove(&mut self.instance, input)
146    }
147
148    /// Generates proof for every continuation segment
149    ///
150    /// This function internally calls [verify_segments] to verify the result before returning the
151    /// proof.
152    ///
153    /// **Note**: This function calls [`app_commit`](Self::app_commit), which is computationally
154    /// intensive if it is the first time it is called within an `AppProver` instance.
155    #[instrument(name = "app_prove_and_verify", skip_all)]
156    pub fn prove_and_verify(
157        &mut self,
158        input: StdIn<Val<E::SC>>,
159    ) -> Result<ContinuationVmProof<E::SC>, VirtualMachineError>
160    where
161        <VB::VmConfig as VmExecutionConfig<Val<E::SC>>>::Executor: Executor<Val<E::SC>>
162            + MeteredExecutor<Val<E::SC>>
163            + PreflightExecutor<Val<E::SC>, VB::RecordArena>,
164    {
165        let proofs = self.prove(input)?;
166        // We skip verification of the user public values proof here because it is directly computed
167        // from the merkle tree above
168        let res = verify_segments(
169            &self.instance.vm.engine,
170            &self.app_vm_vk,
171            &proofs.per_segment,
172        )?;
173        let app_exe_commit_u32s = self.app_commit().app_exe_commit.to_u32_digest();
174        let exe_commit_u32s = res.exe_commit.map(|x| x.as_canonical_u32());
175        if exe_commit_u32s != app_exe_commit_u32s {
176            return Err(VmVerificationError::ExeCommitMismatch {
177                expected: app_exe_commit_u32s,
178                actual: exe_commit_u32s,
179            }
180            .into());
181        }
182        Ok(proofs)
183    }
184
185    /// App Exe
186    pub fn exe(&self) -> Arc<VmExe<Val<E::SC>>> {
187        self.instance.exe().clone()
188    }
189
190    /// App VM
191    pub fn vm(&self) -> &VirtualMachine<E, VB> {
192        &self.instance.vm
193    }
194
195    /// App VM config
196    pub fn vm_config(&self) -> &VB::VmConfig {
197        self.instance.vm.config()
198    }
199}
200
201/// The payload of a verified guest VM execution with user public values extracted and
202/// verified.
203pub struct VerifiedAppArtifacts {
204    /// The Merklelized hash of:
205    /// - Program code commitment (commitment of the cached trace)
206    /// - Merkle root of the initial memory
207    /// - Starting program counter (`pc_start`)
208    ///
209    /// The Merklelization uses Poseidon2 as a cryptographic hash function (for the leaves)
210    /// and a cryptographic compression function (for internal nodes).
211    pub app_exe_commit: CommitBytes,
212    pub user_public_values: Vec<u8>,
213}
214
215/// Verifies the [ContinuationVmProof], which is a collection of STARK proofs as well as
216/// additional Merkle proof for user public values.
217///
218/// This function verifies the STARK proofs and additional conditions to ensure that the
219/// `proof` is a valid proof of guest VM execution that terminates successfully (exit code 0)
220/// _with respect to_ a commitment to some VM executable.
221/// It is the responsibility of the caller to check that the commitment matches the expected
222/// VM executable.
223pub fn verify_app_proof(
224    app_vk: &AppVerifyingKey,
225    proof: &ContinuationVmProof<SC>,
226) -> Result<VerifiedAppArtifacts, VmVerificationError> {
227    static POSEIDON2_HASHER: OnceLock<Poseidon2Hasher<F>> = OnceLock::new();
228    let engine = BabyBearPoseidon2Engine::new(app_vk.fri_params);
229    let VerifiedExecutionPayload {
230        exe_commit,
231        final_memory_root,
232    } = verify_segments(&engine, &app_vk.vk, &proof.per_segment)?;
233
234    proof.user_public_values.verify(
235        POSEIDON2_HASHER.get_or_init(vm_poseidon2_hasher),
236        app_vk.memory_dimensions,
237        final_memory_root,
238    )?;
239
240    let app_exe_commit = CommitBytes::from_u32_digest(&exe_commit.map(|x| x.as_canonical_u32()));
241    // The user public values address space has cells have type u8
242    let user_public_values = proof
243        .user_public_values
244        .public_values
245        .iter()
246        .map(|x| x.as_canonical_u32().try_into().unwrap())
247        .collect_vec();
248    Ok(VerifiedAppArtifacts {
249        app_exe_commit,
250        user_public_values,
251    })
252}
253
254#[cfg(feature = "async")]
255mod async_prover {
256    use derivative::Derivative;
257    use eyre::eyre;
258    use openvm_circuit::{
259        arch::ExecutionError, system::memory::merkle::public_values::UserPublicValuesProof,
260    };
261    use openvm_stark_sdk::config::FriParameters;
262    use tokio::{spawn, sync::Semaphore, task::spawn_blocking};
263    use tracing::{info_span, instrument, Instrument};
264
265    use super::*;
266
267    /// Thread-safe asynchronous app prover.
268    #[derive(Derivative, Getters)]
269    #[derivative(Clone)]
270    pub struct AsyncAppProver<E, VB>
271    where
272        E: StarkEngine,
273        VB: VmBuilder<E>,
274    {
275        pub program_name: Option<String>,
276        #[getset(get = "pub")]
277        vm_builder: VB,
278        #[getset(get = "pub")]
279        app_vm_pk: Arc<VmProvingKey<E::SC, VB::VmConfig>>,
280        app_exe: Arc<VmExe<Val<E::SC>>>,
281        #[getset(get = "pub")]
282        leaf_verifier_program_commit: Com<E::SC>,
283
284        semaphore: Arc<Semaphore>,
285    }
286
287    impl<E, VB> AsyncAppProver<E, VB>
288    where
289        E: StarkFriEngine + 'static,
290        VB: VmBuilder<E> + Clone + Send + Sync + 'static,
291        VB::VmConfig: Send + Sync,
292        <VB::VmConfig as VmExecutionConfig<Val<E::SC>>>::Executor: Executor<Val<E::SC>>
293            + MeteredExecutor<Val<E::SC>>
294            + PreflightExecutor<Val<E::SC>, VB::RecordArena>,
295        Val<E::SC>: PrimeField32,
296        Com<E::SC>:
297            AsRef<[Val<E::SC>; CHUNK]> + From<[Val<E::SC>; CHUNK]> + Into<[Val<E::SC>; CHUNK]>,
298    {
299        pub fn new(
300            vm_builder: VB,
301            app_vm_pk: Arc<VmProvingKey<E::SC, VB::VmConfig>>,
302            app_exe: Arc<VmExe<Val<E::SC>>>,
303            leaf_verifier_program_commit: Com<E::SC>,
304            max_concurrency: usize,
305        ) -> Result<Self, VirtualMachineError> {
306            Ok(Self {
307                program_name: None,
308                vm_builder,
309                app_vm_pk,
310                app_exe,
311                leaf_verifier_program_commit,
312                semaphore: Arc::new(Semaphore::new(max_concurrency)),
313            })
314        }
315
316        pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
317            self.program_name = Some(program_name.as_ref().to_string());
318            self
319        }
320        pub fn with_program_name(mut self, program_name: impl AsRef<str>) -> Self {
321            self.set_program_name(program_name);
322            self
323        }
324
325        /// App Exe
326        pub fn exe(&self) -> Arc<VmExe<Val<E::SC>>> {
327            self.app_exe.clone()
328        }
329
330        /// App VM config
331        pub fn vm_config(&self) -> &VB::VmConfig {
332            &self.app_vm_pk.vm_config
333        }
334
335        pub fn fri_params(&self) -> FriParameters {
336            self.app_vm_pk.fri_params
337        }
338
339        /// Creates an [AppProver] within a particular thread. The former instance is not
340        /// thread-safe and should **not** be moved between threads.
341        pub fn local(&self) -> Result<AppProver<E, VB>, VirtualMachineError> {
342            AppProver::new(
343                self.vm_builder.clone(),
344                &self.app_vm_pk,
345                self.app_exe.clone(),
346                self.leaf_verifier_program_commit.clone(),
347            )
348        }
349
350        #[instrument(
351            name = "app proof",
352            skip_all,
353            fields(
354                group = self.program_name.as_ref().unwrap_or(&"app_proof".to_string())
355            )
356        )]
357        pub async fn prove(
358            self,
359            input: StdIn<Val<E::SC>>,
360        ) -> eyre::Result<ContinuationVmProof<E::SC>> {
361            assert!(self.vm_config().as_ref().continuation_enabled);
362            check_max_constraint_degrees(self.vm_config().as_ref(), &self.fri_params());
363            #[cfg(feature = "metrics")]
364            metrics::counter!("fri.log_blowup").absolute(self.fri_params().log_blowup as u64);
365
366            // PERF[jpw]: it is possible to create metered_interpreter without creating vm. The
367            // latter is more convenient, but does unnecessary setup (e.g., transfer pk to
368            // device). Also, app_commit should be cached.
369            let mut local_prover = self.local()?;
370            let app_commit = local_prover.app_commit();
371            local_prover.instance.reset_state(input.clone());
372            let mut state = local_prover.instance.state_mut().take().unwrap();
373            let vm = &mut local_prover.instance.vm;
374            let metered_ctx = vm.build_metered_ctx(&self.app_exe);
375            let metered_interpreter = vm.metered_interpreter(&self.app_exe)?;
376            let (segments, _) = metered_interpreter.execute_metered(input, metered_ctx)?;
377            drop(metered_interpreter);
378            let pure_interpreter = vm.interpreter(&self.app_exe)?;
379            let mut tasks = Vec::with_capacity(segments.len());
380            let mut num_ins_last = 0;
381            let user_pv_proof: Arc<OnceLock<UserPublicValuesProof<CHUNK, Val<E::SC>>>> =
382                Arc::new(OnceLock::new());
383            for (seg_idx, segment) in segments.into_iter().enumerate() {
384                let semaphore = self.semaphore.clone();
385                let async_worker = self.clone();
386                state = pure_interpreter.execute_from_state(state, Some(num_ins_last))?;
387                num_ins_last = segment.num_insns;
388                let start_state = state.clone();
389                let user_pv_proof = user_pv_proof.clone();
390                let task = spawn(
391                    async move {
392                        let _permit = semaphore.acquire().await?;
393                        let span = tracing::Span::current();
394                        spawn_blocking(move || {
395                            let _span = span.enter();
396                            info_span!("prove_segment", segment = seg_idx).in_scope(
397                                || -> eyre::Result<_> {
398                                    // We need a separate span so the metric label includes
399                                    // "segment"
400                                    // from _segment_span
401                                    let _prove_span = info_span!(
402                                        "vm_prove",
403                                        thread_id = ?std::thread::current().id()
404                                    )
405                                    .entered();
406                                    let mut worker = async_worker.local()?;
407                                    let instance = &mut worker.instance;
408                                    let vm = &mut instance.vm;
409                                    let preflight_interpreter = &mut instance.interpreter;
410                                    let (segment_proof, final_memory) = vm.prove(
411                                        preflight_interpreter,
412                                        start_state,
413                                        Some(segment.num_insns),
414                                        &segment.trace_heights,
415                                    )?;
416                                    if let Some(final_memory) = final_memory {
417                                        let top_tree = vm.memory_top_tree().unwrap();
418                                        let proof = UserPublicValuesProof::compute(
419                                            vm.config().as_ref().memory_config.memory_dimensions(),
420                                            vm.config().as_ref().num_public_values,
421                                            &vm_poseidon2_hasher(),
422                                            &final_memory.memory,
423                                            top_tree,
424                                        );
425                                        user_pv_proof.set(proof).map_err(|_| {
426                                            eyre!("Only one segment should be terminal")
427                                        })?;
428                                    }
429                                    Ok(segment_proof)
430                                },
431                            )
432                        })
433                        .await?
434                    }
435                    .in_current_span(),
436                );
437                tasks.push(task);
438            }
439
440            let mut proofs = Vec::with_capacity(tasks.len());
441            for task in tasks {
442                let proof = task.await??;
443                proofs.push(proof);
444            }
445            let user_public_values = user_pv_proof
446                .get()
447                .ok_or(ExecutionError::DidNotTerminate)?
448                .clone();
449            let cont_proof = ContinuationVmProof {
450                per_segment: proofs,
451                user_public_values,
452            };
453
454            // We skip verification of the user public values proof here because it is directly
455            // computed from the merkle tree above
456            let engine = E::new(self.fri_params());
457            let res = verify_segments(
458                &engine,
459                &self.app_vm_pk.vm_pk.get_vk(),
460                &cont_proof.per_segment,
461            )?;
462            let app_exe_commit_u32s = app_commit.app_exe_commit.to_u32_digest();
463            let exe_commit_u32s = res.exe_commit.map(|x| x.as_canonical_u32());
464            if exe_commit_u32s != app_exe_commit_u32s {
465                return Err(VmVerificationError::ExeCommitMismatch {
466                    expected: app_exe_commit_u32s,
467                    actual: exe_commit_u32s,
468                }
469                .into());
470            }
471            Ok(cont_proof)
472        }
473    }
474}
openvm_sdk/prover/app.rs

openvm_sdk/prover/
app.rs
