openvm_verify_stark_circuit/
trace.rs

1use std::{borrow::Borrow, iter::once};
2
3use itertools::Itertools;
4use openvm_circuit::system::memory::{
5    dimensions::MemoryDimensions, merkle::public_values::UserPublicValuesProof,
6};
7use openvm_continuations::circuit::{
8    deferral::DeferralMerkleProofs,
9    root::{def_paths, memory},
10};
11use openvm_cpu_backend::CpuBackend;
12use openvm_poseidon2_air::POSEIDON2_WIDTH;
13use openvm_recursion_circuit::prelude::{DIGEST_SIZE, F, SC};
14use openvm_stark_backend::{
15    proof::Proof,
16    prover::{AirProvingContext, ProverBackend},
17};
18use openvm_verify_stark_host::pvs::{DeferralPvs, DEF_PVS_AIR_ID};
19use p3_field::PrimeField32;
20#[cfg(feature = "cuda")]
21use {
22    openvm_circuit_primitives::hybrid_chip::cpu_proving_ctx_to_gpu,
23    openvm_cuda_backend::GpuBackend, openvm_cuda_common::stream::GpuDeviceCtx,
24};
25
26use crate::{
27    output::DeferralOutputCtx,
28    verifier::{generate_record, DeferredVerifyPvsRecord},
29};
30
31pub struct PreVerifierData<PB: ProverBackend> {
32    pub pre_verifier_ctxs: [AirProvingContext<PB>; 2],
33    pub post_verifier_ctxs: Vec<AirProvingContext<PB>>,
34    pub poseidon2_compress_inputs: Vec<[PB::Val; POSEIDON2_WIDTH]>,
35    pub poseidon2_permute_inputs: Vec<[PB::Val; POSEIDON2_WIDTH]>,
36    pub range_inputs: Vec<usize>,
37    pub verifier_pvs_record: DeferredVerifyPvsRecord<PB::Val>,
38    pub output_commit: [PB::Val; DIGEST_SIZE],
39}
40
41// Trait used to remain generic in PB
42pub trait DeferredVerifyTraceGen<PB: ProverBackend, DC: Clone + Send + Sync> {
43    fn new(deferral_enabled: bool) -> Self;
44
45    // Returns the AIR proving contexts, Poseidon2 and range inputs, and the data
46    // needed to compute the DeferredVerifyPvsAir trace later
47    fn pre_verifier_subcircuit_tracegen(
48        &self,
49        proof: &Proof<SC>,
50        user_pvs_proof: &UserPublicValuesProof<DIGEST_SIZE, PB::Val>,
51        memory_dimensions: MemoryDimensions,
52        def_idx: usize,
53        deferral_merkle_proofs: Option<&DeferralMerkleProofs<F>>,
54        device_ctx: &DC,
55    ) -> PreVerifierData<PB>;
56
57    fn generate_verifier_pvs_ctx(
58        &self,
59        proof: &Proof<SC>,
60        record: DeferredVerifyPvsRecord<PB::Val>,
61        final_transcript_state: [PB::Val; POSEIDON2_WIDTH],
62        output_commit: [PB::Val; DIGEST_SIZE],
63        def_idx: usize,
64        device_ctx: &DC,
65    ) -> AirProvingContext<PB>;
66}
67
68pub struct DeferredVerifyTraceGenImpl {
69    pub deferral_enabled: bool,
70}
71
72impl DeferredVerifyTraceGen<CpuBackend<SC>, ()> for DeferredVerifyTraceGenImpl {
73    fn new(deferral_enabled: bool) -> Self {
74        Self { deferral_enabled }
75    }
76
77    fn pre_verifier_subcircuit_tracegen(
78        &self,
79        proof: &Proof<SC>,
80        user_pvs_proof: &UserPublicValuesProof<DIGEST_SIZE, F>,
81        memory_dimensions: MemoryDimensions,
82        def_idx: usize,
83        deferral_merkle_proofs: Option<&DeferralMerkleProofs<F>>,
84        _device_ctx: &(),
85    ) -> PreVerifierData<CpuBackend<SC>> {
86        let (
87            verifier_pvs_record,
88            verifier_p2_compress_inputs,
89            verifier_p2_permute_inputs,
90            verifier_range_inputs,
91        ) = generate_record(proof);
92        let (commit_ctx, commit_p2_inputs) =
93            super::commit::generate_proving_ctx(user_pvs_proof.public_values.clone());
94        let (memory_ctx, memory_p2_inputs) = memory::generate_proving_input(
95            user_pvs_proof.public_values_commit,
96            &user_pvs_proof.proof,
97            memory_dimensions,
98            user_pvs_proof.public_values.len(),
99        );
100        let DeferralOutputCtx {
101            proving_ctx: output_ctx,
102            poseidon2_inputs: output_p2_inputs,
103            range_inputs,
104            output_commit,
105        } = super::output::generate_proving_ctx(
106            verifier_pvs_record.app_exe_commit,
107            verifier_pvs_record.app_vm_commit,
108            user_pvs_proof.public_values.clone(),
109            def_idx,
110        );
111
112        let (paths_ctx, paths_p2_inputs) = if let Some(deferral_merkle_proofs) =
113            deferral_merkle_proofs
114        {
115            assert!(self.deferral_enabled);
116            let def_pvs: &DeferralPvs<_> = proof.public_values[DEF_PVS_AIR_ID].as_slice().borrow();
117            let depth = def_pvs.depth.as_canonical_u32() as usize;
118            let (acc_merkle_paths_ctx, acc_merkle_paths_p2_inputs) =
119                def_paths::generate_proving_input(
120                    def_pvs.initial_acc_hash,
121                    def_pvs.final_acc_hash,
122                    &deferral_merkle_proofs.initial_merkle_proof,
123                    &deferral_merkle_proofs.final_merkle_proof,
124                    memory_dimensions,
125                    depth,
126                    depth == 0,
127                );
128            (Some(acc_merkle_paths_ctx), acc_merkle_paths_p2_inputs)
129        } else {
130            assert!(!self.deferral_enabled);
131            (None, vec![])
132        };
133
134        PreVerifierData {
135            pre_verifier_ctxs: [commit_ctx, memory_ctx],
136            post_verifier_ctxs: once(output_ctx).chain(paths_ctx).collect_vec(),
137            poseidon2_compress_inputs: verifier_p2_compress_inputs
138                .into_iter()
139                .chain(commit_p2_inputs)
140                .chain(memory_p2_inputs)
141                .chain(paths_p2_inputs)
142                .collect_vec(),
143            poseidon2_permute_inputs: verifier_p2_permute_inputs
144                .into_iter()
145                .chain(output_p2_inputs)
146                .collect_vec(),
147            range_inputs: verifier_range_inputs
148                .into_iter()
149                .chain(range_inputs)
150                .collect_vec(),
151            verifier_pvs_record,
152            output_commit,
153        }
154    }
155
156    fn generate_verifier_pvs_ctx(
157        &self,
158        proof: &Proof<SC>,
159        record: DeferredVerifyPvsRecord<F>,
160        final_transcript_state: [F; POSEIDON2_WIDTH],
161        output_commit: [F; DIGEST_SIZE],
162        def_idx: usize,
163        _device_ctx: &(),
164    ) -> AirProvingContext<CpuBackend<SC>> {
165        super::verifier::generate_proving_ctx(
166            proof,
167            record,
168            final_transcript_state,
169            output_commit,
170            def_idx,
171            self.deferral_enabled,
172        )
173    }
174}
175
176#[cfg(feature = "cuda")]
177impl DeferredVerifyTraceGen<GpuBackend, GpuDeviceCtx> for DeferredVerifyTraceGenImpl {
178    fn new(deferral_enabled: bool) -> Self {
179        Self { deferral_enabled }
180    }
181
182    fn pre_verifier_subcircuit_tracegen(
183        &self,
184        proof: &Proof<SC>,
185        user_pvs_proof: &UserPublicValuesProof<DIGEST_SIZE, F>,
186        memory_dimensions: MemoryDimensions,
187        def_idx: usize,
188        deferral_merkle_proofs: Option<&DeferralMerkleProofs<F>>,
189        device_ctx: &GpuDeviceCtx,
190    ) -> PreVerifierData<GpuBackend> {
191        let PreVerifierData {
192            pre_verifier_ctxs,
193            post_verifier_ctxs: post_verifier_ctx,
194            poseidon2_compress_inputs,
195            poseidon2_permute_inputs,
196            range_inputs,
197            verifier_pvs_record,
198            output_commit,
199        } = <Self as DeferredVerifyTraceGen<CpuBackend<SC>, ()>>::pre_verifier_subcircuit_tracegen(
200            self,
201            proof,
202            user_pvs_proof,
203            memory_dimensions,
204            def_idx,
205            deferral_merkle_proofs,
206            &(),
207        );
208
209        PreVerifierData {
210            pre_verifier_ctxs: pre_verifier_ctxs
211                .map(|air_ctx| cpu_proving_ctx_to_gpu(air_ctx, device_ctx)),
212            post_verifier_ctxs: post_verifier_ctx
213                .into_iter()
214                .map(|air_ctx| cpu_proving_ctx_to_gpu(air_ctx, device_ctx))
215                .collect_vec(),
216            poseidon2_compress_inputs,
217            poseidon2_permute_inputs,
218            range_inputs,
219            verifier_pvs_record,
220            output_commit,
221        }
222    }
223
224    fn generate_verifier_pvs_ctx(
225        &self,
226        proof: &Proof<SC>,
227        record: DeferredVerifyPvsRecord<F>,
228        final_transcript_state: [F; POSEIDON2_WIDTH],
229        output_commit: [F; DIGEST_SIZE],
230        def_idx: usize,
231        device_ctx: &GpuDeviceCtx,
232    ) -> AirProvingContext<GpuBackend> {
233        cpu_proving_ctx_to_gpu(
234            super::verifier::generate_proving_ctx(
235                proof,
236                record,
237                final_transcript_state,
238                output_commit,
239                def_idx,
240                self.deferral_enabled,
241            ),
242            device_ctx,
243        )
244    }
245}