openvm_continuations/circuit/deferral/hook/
trace.rs

1use std::borrow::Borrow;
2
3#[cfg(feature = "cuda")]
4use openvm_circuit_primitives::hybrid_chip::cpu_proving_ctx_to_gpu;
5use openvm_cpu_backend::CpuBackend;
6#[cfg(feature = "cuda")]
7use openvm_cuda_backend::GpuBackend;
8#[cfg(feature = "cuda")]
9use openvm_cuda_common::stream::GpuDeviceCtx;
10use openvm_poseidon2_air::POSEIDON2_WIDTH;
11use openvm_stark_backend::{
12    proof::Proof,
13    prover::{AirProvingContext, ProverBackend},
14};
15use openvm_stark_sdk::config::baby_bear_poseidon2::{BabyBearPoseidon2Config, DIGEST_SIZE, F};
16use openvm_verify_stark_host::pvs::VerifierBasePvs;
17use p3_field::{PrimeCharacteristicRing, PrimeField32};
18
19use crate::{
20    circuit::{deferral::MAX_DEF_AGG_MERKLE_DEPTH, SingleAirTraceData},
21    SC,
22};
23
24pub type DeferralIoCommit<F> = ([F; DIGEST_SIZE], [F; DIGEST_SIZE]);
25
26pub struct DeferralHookPreCtx<PB: ProverBackend> {
27    pub verifier_pvs_ctx: AirProvingContext<PB>,
28    pub decommit_ctx: AirProvingContext<PB>,
29    pub onion_ctx: AirProvingContext<PB>,
30    pub poseidon2_compress_inputs: Vec<[PB::Val; POSEIDON2_WIDTH]>,
31    pub poseidon2_permute_inputs: Vec<[PB::Val; POSEIDON2_WIDTH]>,
32    pub range_check_inputs: Vec<usize>,
33    pub power_check_inputs: Vec<usize>,
34}
35
36// Trait used to remain generic in PB.
37pub trait DeferralHookTraceGen<PB: ProverBackend, DC: Clone + Send + Sync> {
38    fn new() -> Self;
39
40    fn pre_verifier_subcircuit_tracegen(
41        &self,
42        proof: &Proof<SC>,
43        leaf_children: Vec<DeferralIoCommit<PB::Val>>,
44        device_ctx: &DC,
45    ) -> DeferralHookPreCtx<PB>;
46}
47
48pub struct DeferralHookTraceGenImpl;
49
50fn normalize_leaf_children(
51    mut leaf_children: Vec<DeferralIoCommit<F>>,
52) -> (Vec<DeferralIoCommit<F>>, usize) {
53    assert!(
54        !leaf_children.is_empty(),
55        "deferral hook requires at least one leaf commit"
56    );
57    let num_real_leaves = leaf_children.len();
58    let target_len = leaf_children.len().next_power_of_two();
59    leaf_children.resize(target_len, ([F::ZERO; DIGEST_SIZE], [F::ZERO; DIGEST_SIZE]));
60    (leaf_children, num_real_leaves)
61}
62
63impl DeferralHookTraceGen<CpuBackend<BabyBearPoseidon2Config>, ()> for DeferralHookTraceGenImpl {
64    fn new() -> Self {
65        Self
66    }
67
68    fn pre_verifier_subcircuit_tracegen(
69        &self,
70        proof: &Proof<SC>,
71        leaf_children: Vec<DeferralIoCommit<F>>,
72        _device_ctx: &(),
73    ) -> DeferralHookPreCtx<CpuBackend<BabyBearPoseidon2Config>> {
74        let (leaf_children, num_real_leaves) = normalize_leaf_children(leaf_children);
75        let super::decommit::MerkleDecommitTraceCtx {
76            proving_ctx: decommit_ctx,
77            poseidon2_inputs: decommit_p2_inputs,
78            io_commits,
79            merkle_root: computed_merkle_root,
80        } = super::decommit::generate_proving_ctx(leaf_children, num_real_leaves);
81
82        let def_pvs: &crate::circuit::deferral::DeferralAggregationPvs<F> =
83            proof.public_values[1].as_slice().borrow();
84        let merkle_depth = def_pvs.merkle_depth.as_canonical_u32() as usize;
85        let max_depth_minus_merkle_depth = MAX_DEF_AGG_MERKLE_DEPTH
86            .checked_sub(merkle_depth)
87            .expect("deferral merkle depth exceeds max depth");
88        assert_eq!(
89            computed_merkle_root, def_pvs.merkle_commit,
90            "leaf_children do not match the child proof merkle_commit"
91        );
92
93        let verifier_pvs: &VerifierBasePvs<F> = proof.public_values[0].as_slice().borrow();
94        let def_circuit_commit =
95            super::verifier::def_circuit_commit_from_verifier_pvs(verifier_pvs);
96
97        let super::onion::OnionTraceCtx {
98            proving_ctx: onion_ctx,
99            poseidon2_inputs: onion_p2_inputs,
100            input_onion,
101            output_onion,
102        } = super::onion::generate_proving_ctx(def_circuit_commit, io_commits);
103
104        let super::verifier::DeferralHookVerifierTraceCtx {
105            trace_data:
106                SingleAirTraceData {
107                    air_proving_ctx: verifier_pvs_ctx,
108                    poseidon2_compress_inputs: verifier_p2_compress_inputs,
109                    poseidon2_permute_inputs: verifier_p2_permute_inputs,
110                    mut range_check_inputs,
111                },
112            ..
113        } = super::verifier::generate_proving_ctx(proof, input_onion, output_onion);
114
115        range_check_inputs.push(max_depth_minus_merkle_depth);
116
117        DeferralHookPreCtx {
118            verifier_pvs_ctx,
119            decommit_ctx,
120            onion_ctx,
121            poseidon2_compress_inputs: verifier_p2_compress_inputs
122                .into_iter()
123                .chain(decommit_p2_inputs)
124                .chain(onion_p2_inputs)
125                .collect(),
126            poseidon2_permute_inputs: verifier_p2_permute_inputs,
127            range_check_inputs,
128            power_check_inputs: vec![merkle_depth],
129        }
130    }
131}
132
133#[cfg(feature = "cuda")]
134impl DeferralHookTraceGen<GpuBackend, GpuDeviceCtx> for DeferralHookTraceGenImpl {
135    fn new() -> Self {
136        Self
137    }
138
139    fn pre_verifier_subcircuit_tracegen(
140        &self,
141        proof: &Proof<SC>,
142        leaf_children: Vec<DeferralIoCommit<F>>,
143        device_ctx: &GpuDeviceCtx,
144    ) -> DeferralHookPreCtx<GpuBackend> {
145        let DeferralHookPreCtx {
146            verifier_pvs_ctx,
147            decommit_ctx,
148            onion_ctx,
149            poseidon2_compress_inputs,
150            poseidon2_permute_inputs,
151            range_check_inputs,
152            power_check_inputs,
153        } = <Self as DeferralHookTraceGen<CpuBackend<BabyBearPoseidon2Config>, ()>>::pre_verifier_subcircuit_tracegen(self, proof, leaf_children, &());
154
155        DeferralHookPreCtx {
156            verifier_pvs_ctx: cpu_proving_ctx_to_gpu(verifier_pvs_ctx, device_ctx),
157            decommit_ctx: cpu_proving_ctx_to_gpu(decommit_ctx, device_ctx),
158            onion_ctx: cpu_proving_ctx_to_gpu(onion_ctx, device_ctx),
159            poseidon2_compress_inputs,
160            poseidon2_permute_inputs,
161            range_check_inputs,
162            power_check_inputs,
163        }
164    }
165}
openvm_continuations/circuit/deferral/hook/trace.rs

openvm_continuations/circuit/deferral/hook/
trace.rs
