openvm_recursion_circuit/batch_constraint/sumcheck/multilinear/
air.rs

1use std::borrow::Borrow;
2
3use openvm_circuit_primitives::{
4    utils::assert_array_eq, ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
5};
6use openvm_recursion_circuit_derive::AlignedBorrow;
7use openvm_stark_backend::{
8    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
9};
10use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF;
11use p3_air::{Air, AirBuilder, BaseAir};
12use p3_field::{extension::BinomiallyExtendable, Field, PrimeCharacteristicRing};
13use p3_matrix::Matrix;
14
15use crate::{
16    batch_constraint::bus::{
17        BatchConstraintConductorBus, BatchConstraintConductorMessage,
18        BatchConstraintInnerMessageType, SumcheckClaimBus, SumcheckClaimMessage,
19    },
20    bus::{
21        ConstraintSumcheckRandomness, ConstraintSumcheckRandomnessBus, StackingModuleBus,
22        StackingModuleMessage, TranscriptBus,
23    },
24    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
25    utils::{
26        assert_one_ext, ext_field_add, ext_field_multiply, ext_field_multiply_scalar,
27        ext_field_subtract_scalar, scalar_subtract_ext_field,
28    },
29};
30
31#[derive(AlignedBorrow, Clone, Copy, Debug, StructReflection)]
32#[repr(C)]
33pub struct MultilinearSumcheckCols<T> {
34    pub is_valid: T,
35    pub proof_idx: T,
36    pub round_idx: T,
37    pub is_proof_start: T,
38    pub is_first_eval: T,
39
40    /// A valid row which is not involved in any interactions
41    /// but should satisfy air constraints
42    pub is_dummy: T,
43
44    pub eval_idx: T,
45
46    pub cur_sum: [T; D_EF],
47    pub eval: [T; D_EF],
48
49    pub prefix_product: [T; D_EF],
50    pub suffix_product: [T; D_EF],
51    // 1 / i!(d - i)!
52    pub denom_inv: T,
53
54    // Lagrange coefficients for the interpolated polynomial
55    // prefix_product * suffix_product * denom_inv
56    pub lagrange_coeff: [T; D_EF],
57
58    pub r: [T; D_EF],
59
60    pub tidx: T,
61}
62
63#[derive(ColumnsAir)]
64#[columns_via(MultilinearSumcheckCols<u8>)]
65pub struct MultilinearSumcheckAir {
66    pub max_constraint_degree: usize,
67    pub claim_bus: SumcheckClaimBus,
68    pub transcript_bus: TranscriptBus,
69    pub randomness_bus: ConstraintSumcheckRandomnessBus,
70    pub batch_constraint_conductor_bus: BatchConstraintConductorBus,
71    pub stacking_module_bus: StackingModuleBus,
72}
73
74impl<F> BaseAirWithPublicValues<F> for MultilinearSumcheckAir {}
75impl<F> PartitionedBaseAir<F> for MultilinearSumcheckAir {}
76
77impl<F> BaseAir<F> for MultilinearSumcheckAir {
78    fn width(&self) -> usize {
79        MultilinearSumcheckCols::<F>::width()
80    }
81}
82
83impl<AB: AirBuilder + InteractionBuilder> Air<AB> for MultilinearSumcheckAir
84where
85    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
86{
87    fn eval(&self, builder: &mut AB) {
88        let main = builder.main();
89
90        let (local, next) = (
91            main.row_slice(0).expect("window should have two elements"),
92            main.row_slice(1).expect("window should have two elements"),
93        );
94        let local: &MultilinearSumcheckCols<AB::Var> = (*local).borrow();
95        let next: &MultilinearSumcheckCols<AB::Var> = (*next).borrow();
96
97        let s_deg = self.max_constraint_degree + 1;
98
99        ///////////////////////////////////////////////////////////////////////
100        // Loop Constraints
101        ///////////////////////////////////////////////////////////////////////
102
103        type LoopSubAir = NestedForLoopSubAir<2>;
104        LoopSubAir {}.eval(
105            builder,
106            (
107                NestedForLoopIoCols {
108                    is_enabled: local.is_valid,
109                    counter: [local.proof_idx, local.round_idx],
110                    is_first: [local.is_proof_start, local.is_first_eval],
111                }
112                .map_into(),
113                NestedForLoopIoCols {
114                    is_enabled: next.is_valid,
115                    counter: [next.proof_idx, next.round_idx],
116                    is_first: [next.is_proof_start, next.is_first_eval],
117                }
118                .map_into(),
119            ),
120        );
121
122        let is_transition_eval = LoopSubAir::local_is_transition(next.is_valid, next.is_first_eval);
123        let is_last_eval =
124            LoopSubAir::local_is_last(local.is_valid, next.is_valid, next.is_first_eval);
125
126        let is_same_proof = LoopSubAir::local_is_transition(next.is_valid, next.is_proof_start);
127        let is_proof_end =
128            LoopSubAir::local_is_last(local.is_valid, next.is_valid, next.is_proof_start);
129
130        // is_dummy forces a proof to be a single row
131        builder.assert_bool(local.is_dummy);
132        builder.when(local.is_dummy).assert_one(local.is_valid);
133        builder
134            .when(local.is_dummy)
135            .assert_one(local.is_proof_start);
136        builder
137            .when(local.is_dummy)
138            .assert_one(is_proof_end.clone());
139
140        let is_not_dummy = AB::Expr::ONE - local.is_dummy;
141
142        // Round idx starts at 0
143        builder
144            .when(local.is_proof_start)
145            .assert_zero(local.round_idx);
146
147        // Eval idx starts at 0
148        builder
149            .when(local.is_first_eval)
150            .assert_zero(local.eval_idx);
151        // Eval idx increments by 1
152        builder
153            .when(is_transition_eval.clone())
154            .assert_eq(next.eval_idx, local.eval_idx + AB::Expr::ONE);
155        // Eval idx ends at s_deg if not dummy
156        builder
157            .when(is_last_eval.clone() * is_not_dummy.clone())
158            .assert_eq(local.eval_idx, AB::Expr::from_usize(s_deg));
159
160        ///////////////////////////////////////////////////////////////////////
161        // Factorials Constraints
162        ///////////////////////////////////////////////////////////////////////
163
164        let d_factorial_inv = AB::Expr::from_prime_subfield(
165            <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield::from_usize((1..=s_deg).product())
166                .inverse(),
167        );
168        // Starts at d!
169        builder
170            .when(local.is_first_eval)
171            .assert_eq(local.denom_inv, d_factorial_inv);
172        // 1 / (i + 1)!(d - i - 1)!  (i + 1) = 1 / i!(d - i)! * (d - i)
173        builder.when(is_transition_eval.clone()).assert_eq(
174            next.denom_inv * next.eval_idx,
175            local.denom_inv * (AB::Expr::from_usize(s_deg) - local.eval_idx),
176        );
177
178        ///////////////////////////////////////////////////////////////////////
179        // Prefix/Suffix Product Constraints
180        ///////////////////////////////////////////////////////////////////////
181
182        assert_array_eq(
183            &mut builder.when(is_transition_eval.clone()),
184            next.r,
185            local.r,
186        );
187
188        // Prefix Product
189        // Starts at 1
190        assert_one_ext(&mut builder.when(local.is_first_eval), local.prefix_product);
191        // p' = p * (r - i)
192        assert_array_eq(
193            &mut builder.when(is_transition_eval.clone()),
194            next.prefix_product,
195            ext_field_multiply(
196                local.prefix_product,
197                ext_field_subtract_scalar(local.r, local.eval_idx),
198            ),
199        );
200
201        // Suffix Product
202        // Ends at 1
203        assert_one_ext(
204            &mut builder.when(is_last_eval.clone()),
205            local.suffix_product,
206        );
207        // s = s' * (i + 1 - r)
208        assert_array_eq(
209            &mut builder.when(is_transition_eval.clone()),
210            local.suffix_product,
211            ext_field_multiply(
212                next.suffix_product,
213                scalar_subtract_ext_field(local.eval_idx + AB::Expr::ONE, local.r),
214            ),
215        );
216
217        ///////////////////////////////////////////////////////////////////////
218        // Sumcheck evaluation constraints
219        ///////////////////////////////////////////////////////////////////////
220
221        // Lagrange coefficient
222        assert_array_eq(
223            &mut builder.when(local.is_valid),
224            local.lagrange_coeff,
225            ext_field_multiply_scalar(
226                ext_field_multiply(local.prefix_product, local.suffix_product),
227                local.denom_inv,
228            ),
229        );
230
231        // Initialize at first evaluation
232        assert_array_eq(
233            &mut builder.when(local.is_first_eval),
234            local.cur_sum,
235            ext_field_multiply(local.eval, local.lagrange_coeff),
236        );
237
238        // Cumulative sum
239        assert_array_eq(
240            &mut builder.when(is_transition_eval.clone()),
241            next.cur_sum,
242            ext_field_add(
243                local.cur_sum,
244                ext_field_multiply(next.eval, next.lagrange_coeff),
245            ),
246        );
247
248        ///////////////////////////////////////////////////////////////////////
249        // Transition constraints
250        ///////////////////////////////////////////////////////////////////////
251
252        builder
253            .when(is_same_proof)
254            .assert_eq(next.tidx, local.tidx + AB::Expr::from_usize(D_EF));
255
256        ///////////////////////////////////////////////////////////////////////
257        // Interactions
258        ///////////////////////////////////////////////////////////////////////
259
260        // Observe evaluations s(1), s(2) etc.
261        self.transcript_bus.observe_ext(
262            builder,
263            local.proof_idx,
264            local.tidx,
265            next.eval,
266            is_transition_eval.clone() * is_not_dummy.clone(),
267        );
268        // Sample challenge r
269        self.transcript_bus.sample_ext(
270            builder,
271            local.proof_idx,
272            local.tidx,
273            local.r,
274            is_last_eval.clone() * is_not_dummy.clone(),
275        );
276
277        // Receive tidx from univariate sumcheck and send it to stacking module
278        self.stacking_module_bus.receive(
279            builder,
280            local.proof_idx,
281            StackingModuleMessage { tidx: local.tidx },
282            local.is_proof_start * is_not_dummy.clone(),
283        );
284        self.stacking_module_bus.send(
285            builder,
286            local.proof_idx,
287            StackingModuleMessage {
288                tidx: local.tidx + AB::Expr::from_usize(D_EF),
289            },
290            is_proof_end.clone() * is_not_dummy.clone(),
291        );
292
293        self.claim_bus.receive(
294            builder,
295            local.proof_idx,
296            SumcheckClaimMessage {
297                round: local.round_idx.into(),
298                value: ext_field_add(local.eval, next.eval),
299            },
300            local.is_first_eval * is_not_dummy.clone(),
301        );
302        self.claim_bus.send(
303            builder,
304            local.proof_idx,
305            SumcheckClaimMessage {
306                round: local.round_idx + AB::Expr::ONE,
307                value: local.cur_sum.map(Into::into),
308            },
309            is_last_eval.clone() * is_not_dummy.clone(),
310        );
311        self.randomness_bus.send(
312            builder,
313            local.proof_idx,
314            ConstraintSumcheckRandomness {
315                idx: local.round_idx + AB::Expr::ONE,
316                challenge: local.r.map(|x| x.into()),
317            },
318            local.is_first_eval * is_not_dummy.clone(),
319        );
320        // Here idx > 0 and all idx are distinct within one proof_idx
321        self.batch_constraint_conductor_bus.add_key_with_lookups(
322            builder,
323            local.proof_idx,
324            BatchConstraintConductorMessage {
325                msg_type: BatchConstraintInnerMessageType::R.to_field(),
326                idx: local.round_idx + AB::Expr::ONE,
327                value: local.r.map(|x| x.into()),
328            },
329            local.is_first_eval * is_not_dummy,
330        );
331    }
332}
openvm_recursion_circuit/batch_constraint/sumcheck/multilinear/air.rs

openvm_recursion_circuit/batch_constraint/sumcheck/multilinear/
air.rs
