openvm_continuations/circuit/inner/def_pvs/
air.rs

1use std::{array::from_fn, borrow::Borrow};
2
3use openvm_circuit_primitives::{
4    utils::{assert_array_eq, not},
5    ColumnsAir, StructReflection, StructReflectionHelper,
6};
7use openvm_recursion_circuit::{
8    bus::{
9        CachedCommitBus, CachedCommitBusMessage, Poseidon2CompressBus, Poseidon2CompressMessage,
10        PublicValuesBus, PublicValuesBusMessage,
11    },
12    primitives::bus::{RangeCheckerBus, RangeCheckerBusMessage},
13};
14use openvm_recursion_circuit_derive::AlignedBorrow;
15use openvm_stark_backend::{
16    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
17};
18use openvm_verify_stark_host::pvs::{
19    DeferralPvs, CONSTRAINT_EVAL_AIR_ID, CONSTRAINT_EVAL_CACHED_INDEX, DEF_PVS_AIR_ID,
20};
21use p3_air::{Air, AirBuilder, AirBuilderWithPublicValues, BaseAir};
22use p3_field::PrimeCharacteristicRing;
23use p3_matrix::Matrix;
24
25use crate::{
26    circuit::{
27        deferral::DEF_HOOK_PVS_AIR_ID,
28        inner::bus::{PvsAirConsistencyBus, PvsAirConsistencyMessage},
29    },
30    utils::digests_to_poseidon2_input,
31    CommitBytes,
32};
33
34#[repr(C)]
35#[derive(AlignedBorrow, StructReflection)]
36pub struct DeferralPvsCols<F> {
37    pub row_idx: F,
38    pub deferral_flag: F,
39    pub has_verifier_pvs: F,
40
41    pub proof_idx: F,
42    pub is_present: F,
43    pub single_present_is_right: F,
44
45    pub child_pvs: DeferralPvs<F>,
46}
47
48#[derive(ColumnsAir)]
49#[columns_via(DeferralPvsCols<u8>)]
50pub struct DeferralPvsAir {
51    pub public_values_bus: PublicValuesBus,
52    pub cached_commit_bus: CachedCommitBus,
53    pub poseidon2_bus: Poseidon2CompressBus,
54    pub range_bus: RangeCheckerBus,
55    pub pvs_air_consistency_bus: PvsAirConsistencyBus,
56
57    pub expected_def_hook_cached_commit: CommitBytes,
58}
59
60impl<F> BaseAir<F> for DeferralPvsAir {
61    fn width(&self) -> usize {
62        DeferralPvsCols::<u8>::width()
63    }
64}
65impl<F> BaseAirWithPublicValues<F> for DeferralPvsAir {
66    fn num_public_values(&self) -> usize {
67        DeferralPvs::<u8>::width()
68    }
69}
70impl<F> PartitionedBaseAir<F> for DeferralPvsAir {}
71
72impl<AB: AirBuilder + InteractionBuilder + AirBuilderWithPublicValues> Air<AB> for DeferralPvsAir {
73    fn eval(&self, builder: &mut AB) {
74        let main = builder.main();
75        let (local, next) = (
76            main.row_slice(0).expect("window should have two elements"),
77            main.row_slice(1).expect("window should have two elements"),
78        );
79        let local: &DeferralPvsCols<AB::Var> = (*local).borrow();
80        let next: &DeferralPvsCols<AB::Var> = (*next).borrow();
81
82        /*
83         * This AIR may have 1 or 2 rows. There are 4 valid 1-row cases:
84         * - deferral_flag == 0: child deferral pvs are unset
85         * - deferral_flag == 1 && proof_idx == 0: wrapping a deferral proof
86         * - deferral_flag == 1 && proof_idx == 1: combining a VM and deferral proof
87         * - deferral_flag == 2: wrapping a combined proof
88         *
89         * There are 2 valid 2-row cases, both with deferral_flag == 1:
90         * - Both child proofs are present
91         * - The first proof is present and the second is absent
92         */
93        // constrain that when hash_pvs is set we have exactly 2 def rows
94        builder.assert_bool(local.row_idx);
95        builder.when_first_row().assert_zero(local.row_idx);
96        builder
97            .when_transition()
98            .assert_one(next.row_idx - local.row_idx);
99
100        let has_two_rows = (next.row_idx - local.row_idx).square();
101        let has_one_row = not::<AB::Expr>(has_two_rows.clone());
102
103        // constrain that all the present rows are at the beginning
104        builder.assert_bool(local.is_present);
105        builder
106            .when_transition()
107            .assert_bool(local.is_present - next.is_present);
108
109        // constrain if deferral_flag is set, there is at least one present proof
110        builder.assert_tern(local.deferral_flag);
111        builder.assert_eq(local.deferral_flag, next.deferral_flag);
112        builder
113            .when(local.deferral_flag)
114            .assert_bool(local.is_present + next.is_present - AB::Expr::ONE);
115
116        // basic constraints for consistency columns
117        builder.when_first_row().assert_bool(local.proof_idx);
118        builder
119            .when_first_row()
120            .when(local.proof_idx)
121            .assert_one(has_one_row.clone());
122        builder.assert_bool(local.has_verifier_pvs);
123        builder.assert_eq(local.has_verifier_pvs, next.has_verifier_pvs);
124
125        // constrain single_present_is_right is set when there is 1 present and
126        // 1 absent row
127        builder.assert_bool(local.single_present_is_right);
128        builder.assert_eq(local.single_present_is_right, next.single_present_is_right);
129        builder
130            .when(local.single_present_is_right)
131            .assert_one(local.is_present + next.is_present);
132
133        /*
134         * When deferral_flag is unset, there must be a single row with zeros for
135         * public values.
136         */
137        let mut when_flag_not_one = builder.when_ne(local.deferral_flag, AB::Expr::ONE);
138        when_flag_not_one.assert_one(has_one_row.clone());
139
140        let mut when_invalid = when_flag_not_one.when_ne(local.deferral_flag, AB::Expr::TWO);
141        when_invalid.assert_zero(local.is_present);
142        for child_pv in local.child_pvs.as_slice() {
143            when_invalid.assert_zero(*child_pv);
144        }
145
146        /*
147         * If there are two rows and a proof is absent, it represents an accumulator
148         * Merkle subtree that has been left untouched. We constrain its initial and
149         * final accumulator hashes to be equal. Additionally, if there are two rows
150         * then the child_pvs depth should be equal.
151         */
152        assert_array_eq(
153            &mut builder
154                .when(has_two_rows.clone())
155                .when(not(local.is_present)),
156            local.child_pvs.initial_acc_hash,
157            local.child_pvs.final_acc_hash,
158        );
159
160        builder
161            .when(has_two_rows.clone())
162            .assert_eq(local.child_pvs.depth, next.child_pvs.depth);
163
164        /*
165         * If this row is present then we need to receive the child public values
166         * from ProofShapeModule. At the hook level this is at DEF_HOOK_PVS_AIR_ID,
167         * at every other level it will be at DEF_PVS_AIR_ID.
168         */
169        let def_pvs_air_idx = AB::Expr::from_usize(DEF_PVS_AIR_ID) * local.has_verifier_pvs
170            + AB::Expr::from_usize(DEF_HOOK_PVS_AIR_ID) * not(local.has_verifier_pvs);
171        for (pv_idx, value) in local.child_pvs.as_slice().iter().enumerate() {
172            self.public_values_bus.receive(
173                builder,
174                local.proof_idx,
175                PublicValuesBusMessage {
176                    air_idx: def_pvs_air_idx.clone(),
177                    pv_idx: AB::Expr::from_usize(pv_idx),
178                    value: (*value).into(),
179                },
180                local.is_present,
181            );
182        }
183
184        /*
185         * We look up proof metadata from VerifierPvsAir here to ensure consistency
186         * on each row.
187         */
188        self.pvs_air_consistency_bus.lookup_key(
189            builder,
190            local.proof_idx,
191            PvsAirConsistencyMessage {
192                deferral_flag: local.deferral_flag,
193                has_verifier_pvs: local.has_verifier_pvs,
194            },
195            local.is_present,
196        );
197
198        /*
199         * If this row corresponds to a direct deferral hook circuit child (i.e.
200         * has_verifier_pvs == 0), receive the child's cached trace commit and
201         * constrain it to an expected constant.
202         */
203        self.cached_commit_bus.receive(
204            builder,
205            local.proof_idx,
206            CachedCommitBusMessage {
207                air_idx: AB::Expr::from_usize(CONSTRAINT_EVAL_AIR_ID),
208                cached_idx: AB::Expr::from_usize(CONSTRAINT_EVAL_CACHED_INDEX),
209                global_cached_idx: AB::Expr::ZERO,
210                cached_commit: self.expected_def_hook_cached_commit.into(),
211            },
212            local.is_present * not(local.has_verifier_pvs),
213        );
214
215        /*
216         * Finally, we constrain the public values to be consistent with the
217         * child's. If there is one row then the pvs are simply passed through
218         * and node_idx must be 0.
219         */
220        let &DeferralPvs::<_> {
221            initial_acc_hash,
222            final_acc_hash,
223            depth,
224            node_idx,
225        } = builder.public_values().borrow();
226
227        // constrain that pvs are passed through if there is one row
228        let mut when_one_row = builder.when(has_one_row);
229        when_one_row.assert_eq(local.child_pvs.depth, depth);
230        when_one_row.assert_eq(local.child_pvs.node_idx, node_idx);
231        when_one_row.assert_zero(node_idx);
232
233        assert_array_eq(
234            &mut when_one_row,
235            local.child_pvs.initial_acc_hash,
236            initial_acc_hash,
237        );
238        assert_array_eq(
239            &mut when_one_row,
240            local.child_pvs.final_acc_hash,
241            final_acc_hash,
242        );
243
244        /*
245         * If there are two rows, then initial_acc_hash and final_acc_hash are
246         * combined, depth is incremented by 1, and child node indices are
247         * constrained to be adjacent. The left child node_idx must be even,
248         * and the parent node_idx is the left child node_idx divided by 2.
249         * Note that we range check the parent node_idx to be in [0, 256),
250         * which forces each def_idx into [0, 512).
251         */
252        let is_first_of_two_rows = next.row_idx;
253        let single_present_is_left = not(local.single_present_is_right);
254
255        let left_init_child = from_fn(|i| {
256            single_present_is_left.clone() * local.child_pvs.initial_acc_hash[i]
257                + local.single_present_is_right * next.child_pvs.initial_acc_hash[i]
258        });
259        let right_init_child = from_fn(|i| {
260            local.single_present_is_right * local.child_pvs.initial_acc_hash[i]
261                + single_present_is_left.clone() * next.child_pvs.initial_acc_hash[i]
262        });
263
264        self.poseidon2_bus.lookup_key(
265            builder,
266            Poseidon2CompressMessage {
267                input: digests_to_poseidon2_input(left_init_child, right_init_child),
268                output: initial_acc_hash.map(Into::into),
269            },
270            is_first_of_two_rows,
271        );
272
273        let left_final_child = from_fn(|i| {
274            single_present_is_left.clone() * local.child_pvs.final_acc_hash[i]
275                + local.single_present_is_right * next.child_pvs.final_acc_hash[i]
276        });
277        let right_final_child = from_fn(|i| {
278            local.single_present_is_right * local.child_pvs.final_acc_hash[i]
279                + single_present_is_left.clone() * next.child_pvs.final_acc_hash[i]
280        });
281
282        self.poseidon2_bus.lookup_key(
283            builder,
284            Poseidon2CompressMessage {
285                input: digests_to_poseidon2_input(left_final_child, right_final_child),
286                output: final_acc_hash.map(Into::into),
287            },
288            is_first_of_two_rows,
289        );
290
291        builder
292            .when(has_two_rows)
293            .assert_one(depth.into() - local.child_pvs.depth);
294
295        let canonical_child_node_idx =
296            AB::Expr::TWO * node_idx.into() + local.single_present_is_right;
297
298        builder
299            .when(is_first_of_two_rows)
300            .when(next.is_present)
301            .assert_one(next.child_pvs.node_idx - local.child_pvs.node_idx);
302        builder
303            .when(is_first_of_two_rows)
304            .assert_eq(local.child_pvs.node_idx, canonical_child_node_idx);
305
306        self.range_bus.lookup_key(
307            builder,
308            RangeCheckerBusMessage {
309                value: node_idx.into(),
310                max_bits: AB::Expr::from_u8(8),
311            },
312            is_first_of_two_rows,
313        );
314    }
315}
openvm_continuations/circuit/inner/def_pvs/air.rs

openvm_continuations/circuit/inner/def_pvs/
air.rs
