openvm_continuations/circuit/deferral/inner/def_pvs/
air.rs

1use std::{array::from_fn, borrow::Borrow};
2
3use itertools::Itertools;
4use openvm_circuit_primitives::{
5    encoder::Encoder,
6    utils::{assert_array_eq, not},
7    ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
8};
9use openvm_recursion_circuit::{
10    bus::{
11        Poseidon2CompressBus, Poseidon2CompressMessage, PublicValuesBus, PublicValuesBusMessage,
12    },
13    prelude::DIGEST_SIZE,
14    primitives::bus::{RangeCheckerBus, RangeCheckerBusMessage},
15};
16use openvm_recursion_circuit_derive::AlignedBorrow;
17use openvm_stark_backend::{
18    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
19};
20use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
21use p3_field::PrimeCharacteristicRing;
22use p3_matrix::Matrix;
23
24use crate::{
25    circuit::deferral::{
26        inner::bus::{
27            DefPvsConsistencyBus, DefPvsConsistencyMessage, InputOrMerkleCommitBus,
28            InputOrMerkleCommitMessage,
29        },
30        utils::def_zero_hashes_from_depth_one,
31        DeferralAggregationPvs, DeferralCircuitPvs, DEF_AGG_PVS_AIR_ID, DEF_CIRCUIT_PVS_AIR_ID,
32        DEF_INTERNAL_TAG, DEF_LEAF_TAG, MAX_DEF_AGG_MERKLE_DEPTH,
33    },
34    utils::digests_to_poseidon2_input,
35    CommitBytes,
36};
37
38const ENCODER_MAX_DEGREE: u32 = 2;
39#[repr(C)]
40#[derive(AlignedBorrow, StructReflection)]
41pub struct DeferralAggPvsCols<F> {
42    pub proof_idx: F,
43    pub is_present: F,
44    pub has_verifier_pvs: F,
45
46    pub merkle_commit: [F; DIGEST_SIZE],
47    pub tagged_input_commit: [F; DIGEST_SIZE],
48    pub tagged_left_merkle: [F; DIGEST_SIZE],
49
50    pub child_pvs: DeferralCircuitPvs<F>,
51}
52
53#[derive(ColumnsAir)]
54#[columns_via(DeferralAggPvsCols<u8>)]
55pub struct DeferralAggPvsAir {
56    pub public_values_bus: PublicValuesBus,
57    pub poseidon2_bus: Poseidon2CompressBus,
58    pub range_bus: RangeCheckerBus,
59    pub input_or_merkle_commit_bus: InputOrMerkleCommitBus,
60    pub def_pvs_consistency_bus: DefPvsConsistencyBus,
61
62    pub encoder: Encoder,
63    pub zero_hashes: [CommitBytes; MAX_DEF_AGG_MERKLE_DEPTH + 1],
64}
65
66impl DeferralAggPvsAir {
67    pub(super) fn depth_encoder() -> Encoder {
68        Encoder::new(MAX_DEF_AGG_MERKLE_DEPTH + 1, ENCODER_MAX_DEGREE, false)
69    }
70
71    pub fn new(
72        public_values_bus: PublicValuesBus,
73        poseidon2_bus: Poseidon2CompressBus,
74        range_bus: RangeCheckerBus,
75        input_or_merkle_commit_bus: InputOrMerkleCommitBus,
76        def_pvs_consistency_bus: DefPvsConsistencyBus,
77    ) -> Self {
78        let encoder = Self::depth_encoder();
79        assert!(encoder.width() < DIGEST_SIZE);
80        Self {
81            public_values_bus,
82            poseidon2_bus,
83            range_bus,
84            input_or_merkle_commit_bus,
85            def_pvs_consistency_bus,
86            encoder,
87            zero_hashes: def_zero_hashes_from_depth_one().map(Into::into),
88        }
89    }
90}
91
92impl<F> BaseAir<F> for DeferralAggPvsAir {
93    fn width(&self) -> usize {
94        DeferralAggPvsCols::<u8>::width()
95    }
96}
97impl<F> BaseAirWithPublicValues<F> for DeferralAggPvsAir {
98    fn num_public_values(&self) -> usize {
99        DeferralAggregationPvs::<u8>::width()
100    }
101}
102impl<F> PartitionedBaseAir<F> for DeferralAggPvsAir {}
103
104impl<AB: AirBuilder + InteractionBuilder + AirBuilderWithPublicValues> Air<AB>
105    for DeferralAggPvsAir
106{
107    fn eval(&self, builder: &mut AB) {
108        let main = builder.main();
109        let (local, next) = (
110            main.row_slice(0).expect("window should have two elements"),
111            main.row_slice(1).expect("window should have two elements"),
112        );
113        let local: &DeferralAggPvsCols<AB::Var> = (*local).borrow();
114        let next: &DeferralAggPvsCols<AB::Var> = (*next).borrow();
115
116        /*
117         * This AIR may have 1 or 2 rows. The first row is always present, while the second row
118         * may be padding (is_present = 0). It is assumed these rows are adjacent leaves in the
119         * deferral aggregation tree; otherwise the final root-layer Merkle check will fail.
120         */
121        builder.assert_bool(local.is_present);
122        builder.when_first_row().assert_one(local.is_present);
123
124        builder.assert_bool(local.proof_idx);
125        builder.when_first_row().assert_zero(local.proof_idx);
126        builder
127            .when_transition()
128            .assert_one(next.proof_idx - local.proof_idx);
129
130        builder.assert_bool(local.has_verifier_pvs);
131        builder.assert_eq(local.has_verifier_pvs, next.has_verifier_pvs);
132
133        /*
134         * We need to receive public values here to ensure the values read are correct.
135         * output_commit is received from ProofShapeModule, and the (possibly cached-folded)
136         * input_commit/merkle_commit is received from InputCommitAir. output_commit is only
137         * present for leaf children.
138         *
139         * A row may be non-present (is_present = 0). For those rows, child-specific public
140         * values are intentionally not constrained here; consistency is enforced by the
141         * final Merkle root checks at the deferral hook layer.
142         */
143        let is_leaf = not(local.has_verifier_pvs);
144        let is_present_leaf = is_leaf.clone() * local.is_present;
145        let is_internal = local.has_verifier_pvs;
146        let air_idx = AB::Expr::from_usize(DEF_CIRCUIT_PVS_AIR_ID);
147
148        for (pv_idx, value) in local.child_pvs.output_commit.iter().enumerate() {
149            self.public_values_bus.receive(
150                builder,
151                local.proof_idx,
152                PublicValuesBusMessage {
153                    air_idx: air_idx.clone(),
154                    pv_idx: AB::Expr::from_usize(pv_idx + DIGEST_SIZE),
155                    value: (*value).into(),
156                },
157                is_present_leaf.clone(),
158            );
159        }
160
161        self.public_values_bus.receive(
162            builder,
163            local.proof_idx,
164            PublicValuesBusMessage {
165                air_idx,
166                pv_idx: AB::Expr::from_usize(2 * DIGEST_SIZE),
167                value: local.child_pvs.def_idx.into(),
168            },
169            is_present_leaf.clone(),
170        );
171
172        self.input_or_merkle_commit_bus.receive(
173            builder,
174            local.proof_idx,
175            InputOrMerkleCommitMessage {
176                has_verifier_pvs: local.has_verifier_pvs.into(),
177                commit: from_fn(|i| {
178                    is_leaf.clone() * local.child_pvs.input_commit[i]
179                        + is_internal * local.merkle_commit[i]
180                }),
181            },
182            local.is_present,
183        );
184
185        /*
186         * On the leaf layer we need to constrain that merkle_commit is the Poseidon2 compression
187         * of the tagged input_commit and output_commit.
188         */
189        self.poseidon2_bus.lookup_key(
190            builder,
191            Poseidon2CompressMessage {
192                input: digests_to_poseidon2_input(
193                    DEF_LEAF_TAG.map(AB::Expr::from_u8),
194                    local.child_pvs.input_commit.map(Into::into),
195                ),
196                output: local.tagged_input_commit.map(Into::into),
197            },
198            is_present_leaf.clone(),
199        );
200
201        self.poseidon2_bus.lookup_key(
202            builder,
203            Poseidon2CompressMessage {
204                input: digests_to_poseidon2_input(
205                    local.tagged_input_commit.map(Into::into),
206                    local.child_pvs.output_commit.map(Into::into),
207                ),
208                output: local.merkle_commit.map(Into::into),
209            },
210            is_present_leaf,
211        );
212
213        /*
214         * We want to ensure consistency between AIRs that process public values, and we do so
215         * using the def_pvs_consistency_bus.
216         */
217        self.def_pvs_consistency_bus.receive(
218            builder,
219            local.proof_idx,
220            DefPvsConsistencyMessage {
221                has_verifier_pvs: local.has_verifier_pvs,
222            },
223            local.is_present,
224        );
225
226        /*
227         * On internal layers we receive num_def_circuit_proofs, merkle_depth, and def_idx from
228         * PublicValuesBus. To save columns, we repurpose the first two elements of the unused
229         * local.child_pvs.input_commit for num_def_circuit_proofs and merkle_depth.
230         */
231        let local_num_def_circuit_proofs = local.child_pvs.input_commit[0];
232        let local_merkle_depth = local.child_pvs.input_commit[1];
233
234        let mut when_dummy_internal = builder.when(not(local.is_present) * is_internal);
235        when_dummy_internal.assert_zero(local_num_def_circuit_proofs);
236        when_dummy_internal.assert_zero(local_merkle_depth);
237
238        self.public_values_bus.receive(
239            builder,
240            local.proof_idx,
241            PublicValuesBusMessage {
242                air_idx: AB::Expr::from_usize(DEF_AGG_PVS_AIR_ID),
243                pv_idx: AB::Expr::from_usize(DIGEST_SIZE),
244                value: local_num_def_circuit_proofs.into(),
245            },
246            is_internal * local.is_present,
247        );
248
249        self.public_values_bus.receive(
250            builder,
251            local.proof_idx,
252            PublicValuesBusMessage {
253                air_idx: AB::Expr::from_usize(DEF_AGG_PVS_AIR_ID),
254                pv_idx: AB::Expr::from_usize(DIGEST_SIZE + 1),
255                value: local_merkle_depth.into(),
256            },
257            is_internal * local.is_present,
258        );
259
260        self.public_values_bus.receive(
261            builder,
262            local.proof_idx,
263            PublicValuesBusMessage {
264                air_idx: AB::Expr::from_usize(DEF_AGG_PVS_AIR_ID),
265                pv_idx: AB::Expr::from_usize(DIGEST_SIZE + 2),
266                value: local.child_pvs.def_idx.into(),
267            },
268            is_internal * local.is_present,
269        );
270
271        /*
272         * If there is only one row, then this proof is a wrapper and its public values should
273         * match those of its child's. Otherwise, constrain that num_def_circuit_proofs is the
274         * sum of the present children's and that merkle_depth is the child merkle_depth + 1.
275         * If both children are present, we constrain that their merkle_depth values are equal.
276         * Every present child's def_idx should match this proof's.
277         */
278        let &DeferralAggregationPvs::<_> {
279            merkle_commit,
280            num_def_circuit_proofs,
281            merkle_depth,
282            def_idx,
283        } = builder.public_values().borrow();
284
285        let is_first = not(local.proof_idx);
286        let is_first_of_two_rows = is_first.clone() * next.proof_idx;
287
288        let local_num_proofs =
289            local.is_present * is_leaf.clone() + is_internal * local_num_def_circuit_proofs;
290        let next_num_proofs =
291            next.is_present * is_leaf + is_internal * next.child_pvs.input_commit[0];
292        let local_merkle_depth = is_internal * local_merkle_depth;
293        let next_merkle_depth = is_internal * next.child_pvs.input_commit[1];
294
295        let mut when_one_row = builder.when(is_first.clone() * not(next.proof_idx));
296        when_one_row.assert_eq(local_num_proofs.clone(), num_def_circuit_proofs);
297        when_one_row.assert_eq(local_merkle_depth.clone(), merkle_depth);
298        assert_array_eq(&mut when_one_row, local.merkle_commit, merkle_commit);
299
300        builder
301            .when_transition()
302            .assert_eq(local_num_proofs + next_num_proofs, num_def_circuit_proofs);
303        builder
304            .when_transition()
305            .assert_eq(local_merkle_depth.clone() + AB::Expr::ONE, merkle_depth);
306        builder
307            .when_transition()
308            .when(next.is_present)
309            .assert_eq(local_merkle_depth, next_merkle_depth);
310        builder
311            .when(local.is_present)
312            .assert_eq(local.child_pvs.def_idx, def_idx);
313
314        /*
315         * MAX_DEF_AGG_MERKLE_DEPTH - merkle_depth is range checked to be in [0, 256). This,
316         * combined with the constraints above that require at least one valid child proof
317         * with merkle_depth - 1 for any merkle_depth > 0, ensures that merkle_depth does not
318         * overflow F. Inductively num_def_circuit_proofs is bounded by 2^merkle_depth, so it
319         * also cannot wrap in F.
320         */
321        self.range_bus.lookup_key(
322            builder,
323            RangeCheckerBusMessage {
324                value: AB::Expr::from_usize(MAX_DEF_AGG_MERKLE_DEPTH) - merkle_depth.into(),
325                max_bits: AB::Expr::from_u8(8),
326            },
327            is_first,
328        );
329
330        /*
331         * If there are two rows, then the second (next) is either present or not. If present
332         * the parent merkle_hash should be the compression of both children's merkle_commit.
333         * If not the second row's merkle_commit columns are repurposed into encoder flags for
334         * the child merkle depth. They're used to select the correct zero hash depth.
335         */
336        let flags = next
337            .merkle_commit
338            .into_iter()
339            .take(self.encoder.width())
340            .collect_vec();
341        self.encoder
342            .eval(&mut builder.when(not(next.is_present)), &flags);
343
344        let mut zero_hash = [AB::Expr::ZERO; DIGEST_SIZE];
345        let mut encoded_parent_depth = AB::Expr::ZERO;
346
347        for i in 0..=MAX_DEF_AGG_MERKLE_DEPTH {
348            let is_current = self.encoder.get_flag_expr::<AB>(i, &flags);
349            let current_zero_hash: [AB::F; DIGEST_SIZE] = self.zero_hashes[i].into();
350            for j in 0..DIGEST_SIZE {
351                zero_hash[j] += is_current.clone() * current_zero_hash[j];
352            }
353            encoded_parent_depth += is_current * AB::Expr::from_usize(i);
354        }
355
356        builder
357            .when(not(next.is_present))
358            .assert_eq(merkle_depth, encoded_parent_depth + AB::Expr::ONE);
359
360        let right_child = from_fn(|i| {
361            next.is_present * next.merkle_commit[i] + not(next.is_present) * zero_hash[i].clone()
362        });
363
364        self.poseidon2_bus.lookup_key(
365            builder,
366            Poseidon2CompressMessage {
367                input: digests_to_poseidon2_input(
368                    DEF_INTERNAL_TAG.map(AB::Expr::from_u8),
369                    local.merkle_commit.map(Into::into),
370                ),
371                output: local.tagged_left_merkle.map(Into::into),
372            },
373            is_first_of_two_rows.clone(),
374        );
375
376        self.poseidon2_bus.lookup_key(
377            builder,
378            Poseidon2CompressMessage {
379                input: digests_to_poseidon2_input(
380                    local.tagged_left_merkle.map(Into::into),
381                    right_child,
382                ),
383                output: merkle_commit.map(Into::into),
384            },
385            is_first_of_two_rows,
386        );
387    }
388}
openvm_continuations/circuit/deferral/inner/def_pvs/air.rs

openvm_continuations/circuit/deferral/inner/def_pvs/
air.rs
