openvm_continuations/circuit/deferral/hook/verifier/
air.rs

1use std::{array::from_fn, borrow::Borrow};
2
3use openvm_circuit::arch::POSEIDON2_WIDTH;
4use openvm_circuit_primitives::{ColumnsAir, StructReflection, StructReflectionHelper, SubAir};
5use openvm_recursion_circuit::{
6    bus::{
7        CachedCommitBus, CachedCommitBusMessage, Poseidon2CompressBus, Poseidon2CompressMessage,
8        PreHashBus, PreHashMessage, PublicValuesBus, PublicValuesBusMessage,
9    },
10    primitives::bus::{
11        PowerCheckerBus, PowerCheckerBusMessage, RangeCheckerBus, RangeCheckerBusMessage,
12    },
13};
14use openvm_recursion_circuit_derive::AlignedBorrow;
15use openvm_stark_backend::{
16    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
17};
18use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE;
19use openvm_verify_stark_host::pvs::{
20    DeferralPvs, VerifierBasePvs, VkCommit, CONSTRAINT_EVAL_AIR_ID, CONSTRAINT_EVAL_CACHED_INDEX,
21    LOG_MAX_RECURSION_DEPTH, VERIFIER_PVS_AIR_ID,
22};
23use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
24use p3_field::{Field, PrimeCharacteristicRing};
25use p3_matrix::Matrix;
26
27use crate::{
28    circuit::{
29        deferral::{
30            hook::bus::{
31                DefCircuitCommitBus, DefCircuitCommitMessage, OnionResultBus, OnionResultMessage,
32            },
33            DeferralAggregationPvs, DEF_AGG_PVS_AIR_ID, MAX_DEF_AGG_MERKLE_DEPTH,
34        },
35        root::NUM_DIGESTS_IN_VM_COMMIT,
36        subair::{HashSliceCtx, HashSliceSubAir, MerkleRootBus, MerkleRootMessage},
37        utils::{assert_vk_commit_eq, assert_vk_commit_unset, vk_commit_components},
38    },
39    utils::{digests_to_poseidon2_input, pad_slice_to_poseidon2_input, zero_hash},
40    CommitBytes, VkCommitBytes,
41};
42
43#[repr(C)]
44#[derive(AlignedBorrow, StructReflection)]
45pub struct DeferralHookPvsCols<F> {
46    pub verifier_pvs: VerifierBasePvs<F>,
47    pub def_pvs: DeferralAggregationPvs<F>,
48
49    pub num_merkle_leaves: F,
50    pub recursion_depth_minus_one_inv: F,
51
52    pub intermediate_vk_states: [[F; POSEIDON2_WIDTH]; NUM_DIGESTS_IN_VM_COMMIT - 1],
53    pub def_circuit_commit: [F; DIGEST_SIZE],
54
55    pub input_onion: [F; DIGEST_SIZE],
56    pub output_onion: [F; DIGEST_SIZE],
57
58    pub def_circuit_commit_padded: [F; DIGEST_SIZE],
59    pub input_onion_padded: [F; DIGEST_SIZE],
60    pub output_onion_padded: [F; DIGEST_SIZE],
61}
62
63#[derive(ColumnsAir)]
64#[columns_via(DeferralHookPvsCols<u8>)]
65pub struct DeferralHookPvsAir {
66    pub public_values_bus: PublicValuesBus,
67    pub cached_commit_bus: CachedCommitBus,
68    pub pre_hash_bus: PreHashBus,
69    pub poseidon2_compress_bus: Poseidon2CompressBus,
70    pub hash_slice_subair: HashSliceSubAir,
71    pub pow_checker_bus: PowerCheckerBus,
72    pub range_checker_bus: RangeCheckerBus,
73
74    pub def_circuit_commit_bus: DefCircuitCommitBus,
75    pub merkle_root_bus: MerkleRootBus,
76    pub onion_res_bus: OnionResultBus,
77
78    pub expected_internal_recursive_vk_commit: VkCommitBytes,
79    pub zero_hash: CommitBytes,
80}
81
82impl DeferralHookPvsAir {
83    #[allow(clippy::too_many_arguments)]
84    pub fn new(
85        public_values_bus: PublicValuesBus,
86        cached_commit_bus: CachedCommitBus,
87        pre_hash_bus: PreHashBus,
88        poseidon2_compress_bus: Poseidon2CompressBus,
89        hash_slice_subair: HashSliceSubAir,
90        pow_checker_bus: PowerCheckerBus,
91        range_checker_bus: RangeCheckerBus,
92        def_circuit_commit_bus: DefCircuitCommitBus,
93        merkle_root_bus: MerkleRootBus,
94        onion_res_bus: OnionResultBus,
95        expected_internal_recursive_vk_commit: VkCommitBytes,
96    ) -> Self {
97        let zero_hash = zero_hash(1).into();
98        Self {
99            public_values_bus,
100            cached_commit_bus,
101            pre_hash_bus,
102            poseidon2_compress_bus,
103            hash_slice_subair,
104            pow_checker_bus,
105            range_checker_bus,
106            def_circuit_commit_bus,
107            merkle_root_bus,
108            onion_res_bus,
109            expected_internal_recursive_vk_commit,
110            zero_hash,
111        }
112    }
113}
114
115impl<F: Field> BaseAir<F> for DeferralHookPvsAir {
116    fn width(&self) -> usize {
117        DeferralHookPvsCols::<u8>::width()
118    }
119}
120impl<F: Field> BaseAirWithPublicValues<F> for DeferralHookPvsAir {
121    fn num_public_values(&self) -> usize {
122        DeferralPvs::<u8>::width()
123    }
124}
125impl<F: Field> PartitionedBaseAir<F> for DeferralHookPvsAir {}
126
127impl<AB: AirBuilder + InteractionBuilder + AirBuilderWithPublicValues> Air<AB>
128    for DeferralHookPvsAir
129{
130    fn eval(&self, builder: &mut AB) {
131        let main = builder.main();
132        let local = main.row_slice(0).expect("window should have one element");
133        let local: &DeferralHookPvsCols<AB::Var> = (*local).borrow();
134
135        /*
136         * Check that the final proof is computed by the internal recursive prover, i.e.
137         * that internal_flag is 2 and that recursion_depth is in [1, MAX_RECURSION_DEPTH].
138         */
139        builder.assert_eq(local.verifier_pvs.internal_flag, AB::F::TWO);
140
141        let depth_minus_one = local.verifier_pvs.recursion_depth - AB::F::ONE;
142        let is_depth_not_one = depth_minus_one.clone() * local.recursion_depth_minus_one_inv;
143        let is_depth_one = AB::Expr::ONE - is_depth_not_one.clone();
144
145        builder.assert_bool(is_depth_one.clone());
146        builder
147            .when(is_depth_one.clone())
148            .assert_one(local.verifier_pvs.recursion_depth);
149
150        self.range_checker_bus.lookup_key(
151            builder,
152            RangeCheckerBusMessage {
153                value: depth_minus_one,
154                max_bits: AB::Expr::from_u8(LOG_MAX_RECURSION_DEPTH),
155            },
156            AB::F::ONE,
157        );
158
159        /*
160         * We need to receive public values from ProofShapeModule to ensure the values read
161         * here are correct. All verifier public values will be at VERIFIER_PVS_AIR_ID,
162         * while the deferral public values will be at DEF_AGG_PVS_AIR_ID.
163         */
164        let verifier_pvs_id = AB::Expr::from_usize(VERIFIER_PVS_AIR_ID);
165        let def_pvs_id = AB::Expr::from_usize(DEF_AGG_PVS_AIR_ID);
166
167        for (pv_idx, value) in local.verifier_pvs.as_slice().iter().enumerate() {
168            self.public_values_bus.receive(
169                builder,
170                AB::F::ZERO,
171                PublicValuesBusMessage {
172                    air_idx: verifier_pvs_id.clone(),
173                    pv_idx: AB::Expr::from_usize(pv_idx),
174                    value: (*value).into(),
175                },
176                AB::F::ONE,
177            );
178        }
179
180        for (pv_idx, value) in local.def_pvs.as_slice().iter().enumerate() {
181            self.public_values_bus.receive(
182                builder,
183                AB::F::ZERO,
184                PublicValuesBusMessage {
185                    air_idx: def_pvs_id.clone(),
186                    pv_idx: AB::Expr::from_usize(pv_idx),
187                    value: (*value).into(),
188                },
189                AB::F::ONE,
190            );
191        }
192
193        /*
194         * We also need to receive the cached commit from ProofShapeModule, which is either for
195         * the internal-for-leaf (i.e. if recursion_depth == 1) or internal-recursive layer. In
196         * the former case we constrain verifier_pvs.internal_recursive_vk_commit to be unset
197         * (i.e. all 0), and in the latter we constrain it to be equal to a pre-generated
198         * constant as it should be the same regardless of def_vk (provided internal system
199         * params are the same).
200         */
201        let cached_commit = from_fn(|i| {
202            local.verifier_pvs.internal_for_leaf_vk_commit.cached_commit[i] * is_depth_one.clone()
203                + local
204                    .verifier_pvs
205                    .internal_recursive_vk_commit
206                    .cached_commit[i]
207                    * is_depth_not_one.clone()
208        });
209        self.cached_commit_bus.receive(
210            builder,
211            AB::F::ZERO,
212            CachedCommitBusMessage {
213                air_idx: AB::Expr::from_usize(CONSTRAINT_EVAL_AIR_ID),
214                cached_idx: AB::Expr::from_usize(CONSTRAINT_EVAL_CACHED_INDEX),
215                global_cached_idx: AB::Expr::ZERO,
216                cached_commit,
217            },
218            AB::F::ONE,
219        );
220
221        self.pre_hash_bus.receive(
222            builder,
223            AB::F::ZERO,
224            PreHashMessage::<AB::F> {
225                vk_pre_hash: self
226                    .expected_internal_recursive_vk_commit
227                    .vk_pre_hash
228                    .into(),
229            },
230            AB::F::ONE,
231        );
232
233        assert_vk_commit_unset(
234            &mut builder.when(is_depth_one),
235            local.verifier_pvs.internal_recursive_vk_commit,
236        );
237
238        assert_vk_commit_eq(
239            &mut builder.when(is_depth_not_one),
240            local.verifier_pvs.internal_recursive_vk_commit,
241            VkCommit::<AB::Expr>::from(self.expected_internal_recursive_vk_commit),
242        );
243
244        /*
245         * Commit def_circuit_commit is hash_slice of the 6 vk_commit_components (cached_commit
246         * and vk_pre_hash for each of def_vk_commit (called app_vk_commit in the
247         * struct), leaf_vk_commit, and internal_for_leaf_vk_commit). We constrain this
248         * here and send def_circuit_commit to its bus.
249         */
250        let vk_commit_components: Vec<_> = vk_commit_components(&local.verifier_pvs)
251            .into_iter()
252            .map(|c| c.map(Into::into))
253            .collect();
254        self.hash_slice_subair.eval(
255            builder,
256            HashSliceCtx {
257                elements: vk_commit_components.as_slice(),
258                intermediate: local
259                    .intermediate_vk_states
260                    .map(|v| v.map(Into::into))
261                    .as_slice(),
262                result: &local.def_circuit_commit.map(Into::into),
263                enabled: &AB::Expr::ONE,
264            },
265        );
266
267        self.def_circuit_commit_bus.send(
268            builder,
269            DefCircuitCommitMessage {
270                def_circuit_commit: local.def_circuit_commit,
271            },
272            AB::F::ONE,
273        );
274
275        /*
276         * The input_onion and output_onion are computed by decommitting def_pvs.merkle_commit
277         * and performing two onion hashes. Both steps are done in different AIRs, but we must
278         * receive these commits to constrain consistency. Note that merkle_depth is constrained
279         * to be in [0, MAX_DEF_AGG_MERKLE_DEPTH] - the power check forces merkle_depth to be in
280         * [0, 31], and the range check forces merkle_depth <= MAX_DEF_AGG_MERKLE_DEPTH.
281         */
282        self.pow_checker_bus.lookup_key(
283            builder,
284            PowerCheckerBusMessage {
285                log: local.def_pvs.merkle_depth,
286                exp: local.num_merkle_leaves,
287            },
288            AB::F::ONE,
289        );
290
291        self.range_checker_bus.lookup_key(
292            builder,
293            RangeCheckerBusMessage {
294                value: AB::Expr::from_usize(MAX_DEF_AGG_MERKLE_DEPTH) - local.def_pvs.merkle_depth,
295                max_bits: AB::Expr::from_usize(8),
296            },
297            AB::F::ONE,
298        );
299
300        self.merkle_root_bus.receive(
301            builder,
302            MerkleRootMessage {
303                merkle_root: local.def_pvs.merkle_commit.map(Into::into),
304                idx: AB::Expr::ZERO,
305                num_rows_or_zero: local.num_merkle_leaves * AB::Expr::TWO,
306            },
307            AB::F::ONE,
308        );
309
310        self.onion_res_bus.receive(
311            builder,
312            OnionResultMessage {
313                input_onion: local.input_onion,
314                output_onion: local.output_onion,
315                num_elements: local.def_pvs.num_def_circuit_proofs,
316            },
317            AB::F::ONE,
318        );
319
320        /*
321         * Finally, we need to constrain that the public values this AIR produces are consistent
322         * with the child's. initial_acc_hash should be the compression of a padded
323         * def_circuit_commit, and final_acc_hash should be the compression of the input and
324         * output onions. node_idx should be equal to the def_idx of the deferral circuit this
325         * hook proof is for. Note that the Merkle root computation hashes each leaf with the
326         * zero digest prior.
327         */
328        let &DeferralPvs::<_> {
329            initial_acc_hash,
330            final_acc_hash,
331            depth,
332            node_idx,
333        } = builder.public_values().borrow();
334
335        builder.assert_one(depth);
336        builder.assert_eq(node_idx, local.def_pvs.def_idx);
337
338        self.poseidon2_compress_bus.lookup_key(
339            builder,
340            Poseidon2CompressMessage {
341                input: pad_slice_to_poseidon2_input(
342                    &local.def_circuit_commit.map(Into::into),
343                    AB::Expr::ZERO,
344                ),
345                output: local.def_circuit_commit_padded.map(Into::into),
346            },
347            AB::F::ONE,
348        );
349
350        self.poseidon2_compress_bus.lookup_key(
351            builder,
352            Poseidon2CompressMessage {
353                input: pad_slice_to_poseidon2_input(
354                    &local.input_onion.map(Into::into),
355                    AB::Expr::ZERO,
356                ),
357                output: local.input_onion_padded.map(Into::into),
358            },
359            AB::F::ONE,
360        );
361
362        self.poseidon2_compress_bus.lookup_key(
363            builder,
364            Poseidon2CompressMessage {
365                input: pad_slice_to_poseidon2_input(
366                    &local.output_onion.map(Into::into),
367                    AB::Expr::ZERO,
368                ),
369                output: local.output_onion_padded.map(Into::into),
370            },
371            AB::F::ONE,
372        );
373
374        self.poseidon2_compress_bus.lookup_key(
375            builder,
376            Poseidon2CompressMessage {
377                input: digests_to_poseidon2_input(
378                    local.def_circuit_commit_padded.map(Into::into),
379                    self.zero_hash.into(),
380                ),
381                output: initial_acc_hash.map(Into::into),
382            },
383            AB::F::ONE,
384        );
385
386        self.poseidon2_compress_bus.lookup_key(
387            builder,
388            Poseidon2CompressMessage {
389                input: digests_to_poseidon2_input(
390                    local.input_onion_padded.map(Into::into),
391                    local.output_onion_padded.map(Into::into),
392                ),
393                output: final_acc_hash.map(Into::into),
394            },
395            AB::F::ONE,
396        );
397    }
398}
openvm_continuations/circuit/deferral/hook/verifier/air.rs

openvm_continuations/circuit/deferral/hook/verifier/
air.rs
