openvm_continuations/circuit/deferral/inner/
trace.rs

1use std::{borrow::Borrow, iter::once};
2
3use itertools::Itertools;
4use openvm_circuit::arch::POSEIDON2_WIDTH;
5#[cfg(feature = "cuda")]
6use openvm_circuit_primitives::hybrid_chip::cpu_proving_ctx_to_gpu;
7use openvm_cpu_backend::CpuBackend;
8#[cfg(feature = "cuda")]
9use openvm_cuda_backend::GpuBackend;
10#[cfg(feature = "cuda")]
11use openvm_cuda_common::stream::GpuDeviceCtx;
12use openvm_recursion_circuit::utils::poseidon2_hash_slice_with_states;
13use openvm_stark_backend::{
14    proof::Proof,
15    prover::{AirProvingContext, ProverBackend},
16};
17use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, DIGEST_SIZE, F};
18use openvm_verify_stark_host::pvs::VkCommit;
19use p3_field::PrimeCharacteristicRing;
20
21use crate::{
22    circuit::deferral::{
23        utils::{def_leaf_compress, def_tagged_compress, def_zero_hash},
24        DeferralAggregationPvs, DeferralCircuitPvs, DEF_AGG_PVS_AIR_ID, DEF_CIRCUIT_PVS_AIR_ID,
25        DEF_INTERNAL_TAG, DEF_LEAF_TAG,
26    },
27    utils::digests_to_poseidon2_input,
28};
29
30pub struct DeferralInnerPreCtx<PB: ProverBackend> {
31    pub verifier_pvs_ctx: AirProvingContext<PB>,
32    pub def_pvs_ctx: AirProvingContext<PB>,
33    pub input_ctx: 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_check_inputs: Vec<usize>,
37}
38
39fn fold_leaf_input_commit(
40    proof: &Proof<BabyBearPoseidon2Config>,
41    init: [F; DIGEST_SIZE],
42) -> ([F; DIGEST_SIZE], Vec<[F; POSEIDON2_WIDTH]>) {
43    let values = once(init)
44        .chain(
45            proof
46                .trace_vdata
47                .iter()
48                .flatten()
49                .flat_map(|vdata| vdata.cached_commitments.iter().copied()),
50        )
51        .flatten()
52        .collect_vec();
53    let (folded_input_commit, pre_states, _) = poseidon2_hash_slice_with_states(&values);
54    (folded_input_commit, pre_states)
55}
56
57fn child_merkle_commit(
58    proof: &Proof<BabyBearPoseidon2Config>,
59    child_is_agg: bool,
60) -> [F; DIGEST_SIZE] {
61    if child_is_agg {
62        let child_pvs: &DeferralAggregationPvs<F> =
63            proof.public_values[DEF_AGG_PVS_AIR_ID].as_slice().borrow();
64        child_pvs.merkle_commit
65    } else {
66        let child_pvs: &DeferralCircuitPvs<F> = proof.public_values[DEF_CIRCUIT_PVS_AIR_ID]
67            .as_slice()
68            .borrow();
69        let (folded_input_commit, _) = fold_leaf_input_commit(proof, child_pvs.input_commit);
70        def_leaf_compress(folded_input_commit, child_pvs.output_commit).1
71    }
72}
73
74fn push_deferral_compress_inputs(
75    inputs: &mut Vec<[F; POSEIDON2_WIDTH]>,
76    tag: [u8; DIGEST_SIZE],
77    left: [F; DIGEST_SIZE],
78    right: [F; DIGEST_SIZE],
79) {
80    let (tagged_left, _) = def_tagged_compress(tag, left, right);
81    inputs.push(digests_to_poseidon2_input(tag.map(F::from_u8), left));
82    inputs.push(digests_to_poseidon2_input(tagged_left, right));
83}
84
85fn generate_poseidon2_inputs(
86    proofs: &[Proof<BabyBearPoseidon2Config>],
87    child_is_agg: bool,
88    child_merkle_depth: Option<usize>,
89) -> (Vec<[F; POSEIDON2_WIDTH]>, Vec<[F; POSEIDON2_WIDTH]>) {
90    let mut poseidon2_compress_inputs: Vec<[F; POSEIDON2_WIDTH]> = Vec::new();
91    let mut poseidon2_permute_inputs: Vec<[F; POSEIDON2_WIDTH]> = Vec::new();
92
93    for proof in proofs {
94        if child_is_agg {
95            continue;
96        }
97        let child_pvs: &DeferralCircuitPvs<F> = proof.public_values[DEF_CIRCUIT_PVS_AIR_ID]
98            .as_slice()
99            .borrow();
100
101        // InputCommitAir: sponge-hash input_commit and cached trace commitments.
102        let (folded_input_commit, input_permute_inputs) =
103            fold_leaf_input_commit(proof, child_pvs.input_commit);
104        poseidon2_permute_inputs.extend(input_permute_inputs);
105
106        // DeferralAggPvsAir (leaf): hash folded input_commit and output_commit into merkle_commit.
107        push_deferral_compress_inputs(
108            &mut poseidon2_compress_inputs,
109            DEF_LEAF_TAG,
110            folded_input_commit,
111            child_pvs.output_commit,
112        );
113    }
114
115    // DeferralAggPvsAir: hash child merkle commits when this is not a wrapper.
116    if let Some(depth) = child_merkle_depth {
117        let left_merkle = child_merkle_commit(&proofs[0], child_is_agg);
118        let right_merkle = if proofs.len() == 2 {
119            child_merkle_commit(&proofs[1], child_is_agg)
120        } else {
121            def_zero_hash(depth + 1)
122        };
123        push_deferral_compress_inputs(
124            &mut poseidon2_compress_inputs,
125            DEF_INTERNAL_TAG,
126            left_merkle,
127            right_merkle,
128        );
129    }
130
131    (poseidon2_compress_inputs, poseidon2_permute_inputs)
132}
133
134// Trait used to remain generic in PB
135pub trait DeferralInnerTraceGen<PB: ProverBackend, DC: Clone + Send + Sync> {
136    fn new() -> Self;
137    fn pre_verifier_subcircuit_tracegen(
138        &self,
139        proofs: &[Proof<BabyBearPoseidon2Config>],
140        child_is_agg: bool,
141        child_vk_commit: VkCommit<F>,
142        child_merkle_depth: Option<usize>,
143        device_ctx: &DC,
144    ) -> DeferralInnerPreCtx<PB>;
145}
146
147pub struct DeferralInnerTraceGenImpl;
148
149impl DeferralInnerTraceGen<CpuBackend<BabyBearPoseidon2Config>, ()> for DeferralInnerTraceGenImpl {
150    fn new() -> Self {
151        Self
152    }
153
154    fn pre_verifier_subcircuit_tracegen(
155        &self,
156        proofs: &[Proof<BabyBearPoseidon2Config>],
157        child_is_agg: bool,
158        child_vk_commit: VkCommit<F>,
159        child_merkle_depth: Option<usize>,
160        _device_ctx: &(),
161    ) -> DeferralInnerPreCtx<CpuBackend<BabyBearPoseidon2Config>> {
162        let (poseidon2_compress_inputs, poseidon2_permute_inputs) =
163            generate_poseidon2_inputs(proofs, child_is_agg, child_merkle_depth);
164        let super::def_pvs::DeferralAggPvsTraceCtx {
165            proving_ctx: def_pvs_ctx,
166            mut range_check_inputs,
167        } = super::def_pvs::generate_proving_ctx(proofs, child_is_agg, child_merkle_depth);
168        let super::verifier::DeferralVerifierPvsTraceCtx {
169            proving_ctx: verifier_pvs_ctx,
170            range_check_inputs: verifier_range_check_inputs,
171        } = super::verifier::generate_proving_ctx(proofs, child_is_agg, child_vk_commit);
172        range_check_inputs.extend(verifier_range_check_inputs);
173
174        DeferralInnerPreCtx {
175            verifier_pvs_ctx,
176            def_pvs_ctx,
177            input_ctx: super::input::generate_proving_ctx(proofs, child_is_agg),
178            poseidon2_compress_inputs,
179            poseidon2_permute_inputs,
180            range_check_inputs,
181        }
182    }
183}
184
185#[cfg(feature = "cuda")]
186impl DeferralInnerTraceGen<GpuBackend, GpuDeviceCtx> for DeferralInnerTraceGenImpl {
187    fn new() -> Self {
188        Self
189    }
190
191    fn pre_verifier_subcircuit_tracegen(
192        &self,
193        proofs: &[Proof<BabyBearPoseidon2Config>],
194        child_is_agg: bool,
195        child_vk_commit: VkCommit<F>,
196        child_merkle_depth: Option<usize>,
197        device_ctx: &GpuDeviceCtx,
198    ) -> DeferralInnerPreCtx<GpuBackend> {
199        let DeferralInnerPreCtx {
200            verifier_pvs_ctx,
201            def_pvs_ctx,
202            input_ctx,
203            poseidon2_compress_inputs,
204            poseidon2_permute_inputs,
205            range_check_inputs,
206        } = <Self as DeferralInnerTraceGen<CpuBackend<BabyBearPoseidon2Config>, ()>>::pre_verifier_subcircuit_tracegen(
207            self,
208            proofs,
209            child_is_agg,
210            child_vk_commit,
211            child_merkle_depth,
212            &(),
213        );
214        DeferralInnerPreCtx {
215            verifier_pvs_ctx: cpu_proving_ctx_to_gpu(verifier_pvs_ctx, device_ctx),
216            def_pvs_ctx: cpu_proving_ctx_to_gpu(def_pvs_ctx, device_ctx),
217            input_ctx: cpu_proving_ctx_to_gpu(input_ctx, device_ctx),
218            poseidon2_compress_inputs,
219            poseidon2_permute_inputs,
220            range_check_inputs,
221        }
222    }
223}
openvm_continuations/circuit/deferral/inner/trace.rs

openvm_continuations/circuit/deferral/inner/
trace.rs
