openvm_circuit/arch/
vm.rs

1use std::{borrow::Borrow, collections::VecDeque, marker::PhantomData, mem, sync::Arc};
2
3use openvm_circuit::system::program::trace::compute_exe_commit;
4use openvm_instructions::exe::VmExe;
5use openvm_stark_backend::{
6    config::{Com, Domain, StarkGenericConfig, Val},
7    engine::StarkEngine,
8    keygen::types::{LinearConstraint, MultiStarkProvingKey, MultiStarkVerifyingKey},
9    p3_commit::PolynomialSpace,
10    p3_field::{FieldAlgebra, PrimeField32},
11    proof::Proof,
12    prover::types::{CommittedTraceData, ProofInput},
13    utils::metrics_span,
14    verifier::VerificationError,
15    Chip,
16};
17use serde::{Deserialize, Serialize};
18use thiserror::Error;
19use tracing::info_span;
20
21use super::{
22    ExecutionError, VmComplexTraceHeights, VmConfig, CONNECTOR_AIR_ID, MERKLE_AIR_ID,
23    PROGRAM_AIR_ID, PROGRAM_CACHED_TRACE_INDEX,
24};
25#[cfg(feature = "bench-metrics")]
26use crate::metrics::VmMetrics;
27use crate::{
28    arch::{hasher::poseidon2::vm_poseidon2_hasher, segment::ExecutionSegment},
29    system::{
30        connector::{VmConnectorPvs, DEFAULT_SUSPEND_EXIT_CODE},
31        memory::{
32            merkle::MemoryMerklePvs,
33            paged_vec::AddressMap,
34            tree::public_values::{UserPublicValuesProof, UserPublicValuesProofError},
35            MemoryImage, CHUNK,
36        },
37        program::trace::VmCommittedExe,
38    },
39};
40
41#[derive(Error, Debug)]
42pub enum GenerationError {
43    #[error("generated trace heights violate constraints")]
44    TraceHeightsLimitExceeded,
45    #[error(transparent)]
46    Execution(#[from] ExecutionError),
47}
48
49/// VM memory state for continuations.
50pub type VmMemoryState<F> = MemoryImage<F>;
51
52#[derive(Clone, Default, Debug)]
53pub struct Streams<F> {
54    pub input_stream: VecDeque<Vec<F>>,
55    pub hint_stream: VecDeque<F>,
56    pub hint_space: Vec<Vec<F>>,
57}
58
59impl<F> Streams<F> {
60    pub fn new(input_stream: impl Into<VecDeque<Vec<F>>>) -> Self {
61        Self {
62            input_stream: input_stream.into(),
63            hint_stream: VecDeque::default(),
64            hint_space: Vec::default(),
65        }
66    }
67}
68
69impl<F> From<VecDeque<Vec<F>>> for Streams<F> {
70    fn from(value: VecDeque<Vec<F>>) -> Self {
71        Streams::new(value)
72    }
73}
74
75impl<F> From<Vec<Vec<F>>> for Streams<F> {
76    fn from(value: Vec<Vec<F>>) -> Self {
77        Streams::new(value)
78    }
79}
80
81pub struct VmExecutor<F, VC> {
82    pub config: VC,
83    pub overridden_heights: Option<VmComplexTraceHeights>,
84    pub trace_height_constraints: Vec<LinearConstraint>,
85    _marker: PhantomData<F>,
86}
87
88#[repr(i32)]
89pub enum ExitCode {
90    Success = 0,
91    Error = 1,
92    Suspended = -1, // Continuations
93}
94
95pub struct VmExecutorResult<SC: StarkGenericConfig> {
96    pub per_segment: Vec<ProofInput<SC>>,
97    /// When VM is running on persistent mode, public values are stored in a special memory space.
98    pub final_memory: Option<VmMemoryState<Val<SC>>>,
99}
100
101pub struct VmExecutorNextSegmentState<F: PrimeField32> {
102    pub memory: MemoryImage<F>,
103    pub input: Streams<F>,
104    pub pc: u32,
105    #[cfg(feature = "bench-metrics")]
106    pub metrics: VmMetrics,
107}
108
109impl<F: PrimeField32> VmExecutorNextSegmentState<F> {
110    pub fn new(memory: MemoryImage<F>, input: impl Into<Streams<F>>, pc: u32) -> Self {
111        Self {
112            memory,
113            input: input.into(),
114            pc,
115            #[cfg(feature = "bench-metrics")]
116            metrics: VmMetrics::default(),
117        }
118    }
119}
120
121pub struct VmExecutorOneSegmentResult<F: PrimeField32, VC: VmConfig<F>> {
122    pub segment: ExecutionSegment<F, VC>,
123    pub next_state: Option<VmExecutorNextSegmentState<F>>,
124}
125
126impl<F, VC> VmExecutor<F, VC>
127where
128    F: PrimeField32,
129    VC: VmConfig<F>,
130{
131    /// Create a new VM executor with a given config.
132    ///
133    /// The VM will start with a single segment, which is created from the initial state.
134    pub fn new(config: VC) -> Self {
135        Self::new_with_overridden_trace_heights(config, None)
136    }
137
138    pub fn set_override_trace_heights(&mut self, overridden_heights: VmComplexTraceHeights) {
139        self.overridden_heights = Some(overridden_heights);
140    }
141
142    pub fn new_with_overridden_trace_heights(
143        config: VC,
144        overridden_heights: Option<VmComplexTraceHeights>,
145    ) -> Self {
146        Self {
147            config,
148            overridden_heights,
149            trace_height_constraints: vec![],
150            _marker: Default::default(),
151        }
152    }
153
154    pub fn continuation_enabled(&self) -> bool {
155        self.config.system().continuation_enabled
156    }
157
158    /// Executes the program in segments.
159    /// After each segment is executed, call the provided closure on the execution result.
160    /// Returns the results from each closure, one per segment.
161    ///
162    /// The closure takes `f(segment_idx, segment) -> R`.
163    pub fn execute_and_then<R, E>(
164        &self,
165        exe: impl Into<VmExe<F>>,
166        input: impl Into<Streams<F>>,
167        mut f: impl FnMut(usize, ExecutionSegment<F, VC>) -> Result<R, E>,
168        map_err: impl Fn(ExecutionError) -> E,
169    ) -> Result<Vec<R>, E> {
170        let mem_config = self.config.system().memory_config;
171        let exe = exe.into();
172        let mut segment_results = vec![];
173        let memory = AddressMap::from_iter(
174            mem_config.as_offset,
175            1 << mem_config.as_height,
176            1 << mem_config.pointer_max_bits,
177            exe.init_memory.clone(),
178        );
179        let pc = exe.pc_start;
180        let mut state = VmExecutorNextSegmentState::new(memory, input, pc);
181        let mut segment_idx = 0;
182
183        loop {
184            let _span = info_span!("execute_segment", segment = segment_idx).entered();
185            let one_segment_result = self
186                .execute_until_segment(exe.clone(), state)
187                .map_err(&map_err)?;
188            segment_results.push(f(segment_idx, one_segment_result.segment)?);
189            if one_segment_result.next_state.is_none() {
190                break;
191            }
192            state = one_segment_result.next_state.unwrap();
193            segment_idx += 1;
194        }
195        tracing::debug!("Number of continuation segments: {}", segment_results.len());
196        #[cfg(feature = "bench-metrics")]
197        metrics::counter!("num_segments").absolute(segment_results.len() as u64);
198
199        Ok(segment_results)
200    }
201
202    pub fn execute_segments(
203        &self,
204        exe: impl Into<VmExe<F>>,
205        input: impl Into<Streams<F>>,
206    ) -> Result<Vec<ExecutionSegment<F, VC>>, ExecutionError> {
207        self.execute_and_then(exe, input, |_, seg| Ok(seg), |err| err)
208    }
209
210    /// Executes a program until a segmentation happens.
211    /// Returns the last segment and the vm state for next segment.
212    /// This is so that the tracegen and proving of this segment can be immediately started (on a separate machine).
213    pub fn execute_until_segment(
214        &self,
215        exe: impl Into<VmExe<F>>,
216        from_state: VmExecutorNextSegmentState<F>,
217    ) -> Result<VmExecutorOneSegmentResult<F, VC>, ExecutionError> {
218        let exe = exe.into();
219        let mut segment = ExecutionSegment::new(
220            &self.config,
221            exe.program.clone(),
222            from_state.input,
223            Some(from_state.memory),
224            self.trace_height_constraints.clone(),
225            exe.fn_bounds.clone(),
226        );
227        #[cfg(feature = "bench-metrics")]
228        {
229            segment.metrics = from_state.metrics;
230        }
231        if let Some(overridden_heights) = self.overridden_heights.as_ref() {
232            segment.set_override_trace_heights(overridden_heights.clone());
233        }
234        let state = metrics_span("execute_time_ms", || segment.execute_from_pc(from_state.pc))?;
235
236        if state.is_terminated {
237            return Ok(VmExecutorOneSegmentResult {
238                segment,
239                next_state: None,
240            });
241        }
242
243        assert!(
244            self.continuation_enabled(),
245            "multiple segments require to enable continuations"
246        );
247        assert_eq!(
248            state.pc,
249            segment.chip_complex.connector_chip().boundary_states[1]
250                .unwrap()
251                .pc
252        );
253        let final_memory = mem::take(&mut segment.final_memory)
254            .expect("final memory should be set in continuations segment");
255        let streams = segment.chip_complex.take_streams();
256        #[cfg(feature = "bench-metrics")]
257        let metrics = segment.metrics.partial_take();
258        Ok(VmExecutorOneSegmentResult {
259            segment,
260            next_state: Some(VmExecutorNextSegmentState {
261                memory: final_memory,
262                input: streams,
263                pc: state.pc,
264                #[cfg(feature = "bench-metrics")]
265                metrics,
266            }),
267        })
268    }
269
270    pub fn execute(
271        &self,
272        exe: impl Into<VmExe<F>>,
273        input: impl Into<Streams<F>>,
274    ) -> Result<Option<VmMemoryState<F>>, ExecutionError> {
275        let mut last = None;
276        self.execute_and_then(
277            exe,
278            input,
279            |_, seg| {
280                last = Some(seg);
281                Ok(())
282            },
283            |err| err,
284        )?;
285        let last = last.expect("at least one segment must be executed");
286        let final_memory = last.final_memory;
287        let end_state =
288            last.chip_complex.connector_chip().boundary_states[1].expect("end state must be set");
289        if end_state.is_terminate != 1 {
290            return Err(ExecutionError::DidNotTerminate);
291        }
292        if end_state.exit_code != ExitCode::Success as u32 {
293            return Err(ExecutionError::FailedWithExitCode(end_state.exit_code));
294        }
295        Ok(final_memory)
296    }
297
298    pub fn execute_and_generate<SC: StarkGenericConfig>(
299        &self,
300        exe: impl Into<VmExe<F>>,
301        input: impl Into<Streams<F>>,
302    ) -> Result<VmExecutorResult<SC>, GenerationError>
303    where
304        Domain<SC>: PolynomialSpace<Val = F>,
305        VC::Executor: Chip<SC>,
306        VC::Periphery: Chip<SC>,
307    {
308        self.execute_and_generate_impl(exe.into(), None, input)
309    }
310
311    pub fn execute_and_generate_with_cached_program<SC: StarkGenericConfig>(
312        &self,
313        committed_exe: Arc<VmCommittedExe<SC>>,
314        input: impl Into<Streams<F>>,
315    ) -> Result<VmExecutorResult<SC>, GenerationError>
316    where
317        Domain<SC>: PolynomialSpace<Val = F>,
318        VC::Executor: Chip<SC>,
319        VC::Periphery: Chip<SC>,
320    {
321        self.execute_and_generate_impl(
322            committed_exe.exe.clone(),
323            Some(committed_exe.committed_program.clone()),
324            input,
325        )
326    }
327
328    fn execute_and_generate_impl<SC: StarkGenericConfig>(
329        &self,
330        exe: VmExe<F>,
331        committed_program: Option<CommittedTraceData<SC>>,
332        input: impl Into<Streams<F>>,
333    ) -> Result<VmExecutorResult<SC>, GenerationError>
334    where
335        Domain<SC>: PolynomialSpace<Val = F>,
336        VC::Executor: Chip<SC>,
337        VC::Periphery: Chip<SC>,
338    {
339        let mut final_memory = None;
340        let per_segment = self.execute_and_then(
341            exe,
342            input,
343            |seg_idx, mut seg| {
344                // Note: this will only be Some on the last segment; otherwise it is
345                // already moved into next segment state
346                final_memory = mem::take(&mut seg.final_memory);
347                tracing::info_span!("trace_gen", segment = seg_idx)
348                    .in_scope(|| seg.generate_proof_input(committed_program.clone()))
349            },
350            GenerationError::Execution,
351        )?;
352
353        Ok(VmExecutorResult {
354            per_segment,
355            final_memory,
356        })
357    }
358
359    pub fn set_trace_height_constraints(&mut self, constraints: Vec<LinearConstraint>) {
360        self.trace_height_constraints = constraints;
361    }
362}
363
364/// A single segment VM.
365pub struct SingleSegmentVmExecutor<F, VC> {
366    pub config: VC,
367    pub overridden_heights: Option<VmComplexTraceHeights>,
368    pub trace_height_constraints: Vec<LinearConstraint>,
369    _marker: PhantomData<F>,
370}
371
372/// Execution result of a single segment VM execution.
373pub struct SingleSegmentVmExecutionResult<F> {
374    /// All user public values
375    pub public_values: Vec<Option<F>>,
376    /// Heights of each AIR, ordered by AIR ID.
377    pub air_heights: Vec<usize>,
378    /// Heights of (SystemBase, Inventory), in an internal ordering.
379    pub internal_heights: VmComplexTraceHeights,
380}
381
382impl<F, VC> SingleSegmentVmExecutor<F, VC>
383where
384    F: PrimeField32,
385    VC: VmConfig<F>,
386{
387    pub fn new(config: VC) -> Self {
388        Self::new_with_overridden_trace_heights(config, None)
389    }
390
391    pub fn new_with_overridden_trace_heights(
392        config: VC,
393        overridden_heights: Option<VmComplexTraceHeights>,
394    ) -> Self {
395        assert!(
396            !config.system().continuation_enabled,
397            "Single segment VM doesn't support continuation mode"
398        );
399        Self {
400            config,
401            overridden_heights,
402            trace_height_constraints: vec![],
403            _marker: Default::default(),
404        }
405    }
406
407    pub fn set_override_trace_heights(&mut self, overridden_heights: VmComplexTraceHeights) {
408        self.overridden_heights = Some(overridden_heights);
409    }
410
411    pub fn set_trace_height_constraints(&mut self, constraints: Vec<LinearConstraint>) {
412        self.trace_height_constraints = constraints;
413    }
414
415    /// Executes a program, compute the trace heights, and returns the public values.
416    pub fn execute_and_compute_heights(
417        &self,
418        exe: impl Into<VmExe<F>>,
419        input: impl Into<Streams<F>>,
420    ) -> Result<SingleSegmentVmExecutionResult<F>, ExecutionError> {
421        let segment = {
422            let mut segment = self.execute_impl(exe.into(), input.into())?;
423            segment.chip_complex.finalize_memory();
424            segment
425        };
426        let air_heights = segment.chip_complex.current_trace_heights();
427        let internal_heights = segment.chip_complex.get_internal_trace_heights();
428        let public_values = if let Some(pv_chip) = segment.chip_complex.public_values_chip() {
429            pv_chip.core.get_custom_public_values()
430        } else {
431            vec![]
432        };
433        Ok(SingleSegmentVmExecutionResult {
434            public_values,
435            air_heights,
436            internal_heights,
437        })
438    }
439
440    /// Executes a program and returns its proof input.
441    pub fn execute_and_generate<SC: StarkGenericConfig>(
442        &self,
443        committed_exe: Arc<VmCommittedExe<SC>>,
444        input: impl Into<Streams<F>>,
445    ) -> Result<ProofInput<SC>, GenerationError>
446    where
447        Domain<SC>: PolynomialSpace<Val = F>,
448        VC::Executor: Chip<SC>,
449        VC::Periphery: Chip<SC>,
450    {
451        let segment = self.execute_impl(committed_exe.exe.clone(), input)?;
452        let proof_input = tracing::info_span!("trace_gen").in_scope(|| {
453            segment.generate_proof_input(Some(committed_exe.committed_program.clone()))
454        })?;
455        Ok(proof_input)
456    }
457
458    fn execute_impl(
459        &self,
460        exe: VmExe<F>,
461        input: impl Into<Streams<F>>,
462    ) -> Result<ExecutionSegment<F, VC>, ExecutionError> {
463        let pc_start = exe.pc_start;
464        let mut segment = ExecutionSegment::new(
465            &self.config,
466            exe.program.clone(),
467            input.into(),
468            None,
469            self.trace_height_constraints.clone(),
470            exe.fn_bounds.clone(),
471        );
472        if let Some(overridden_heights) = self.overridden_heights.as_ref() {
473            segment.set_override_trace_heights(overridden_heights.clone());
474        }
475        metrics_span("execute_time_ms", || segment.execute_from_pc(pc_start))?;
476        Ok(segment)
477    }
478}
479
480#[derive(Error, Debug)]
481pub enum VmVerificationError {
482    #[error("no proof is provided")]
483    ProofNotFound,
484
485    #[error("program commit mismatch (index of mismatch proof: {index}")]
486    ProgramCommitMismatch { index: usize },
487
488    #[error("initial pc mismatch (initial: {initial}, prev_final: {prev_final})")]
489    InitialPcMismatch { initial: u32, prev_final: u32 },
490
491    #[error("initial memory root mismatch")]
492    InitialMemoryRootMismatch,
493
494    #[error("is terminate mismatch (expected: {expected}, actual: {actual})")]
495    IsTerminateMismatch { expected: bool, actual: bool },
496
497    #[error("exit code mismatch")]
498    ExitCodeMismatch { expected: u32, actual: u32 },
499
500    #[error("AIR has unexpected public values (expected: {expected}, actual: {actual})")]
501    UnexpectedPvs { expected: usize, actual: usize },
502
503    #[error("missing system AIR with ID {air_id}")]
504    SystemAirMissing { air_id: usize },
505
506    #[error("stark verification error: {0}")]
507    StarkError(#[from] VerificationError),
508
509    #[error("user public values proof error: {0}")]
510    UserPublicValuesError(#[from] UserPublicValuesProofError),
511}
512
513pub struct VirtualMachine<SC: StarkGenericConfig, E, VC> {
514    /// Proving engine
515    pub engine: E,
516    /// Runtime executor
517    pub executor: VmExecutor<Val<SC>, VC>,
518    _marker: PhantomData<SC>,
519}
520
521impl<F, SC, E, VC> VirtualMachine<SC, E, VC>
522where
523    F: PrimeField32,
524    SC: StarkGenericConfig,
525    E: StarkEngine<SC>,
526    Domain<SC>: PolynomialSpace<Val = F>,
527    VC: VmConfig<F>,
528    VC::Executor: Chip<SC>,
529    VC::Periphery: Chip<SC>,
530{
531    pub fn new(engine: E, config: VC) -> Self {
532        let executor = VmExecutor::new(config);
533        Self {
534            engine,
535            executor,
536            _marker: PhantomData,
537        }
538    }
539
540    pub fn new_with_overridden_trace_heights(
541        engine: E,
542        config: VC,
543        overridden_heights: Option<VmComplexTraceHeights>,
544    ) -> Self {
545        let executor = VmExecutor::new_with_overridden_trace_heights(config, overridden_heights);
546        Self {
547            engine,
548            executor,
549            _marker: PhantomData,
550        }
551    }
552
553    pub fn config(&self) -> &VC {
554        &self.executor.config
555    }
556
557    pub fn keygen(&self) -> MultiStarkProvingKey<SC> {
558        let mut keygen_builder = self.engine.keygen_builder();
559        let chip_complex = self.config().create_chip_complex().unwrap();
560        for air in chip_complex.airs() {
561            keygen_builder.add_air(air);
562        }
563        keygen_builder.generate_pk()
564    }
565
566    pub fn set_trace_height_constraints(
567        &mut self,
568        trace_height_constraints: Vec<LinearConstraint>,
569    ) {
570        self.executor
571            .set_trace_height_constraints(trace_height_constraints);
572    }
573
574    pub fn commit_exe(&self, exe: impl Into<VmExe<F>>) -> Arc<VmCommittedExe<SC>> {
575        let exe = exe.into();
576        Arc::new(VmCommittedExe::commit(exe, self.engine.config().pcs()))
577    }
578
579    pub fn execute(
580        &self,
581        exe: impl Into<VmExe<F>>,
582        input: impl Into<Streams<F>>,
583    ) -> Result<Option<VmMemoryState<F>>, ExecutionError> {
584        self.executor.execute(exe, input)
585    }
586
587    pub fn execute_and_generate(
588        &self,
589        exe: impl Into<VmExe<F>>,
590        input: impl Into<Streams<F>>,
591    ) -> Result<VmExecutorResult<SC>, GenerationError> {
592        self.executor.execute_and_generate(exe, input)
593    }
594
595    pub fn execute_and_generate_with_cached_program(
596        &self,
597        committed_exe: Arc<VmCommittedExe<SC>>,
598        input: impl Into<Streams<F>>,
599    ) -> Result<VmExecutorResult<SC>, GenerationError>
600    where
601        Domain<SC>: PolynomialSpace<Val = F>,
602    {
603        self.executor
604            .execute_and_generate_with_cached_program(committed_exe, input)
605    }
606
607    pub fn prove_single(
608        &self,
609        pk: &MultiStarkProvingKey<SC>,
610        proof_input: ProofInput<SC>,
611    ) -> Proof<SC> {
612        self.engine.prove(pk, proof_input)
613    }
614
615    pub fn prove(
616        &self,
617        pk: &MultiStarkProvingKey<SC>,
618        results: VmExecutorResult<SC>,
619    ) -> Vec<Proof<SC>> {
620        results
621            .per_segment
622            .into_iter()
623            .enumerate()
624            .map(|(seg_idx, proof_input)| {
625                tracing::info_span!("prove_segment", segment = seg_idx)
626                    .in_scope(|| self.engine.prove(pk, proof_input))
627            })
628            .collect()
629    }
630
631    /// Verify segment proofs, checking continuation boundary conditions between segments if VM memory is persistent
632    /// The behavior of this function differs depending on whether continuations is enabled or not.
633    /// We recommend to call the functions [`verify_segments`] or [`verify_single`] directly instead.
634    pub fn verify(
635        &self,
636        vk: &MultiStarkVerifyingKey<SC>,
637        proofs: Vec<Proof<SC>>,
638    ) -> Result<(), VmVerificationError>
639    where
640        Val<SC>: PrimeField32,
641        Com<SC>: AsRef<[Val<SC>; CHUNK]> + From<[Val<SC>; CHUNK]>,
642    {
643        if self.config().system().continuation_enabled {
644            verify_segments(&self.engine, vk, &proofs).map(|_| ())
645        } else {
646            assert_eq!(proofs.len(), 1);
647            verify_single(&self.engine, vk, &proofs.into_iter().next().unwrap())
648                .map_err(VmVerificationError::StarkError)
649        }
650    }
651}
652
653/// Verifies a single proof. This should be used for proof of VM without continuations.
654///
655/// ## Note
656/// This function does not check any public values or extract the starting pc or commitment
657/// to the [VmCommittedExe].
658pub fn verify_single<SC, E>(
659    engine: &E,
660    vk: &MultiStarkVerifyingKey<SC>,
661    proof: &Proof<SC>,
662) -> Result<(), VerificationError>
663where
664    SC: StarkGenericConfig,
665    E: StarkEngine<SC>,
666{
667    engine.verify(vk, proof)
668}
669
670/// The payload of a verified guest VM execution.
671pub struct VerifiedExecutionPayload<F> {
672    /// The Merklelized hash of:
673    /// - Program code commitment (commitment of the cached trace)
674    /// - Merkle root of the initial memory
675    /// - Starting program counter (`pc_start`)
676    ///
677    /// The Merklelization uses Poseidon2 as a cryptographic hash function (for the leaves)
678    /// and a cryptographic compression function (for internal nodes).
679    pub exe_commit: [F; CHUNK],
680    /// The Merkle root of the final memory state.
681    pub final_memory_root: [F; CHUNK],
682}
683
684/// Verify segment proofs with boundary condition checks for continuation between segments.
685///
686/// Assumption:
687/// - `vk` is a valid verifying key of a VM circuit.
688///
689/// Returns:
690/// - The commitment to the [VmCommittedExe] extracted from `proofs`.
691///   It is the responsibility of the caller to check that the returned commitment matches
692///   the VM executable that the VM was supposed to execute.
693/// - The Merkle root of the final memory state.
694///
695/// ## Note
696/// This function does not extract or verify any user public values from the final memory state.
697/// This verification requires an additional Merkle proof with respect to the Merkle root of
698/// the final memory state.
699// @dev: This function doesn't need to be generic in `VC`.
700pub fn verify_segments<SC, E>(
701    engine: &E,
702    vk: &MultiStarkVerifyingKey<SC>,
703    proofs: &[Proof<SC>],
704) -> Result<VerifiedExecutionPayload<Val<SC>>, VmVerificationError>
705where
706    SC: StarkGenericConfig,
707    E: StarkEngine<SC>,
708    Val<SC>: PrimeField32,
709    Com<SC>: AsRef<[Val<SC>; CHUNK]>,
710{
711    if proofs.is_empty() {
712        return Err(VmVerificationError::ProofNotFound);
713    }
714    let mut prev_final_memory_root = None;
715    let mut prev_final_pc = None;
716    let mut start_pc = None;
717    let mut initial_memory_root = None;
718    let mut program_commit = None;
719
720    for (i, proof) in proofs.iter().enumerate() {
721        let res = engine.verify(vk, proof);
722        match res {
723            Ok(_) => (),
724            Err(e) => return Err(VmVerificationError::StarkError(e)),
725        };
726
727        let mut program_air_present = false;
728        let mut connector_air_present = false;
729        let mut merkle_air_present = false;
730
731        // Check public values.
732        for air_proof_data in proof.per_air.iter() {
733            let pvs = &air_proof_data.public_values;
734            let air_vk = &vk.inner.per_air[air_proof_data.air_id];
735            if air_proof_data.air_id == PROGRAM_AIR_ID {
736                program_air_present = true;
737                if i == 0 {
738                    program_commit =
739                        Some(proof.commitments.main_trace[PROGRAM_CACHED_TRACE_INDEX].as_ref());
740                } else if program_commit.unwrap()
741                    != proof.commitments.main_trace[PROGRAM_CACHED_TRACE_INDEX].as_ref()
742                {
743                    return Err(VmVerificationError::ProgramCommitMismatch { index: i });
744                }
745            } else if air_proof_data.air_id == CONNECTOR_AIR_ID {
746                connector_air_present = true;
747                let pvs: &VmConnectorPvs<_> = pvs.as_slice().borrow();
748
749                if i != 0 {
750                    // Check initial pc matches the previous final pc.
751                    if pvs.initial_pc != prev_final_pc.unwrap() {
752                        return Err(VmVerificationError::InitialPcMismatch {
753                            initial: pvs.initial_pc.as_canonical_u32(),
754                            prev_final: prev_final_pc.unwrap().as_canonical_u32(),
755                        });
756                    }
757                } else {
758                    start_pc = Some(pvs.initial_pc);
759                }
760                prev_final_pc = Some(pvs.final_pc);
761
762                let expected_is_terminate = i == proofs.len() - 1;
763                if pvs.is_terminate != FieldAlgebra::from_bool(expected_is_terminate) {
764                    return Err(VmVerificationError::IsTerminateMismatch {
765                        expected: expected_is_terminate,
766                        actual: pvs.is_terminate.as_canonical_u32() != 0,
767                    });
768                }
769
770                let expected_exit_code = if expected_is_terminate {
771                    ExitCode::Success as u32
772                } else {
773                    DEFAULT_SUSPEND_EXIT_CODE
774                };
775                if pvs.exit_code != FieldAlgebra::from_canonical_u32(expected_exit_code) {
776                    return Err(VmVerificationError::ExitCodeMismatch {
777                        expected: expected_exit_code,
778                        actual: pvs.exit_code.as_canonical_u32(),
779                    });
780                }
781            } else if air_proof_data.air_id == MERKLE_AIR_ID {
782                merkle_air_present = true;
783                let pvs: &MemoryMerklePvs<_, CHUNK> = pvs.as_slice().borrow();
784
785                // Check that initial root matches the previous final root.
786                if i != 0 {
787                    if pvs.initial_root != prev_final_memory_root.unwrap() {
788                        return Err(VmVerificationError::InitialMemoryRootMismatch);
789                    }
790                } else {
791                    initial_memory_root = Some(pvs.initial_root);
792                }
793                prev_final_memory_root = Some(pvs.final_root);
794            } else {
795                if !pvs.is_empty() {
796                    return Err(VmVerificationError::UnexpectedPvs {
797                        expected: 0,
798                        actual: pvs.len(),
799                    });
800                }
801                // We assume the vk is valid, so this is only a debug assert.
802                debug_assert_eq!(air_vk.params.num_public_values, 0);
803            }
804        }
805        if !program_air_present {
806            return Err(VmVerificationError::SystemAirMissing {
807                air_id: PROGRAM_AIR_ID,
808            });
809        }
810        if !connector_air_present {
811            return Err(VmVerificationError::SystemAirMissing {
812                air_id: CONNECTOR_AIR_ID,
813            });
814        }
815        if !merkle_air_present {
816            return Err(VmVerificationError::SystemAirMissing {
817                air_id: MERKLE_AIR_ID,
818            });
819        }
820    }
821    let exe_commit = compute_exe_commit(
822        &vm_poseidon2_hasher(),
823        program_commit.unwrap(),
824        initial_memory_root.as_ref().unwrap(),
825        start_pc.unwrap(),
826    );
827    Ok(VerifiedExecutionPayload {
828        exe_commit,
829        final_memory_root: prev_final_memory_root.unwrap(),
830    })
831}
832
833#[derive(Serialize, Deserialize)]
834#[serde(bound(
835    serialize = "Com<SC>: Serialize",
836    deserialize = "Com<SC>: Deserialize<'de>"
837))]
838pub struct ContinuationVmProof<SC: StarkGenericConfig> {
839    pub per_segment: Vec<Proof<SC>>,
840    pub user_public_values: UserPublicValuesProof<{ CHUNK }, Val<SC>>,
841}
842
843impl<SC: StarkGenericConfig> Clone for ContinuationVmProof<SC>
844where
845    Com<SC>: Clone,
846{
847    fn clone(&self) -> Self {
848        Self {
849            per_segment: self.per_segment.clone(),
850            user_public_values: self.user_public_values.clone(),
851        }
852    }
853}