openvm_recursion_circuit/batch_constraint/sumcheck/univariate/
air.rs

1use std::borrow::Borrow;
2
3use openvm_circuit_primitives::{
4    is_equal::{IsEqSubAir, IsEqualAuxCols, IsEqualIo},
5    utils::{assert_array_eq, not},
6    ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
7};
8use openvm_recursion_circuit_derive::AlignedBorrow;
9use openvm_stark_backend::{
10    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
11};
12use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF;
13use p3_air::{Air, AirBuilder, BaseAir};
14use p3_field::{extension::BinomiallyExtendable, PrimeCharacteristicRing, TwoAdicField};
15use p3_matrix::Matrix;
16
17use crate::{
18    batch_constraint::bus::{
19        BatchConstraintConductorBus, BatchConstraintConductorMessage,
20        BatchConstraintInnerMessageType, SumcheckClaimBus, SumcheckClaimMessage,
21        UnivariateSumcheckInputBus, UnivariateSumcheckInputMessage,
22    },
23    bus::{
24        ConstraintSumcheckRandomness, ConstraintSumcheckRandomnessBus, StackingModuleBus,
25        StackingModuleMessage, TranscriptBus,
26    },
27    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
28    utils::{assert_zeros, ext_field_add, ext_field_multiply, ext_field_multiply_scalar},
29};
30
31#[derive(AlignedBorrow, Clone, Copy, Debug, StructReflection)]
32#[repr(C)]
33pub struct UnivariateSumcheckCols<T> {
34    pub is_valid: T,
35    pub proof_idx: T,
36    pub is_first: T,
37
38    pub coeff_idx: T,
39
40    /// Powers of generator of order 2^{l_skip} to get periodic selector columns
41    pub omega_skip_power: T,
42    pub is_omega_skip_power_equal_to_one: T,
43    pub is_omega_skip_power_equal_to_one_aux: IsEqualAuxCols<T>,
44
45    pub coeff: [T; D_EF],
46    pub sum_at_roots: [T; D_EF],
47
48    pub r: [T; D_EF],
49    pub value_at_r: [T; D_EF],
50
51    pub tidx: T,
52}
53
54#[derive(ColumnsAir)]
55#[columns_via(UnivariateSumcheckCols<u8>)]
56pub struct UnivariateSumcheckAir {
57    /// The univariate domain size is `2^{l_skip}`
58    pub l_skip: usize,
59    /// The degree of the univariate polynomial
60    pub univariate_deg: usize,
61
62    pub univariate_sumcheck_input_bus: UnivariateSumcheckInputBus,
63    pub claim_bus: SumcheckClaimBus,
64    pub stacking_module_bus: StackingModuleBus,
65    pub transcript_bus: TranscriptBus,
66    pub randomness_bus: ConstraintSumcheckRandomnessBus,
67    pub batch_constraint_conductor_bus: BatchConstraintConductorBus,
68}
69
70impl<F> BaseAirWithPublicValues<F> for UnivariateSumcheckAir {}
71impl<F> PartitionedBaseAir<F> for UnivariateSumcheckAir {}
72
73impl<F> BaseAir<F> for UnivariateSumcheckAir {
74    fn width(&self) -> usize {
75        UnivariateSumcheckCols::<F>::width()
76    }
77}
78
79impl<AB: AirBuilder + InteractionBuilder> Air<AB> for UnivariateSumcheckAir
80where
81    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield:
82        BinomiallyExtendable<{ D_EF }> + TwoAdicField,
83{
84    fn eval(&self, builder: &mut AB) {
85        let main = builder.main();
86        let (local, next) = (
87            main.row_slice(0).expect("window should have two elements"),
88            main.row_slice(1).expect("window should have two elements"),
89        );
90        let local: &UnivariateSumcheckCols<AB::Var> = (*local).borrow();
91        let next: &UnivariateSumcheckCols<AB::Var> = (*next).borrow();
92
93        ///////////////////////////////////////////////////////////////////////
94        // Loop Constraints
95        ///////////////////////////////////////////////////////////////////////
96
97        type LoopSubAir = NestedForLoopSubAir<1>;
98        LoopSubAir {}.eval(
99            builder,
100            (
101                NestedForLoopIoCols {
102                    is_enabled: local.is_valid,
103                    counter: [local.proof_idx],
104                    is_first: [local.is_first],
105                }
106                .map_into(),
107                NestedForLoopIoCols {
108                    is_enabled: next.is_valid,
109                    counter: [next.proof_idx],
110                    is_first: [next.is_first],
111                }
112                .map_into(),
113            ),
114        );
115
116        let is_transition = LoopSubAir::local_is_transition(next.is_valid, next.is_first);
117        let is_last = LoopSubAir::local_is_last(local.is_valid, next.is_valid, next.is_first);
118
119        // Coeff index starts at univariate degree
120        builder
121            .when(local.is_first)
122            .assert_eq(local.coeff_idx, AB::Expr::from_usize(self.univariate_deg));
123        // Coeff index decrements by 1
124        builder
125            .when(is_transition.clone())
126            .assert_eq(next.coeff_idx, local.coeff_idx - AB::Expr::ONE);
127        // Coeff index ends at zero
128        builder.when(is_last.clone()).assert_zero(local.coeff_idx);
129
130        ///////////////////////////////////////////////////////////////////////
131        // Powers of omega constraints
132        ///////////////////////////////////////////////////////////////////////
133
134        let omega_skip = AB::Expr::from_prime_subfield(
135            <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield::two_adic_generator(self.l_skip),
136        );
137
138        // Omega power ends at 1
139        builder
140            .when(is_last.clone())
141            .assert_one(local.omega_skip_power);
142        // Powers of omega are calculated properly
143        builder
144            .when(is_transition.clone())
145            .assert_eq(local.omega_skip_power, next.omega_skip_power * omega_skip);
146
147        IsEqSubAir.eval(
148            builder,
149            (
150                IsEqualIo::new(
151                    local.omega_skip_power.into(),
152                    AB::Expr::ONE,
153                    local.is_omega_skip_power_equal_to_one.into(),
154                    local.is_valid.into(),
155                ),
156                local.is_omega_skip_power_equal_to_one_aux.inv,
157            ),
158        );
159
160        ///////////////////////////////////////////////////////////////////////
161        // Sum over Roots Constraints
162        ///////////////////////////////////////////////////////////////////////
163
164        let domain_size = AB::Expr::from_usize(1 << self.l_skip);
165
166        // Initialize sum over roots
167        assert_array_eq(
168            &mut builder
169                .when(local.is_first)
170                .when(local.is_omega_skip_power_equal_to_one),
171            local.sum_at_roots,
172            ext_field_multiply_scalar(local.coeff, domain_size.clone()),
173        );
174        assert_zeros(
175            &mut builder
176                .when(local.is_first)
177                .when(not(local.is_omega_skip_power_equal_to_one)),
178            local.sum_at_roots,
179        );
180        // Add c * 2^{l_skip} at every 2^{l_skip} coefficient
181        assert_array_eq(
182            &mut builder
183                .when(is_transition.clone())
184                .when(next.is_omega_skip_power_equal_to_one),
185            next.sum_at_roots,
186            ext_field_add(
187                local.sum_at_roots,
188                ext_field_multiply_scalar(next.coeff, domain_size.clone()),
189            ),
190        );
191        // Keep the sum over roots unchanged for other values
192        assert_array_eq(
193            &mut builder
194                .when(is_transition.clone())
195                .when(not(next.is_omega_skip_power_equal_to_one)),
196            next.sum_at_roots,
197            local.sum_at_roots,
198        );
199
200        ///////////////////////////////////////////////////////////////////////
201        // Horner evaluation at r
202        ///////////////////////////////////////////////////////////////////////
203
204        assert_array_eq(&mut builder.when(is_transition.clone()), next.r, local.r);
205
206        // Initialize evaluation
207        // e = c
208        assert_array_eq(
209            &mut builder.when(local.is_first),
210            local.value_at_r,
211            local.coeff,
212        );
213        // e' = c + r * e
214        assert_array_eq(
215            &mut builder.when(is_transition.clone()),
216            next.value_at_r,
217            ext_field_add(next.coeff, ext_field_multiply(next.r, local.value_at_r)),
218        );
219
220        ///////////////////////////////////////////////////////////////////////
221        // Transition index
222        ///////////////////////////////////////////////////////////////////////
223
224        builder
225            .when(is_transition.clone())
226            .assert_eq(next.tidx, local.tidx - AB::Expr::from_usize(D_EF));
227
228        ///////////////////////////////////////////////////////////////////////
229        // Interactions
230        ///////////////////////////////////////////////////////////////////////
231
232        // Sample r
233        self.transcript_bus.sample_ext(
234            builder,
235            local.proof_idx,
236            local.tidx + AB::Expr::from_usize(D_EF),
237            local.r,
238            local.is_first,
239        );
240        // Observe coefficients
241        self.transcript_bus.observe_ext(
242            builder,
243            local.proof_idx,
244            local.tidx,
245            local.coeff,
246            local.is_valid,
247        );
248
249        // Receive initial tidx value
250        self.univariate_sumcheck_input_bus.receive(
251            builder,
252            local.proof_idx,
253            UnivariateSumcheckInputMessage { tidx: local.tidx },
254            is_last.clone(),
255        );
256        // Send tidx when there are no multilinear rounds
257        self.stacking_module_bus.send(
258            builder,
259            local.proof_idx,
260            StackingModuleMessage {
261                // Skip r
262                tidx: local.tidx + AB::Expr::from_usize(2 * D_EF),
263            },
264            local.is_first,
265        );
266
267        self.claim_bus.receive(
268            builder,
269            local.proof_idx,
270            SumcheckClaimMessage {
271                round: AB::Expr::NEG_ONE,
272                value: local.sum_at_roots.map(Into::into),
273            },
274            is_last.clone(),
275        );
276        self.claim_bus.send(
277            builder,
278            local.proof_idx,
279            SumcheckClaimMessage {
280                round: AB::Expr::ZERO,
281                value: local.value_at_r.map(Into::into),
282            },
283            is_last,
284        );
285        self.randomness_bus.send(
286            builder,
287            local.proof_idx,
288            ConstraintSumcheckRandomness {
289                idx: AB::Expr::ZERO,
290                challenge: local.r.map(Into::into),
291            },
292            local.is_first,
293        );
294
295        // Here idx = 0 and is called once per proof_idx
296        self.batch_constraint_conductor_bus.add_key_with_lookups(
297            builder,
298            local.proof_idx,
299            BatchConstraintConductorMessage {
300                msg_type: BatchConstraintInnerMessageType::R.to_field(),
301                idx: AB::Expr::ZERO,
302                value: local.r.map(Into::into),
303            },
304            local.is_valid * local.is_first * AB::Expr::TWO,
305        );
306    }
307}
openvm_recursion_circuit/batch_constraint/sumcheck/univariate/air.rs

openvm_recursion_circuit/batch_constraint/sumcheck/univariate/
air.rs
