openvm_recursion_circuit/stacking/sumcheck/
air.rs

1use std::borrow::Borrow;
2
3use openvm_circuit_primitives::{
4    utils::{and, assert_array_eq, not},
5    ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
6};
7use openvm_recursion_circuit_derive::AlignedBorrow;
8use openvm_stark_backend::{
9    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
10};
11use openvm_stark_sdk::config::baby_bear_poseidon2::{D_EF, F};
12use p3_air::{Air, AirBuilder, BaseAir};
13use p3_field::{extension::BinomiallyExtendable, Field, PrimeCharacteristicRing, PrimeField32};
14use p3_matrix::Matrix;
15
16use crate::{
17    bus::{
18        ConstraintSumcheckRandomness, ConstraintSumcheckRandomnessBus, TranscriptBus,
19        TranscriptBusMessage, WhirOpeningPointBus, WhirOpeningPointMessage,
20    },
21    stacking::bus::{
22        EqBaseBus, EqBaseMessage, EqKernelLookupBus, EqKernelLookupMessage, EqRandValuesLookupBus,
23        EqRandValuesLookupMessage, StackingModuleTidxBus, StackingModuleTidxMessage,
24        SumcheckClaimsBus, SumcheckClaimsMessage,
25    },
26    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
27    utils::{
28        assert_zeros, ext_field_add, ext_field_multiply, ext_field_multiply_scalar,
29        ext_field_one_minus, ext_field_subtract,
30    },
31};
32
33#[repr(C)]
34#[derive(AlignedBorrow, StructReflection)]
35pub struct SumcheckRoundsCols<F> {
36    // Proof index columns for continuations
37    pub proof_idx: F,
38    pub is_valid: F,
39    pub is_first: F,
40    pub is_last: F,
41
42    // Sumcheck round this row represents
43    pub round: F,
44
45    // Starting tidx for this sumcheck round
46    pub tidx: F,
47
48    // Evaluations of polynomial s_round
49    pub s_eval_at_0: [F; D_EF],
50    pub s_eval_at_1: [F; D_EF],
51    pub s_eval_at_2: [F; D_EF],
52    pub s_eval_at_u: [F; D_EF],
53
54    // Values of sampled u and r for this round
55    pub u_round: [F; D_EF],
56    pub r_round: [F; D_EF],
57    pub has_r: F,
58    pub u_mult: F,
59
60    // Values of eq(u_0, r_0), eq(u_0, r_0 * omega), and eq(u_0, 1) * eq(r_0 * omega, 1)
61    pub eq_prism_base: [F; D_EF],
62    pub eq_cube_base: [F; D_EF],
63    pub rot_cube_base: [F; D_EF],
64
65    // Value of eq_cube
66    pub eq_cube: [F; D_EF],
67
68    // Intermediate values to compute rot_cube recursively
69    pub r_not_u_prod: [F; D_EF],
70    pub rot_cube_minus_prod: [F; D_EF],
71
72    // Multiplicity of eq_round(u, r) lookup
73    pub eq_rot_mult: F,
74}
75
76#[derive(ColumnsAir)]
77#[columns_via(SumcheckRoundsCols<u8>)]
78pub struct SumcheckRoundsAir {
79    // External buses
80    pub constraint_randomness_bus: ConstraintSumcheckRandomnessBus,
81    pub whir_opening_point_bus: WhirOpeningPointBus,
82    pub transcript_bus: TranscriptBus,
83
84    // Internal buses
85    pub stacking_tidx_bus: StackingModuleTidxBus,
86    pub sumcheck_claims_bus: SumcheckClaimsBus,
87    pub eq_base_bus: EqBaseBus,
88    pub eq_rand_values_bus: EqRandValuesLookupBus,
89    pub eq_kernel_lookup_bus: EqKernelLookupBus,
90
91    pub l_skip: usize,
92}
93
94impl BaseAirWithPublicValues<F> for SumcheckRoundsAir {}
95impl PartitionedBaseAir<F> for SumcheckRoundsAir {}
96
97impl<F> BaseAir<F> for SumcheckRoundsAir {
98    fn width(&self) -> usize {
99        SumcheckRoundsCols::<F>::width()
100    }
101}
102
103impl<AB: AirBuilder + InteractionBuilder> Air<AB> for SumcheckRoundsAir
104where
105    AB::F: PrimeField32,
106    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
107{
108    fn eval(&self, builder: &mut AB) {
109        let main = builder.main();
110        let (local, next) = (
111            main.row_slice(0).expect("window should have two elements"),
112            main.row_slice(1).expect("window should have two elements"),
113        );
114
115        let local: &SumcheckRoundsCols<AB::Var> = (*local).borrow();
116        let next: &SumcheckRoundsCols<AB::Var> = (*next).borrow();
117
118        NestedForLoopSubAir::<1> {}.eval(
119            builder,
120            (
121                NestedForLoopIoCols {
122                    is_enabled: local.is_valid,
123                    counter: [local.proof_idx],
124                    is_first: [local.is_first],
125                }
126                .map_into(),
127                NestedForLoopIoCols {
128                    is_enabled: next.is_valid,
129                    counter: [next.proof_idx],
130                    is_first: [next.is_first],
131                }
132                .map_into(),
133            ),
134        );
135
136        builder.when(local.is_valid).assert_eq(
137            local.is_last,
138            NestedForLoopSubAir::<1>::local_is_last(local.is_valid, next.is_valid, next.is_first),
139        );
140
141        builder.assert_bool(local.is_last);
142        builder
143            .when(and(local.is_valid, local.is_last))
144            .assert_zero((local.proof_idx + AB::F::ONE - next.proof_idx) * next.proof_idx);
145        builder
146            .when(and(not(local.is_valid), local.is_last))
147            .assert_zero(next.proof_idx);
148
149        /*
150         * Constrain that round increments correctly.
151         */
152        builder.when(local.is_first).assert_one(local.round);
153        builder
154            .when(and(not(local.is_last), local.is_valid))
155            .assert_eq(local.round + AB::Expr::ONE, next.round);
156
157        /*
158         * Constrain that s_round(u_round) is the quadratic interpolation using values
159         * s_round(0), s_round(1), and s_round(2). Additionally, constrain that we have
160         * s_round(u_round) = s_{round + 1}(0) + s_{round + 1}(1), and send the value of
161         * s_{n_stack}(u_{n_stack}) to StackingClaimsAir.
162         */
163        self.sumcheck_claims_bus.receive(
164            builder,
165            local.proof_idx,
166            SumcheckClaimsMessage {
167                module_idx: AB::Expr::ONE,
168                value: ext_field_add(local.s_eval_at_0, local.s_eval_at_1),
169            },
170            local.is_first,
171        );
172
173        let s1 = ext_field_subtract(local.s_eval_at_1, local.s_eval_at_0);
174        let s2 = ext_field_subtract(local.s_eval_at_2, local.s_eval_at_1);
175        let p = ext_field_multiply_scalar::<AB::Expr>(
176            ext_field_subtract::<AB::Expr>(s2, s1.clone()),
177            AB::F::TWO.inverse(),
178        );
179        let q = ext_field_subtract::<AB::Expr>(s1, p.clone());
180
181        assert_array_eq(
182            builder,
183            ext_field_add(
184                ext_field_multiply(
185                    ext_field_add::<AB::Expr>(ext_field_multiply(p, local.u_round), q),
186                    local.u_round,
187                ),
188                local.s_eval_at_0,
189            ),
190            local.s_eval_at_u,
191        );
192
193        assert_array_eq(
194            &mut builder.when(not(local.is_last)),
195            local.s_eval_at_u,
196            ext_field_add(next.s_eval_at_0, next.s_eval_at_1),
197        );
198
199        self.sumcheck_claims_bus.send(
200            builder,
201            local.proof_idx,
202            SumcheckClaimsMessage {
203                module_idx: AB::Expr::TWO,
204                value: local.s_eval_at_u.map(Into::into),
205            },
206            and(local.is_last, local.is_valid),
207        );
208
209        /*
210         * Constrain the correctness of eq_cube and rot_cube at each round and provide the
211         * lookups for eq_round(u, r) and k_rot_round(u, r). Computing rot_cube recursively
212         * requires us to store the prefix product of r_round * (1 - u_round), which we
213         * denote r_not_u_prod, and rot_cube - r_not_u_prod.
214         */
215        self.eq_base_bus.receive(
216            builder,
217            local.proof_idx,
218            EqBaseMessage {
219                eq_u_r: local.eq_prism_base,
220                eq_u_r_omega: local.eq_cube_base,
221                eq_u_r_prod: local.rot_cube_base,
222            },
223            local.is_first,
224        );
225
226        assert_array_eq(
227            &mut builder.when(not(local.is_last)),
228            local.eq_prism_base,
229            next.eq_prism_base,
230        );
231
232        assert_array_eq(
233            &mut builder.when(not(local.is_last)),
234            local.eq_cube_base,
235            next.eq_cube_base,
236        );
237
238        assert_array_eq(
239            &mut builder.when(not(local.is_last)),
240            local.rot_cube_base,
241            next.rot_cube_base,
242        );
243
244        let local_u_not_r = ext_field_multiply(local.u_round, ext_field_one_minus(local.r_round));
245        let local_r_not_u = ext_field_multiply(local.r_round, ext_field_one_minus(local.u_round));
246        let next_u_not_r = ext_field_multiply(next.u_round, ext_field_one_minus(next.r_round));
247        let next_r_not_u = ext_field_multiply(next.r_round, ext_field_one_minus(next.u_round));
248
249        assert_array_eq(
250            &mut builder.when(local.is_first),
251            ext_field_one_minus::<AB::Expr>(ext_field_add::<AB::Expr>(
252                local_u_not_r.clone(),
253                local_r_not_u.clone(),
254            )),
255            local.eq_cube,
256        );
257
258        assert_array_eq(
259            &mut builder.when(not(local.is_last)),
260            ext_field_multiply(
261                local.eq_cube,
262                ext_field_one_minus::<AB::Expr>(ext_field_add::<AB::Expr>(
263                    next_u_not_r.clone(),
264                    next_r_not_u.clone(),
265                )),
266            ),
267            next.eq_cube,
268        );
269
270        assert_array_eq(
271            &mut builder.when(local.is_first),
272            local.r_not_u_prod,
273            local_r_not_u,
274        );
275
276        assert_array_eq(
277            &mut builder.when(not(local.is_last)),
278            ext_field_multiply(local.r_not_u_prod, next_r_not_u.clone()),
279            next.r_not_u_prod,
280        );
281
282        assert_array_eq(
283            &mut builder.when(local.is_first),
284            local.rot_cube_minus_prod,
285            local_u_not_r,
286        );
287
288        assert_array_eq(
289            &mut builder.when(not(local.is_last)),
290            ext_field_add::<AB::Expr>(
291                ext_field_multiply(
292                    local.rot_cube_minus_prod,
293                    ext_field_one_minus::<AB::Expr>(ext_field_add::<AB::Expr>(
294                        next_u_not_r.clone(),
295                        next_r_not_u,
296                    )),
297                ),
298                ext_field_multiply(next_u_not_r, local.r_not_u_prod),
299            ),
300            next.rot_cube_minus_prod,
301        );
302
303        self.eq_kernel_lookup_bus.add_key_with_lookups(
304            builder,
305            local.proof_idx,
306            EqKernelLookupMessage {
307                n: local.round.into(),
308                eq_in: ext_field_multiply(local.eq_prism_base, local.eq_cube),
309                k_rot_in: ext_field_add(
310                    ext_field_multiply(local.eq_cube_base, local.eq_cube),
311                    ext_field_multiply(
312                        local.rot_cube_base,
313                        ext_field_subtract(
314                            ext_field_add(local.r_not_u_prod, local.rot_cube_minus_prod),
315                            local.eq_cube,
316                        ),
317                    ),
318                ),
319            },
320            local.is_valid * local.eq_rot_mult,
321        );
322
323        builder.assert_bool(local.has_r);
324        builder
325            .when(and(local.is_valid, not(local.is_last)))
326            .assert_bool(local.has_r - next.has_r);
327        builder
328            .when(not(local.has_r))
329            .assert_zero(local.eq_rot_mult);
330
331        assert_zeros(&mut builder.when(not(local.has_r)), local.r_round);
332
333        /*
334         * Because we sample u_round and r_round from the transcript here, we send
335         * them to other AIRs that need to use it.
336         */
337        self.constraint_randomness_bus.receive(
338            builder,
339            local.proof_idx,
340            ConstraintSumcheckRandomness {
341                idx: local.round,
342                challenge: local.r_round,
343            },
344            and(local.is_valid, local.has_r),
345        );
346
347        self.whir_opening_point_bus.send(
348            builder,
349            local.proof_idx,
350            WhirOpeningPointMessage {
351                idx: local.round + AB::Expr::from_usize(self.l_skip - 1),
352                value: local.u_round.map(Into::into),
353            },
354            local.is_valid,
355        );
356
357        self.eq_rand_values_bus.add_key_with_lookups(
358            builder,
359            local.proof_idx,
360            EqRandValuesLookupMessage {
361                idx: local.round,
362                u: local.u_round,
363            },
364            local.u_mult,
365        );
366        builder.when(not(local.is_valid)).assert_zero(local.u_mult);
367
368        /*
369         * Constrain transcript operations and send the final tidx to StackingClaimsAir.
370         */
371        self.stacking_tidx_bus.receive(
372            builder,
373            local.proof_idx,
374            StackingModuleTidxMessage {
375                module_idx: AB::Expr::ONE,
376                tidx: local.tidx.into(),
377            },
378            local.is_first,
379        );
380
381        builder
382            .when(not(local.is_last) * local.is_valid)
383            .assert_eq(local.tidx + AB::F::from_usize(3 * D_EF), next.tidx);
384
385        for i in 0..D_EF {
386            self.transcript_bus.receive(
387                builder,
388                local.proof_idx,
389                TranscriptBusMessage {
390                    tidx: AB::Expr::from_usize(i) + local.tidx,
391                    value: local.s_eval_at_1[i].into(),
392                    is_sample: AB::Expr::ZERO,
393                },
394                local.is_valid,
395            );
396
397            self.transcript_bus.receive(
398                builder,
399                local.proof_idx,
400                TranscriptBusMessage {
401                    tidx: AB::Expr::from_usize(i + D_EF) + local.tidx,
402                    value: local.s_eval_at_2[i].into(),
403                    is_sample: AB::Expr::ZERO,
404                },
405                local.is_valid,
406            );
407
408            self.transcript_bus.receive(
409                builder,
410                local.proof_idx,
411                TranscriptBusMessage {
412                    tidx: AB::Expr::from_usize(i + 2 * D_EF) + local.tidx,
413                    value: local.u_round[i].into(),
414                    is_sample: AB::Expr::ONE,
415                },
416                local.is_valid,
417            );
418        }
419
420        self.stacking_tidx_bus.send(
421            builder,
422            local.proof_idx,
423            StackingModuleTidxMessage {
424                module_idx: AB::Expr::TWO,
425                tidx: AB::Expr::from_usize(3 * D_EF) + local.tidx,
426            },
427            and(local.is_last, local.is_valid),
428        );
429    }
430}
openvm_recursion_circuit/stacking/sumcheck/air.rs

openvm_recursion_circuit/stacking/sumcheck/
air.rs
