openvm_recursion_circuit/gkr/layer/
air.rs

1use core::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    bus::{TranscriptBus, XiRandomnessBus, XiRandomnessMessage},
17    gkr::{
18        bus::{
19            GkrLayerInputBus, GkrLayerInputMessage, GkrLayerOutputBus, GkrLayerOutputMessage,
20            GkrSumcheckInputBus, GkrSumcheckInputMessage, GkrSumcheckOutputBus,
21            GkrSumcheckOutputMessage,
22        },
23        GkrSumcheckChallengeBus, GkrSumcheckChallengeMessage,
24    },
25    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
26    utils::{assert_zeros, ext_field_add, ext_field_multiply, ext_field_subtract},
27};
28
29#[repr(C)]
30#[derive(AlignedBorrow, Debug, StructReflection)]
31pub struct GkrLayerCols<T> {
32    /// Whether the current row is enabled (i.e. not padding)
33    pub is_enabled: T,
34    pub proof_idx: T,
35    pub is_first: T,
36
37    /// An enabled row which is not involved in any interactions
38    /// but should satisfy air constraints
39    pub is_dummy: T,
40
41    /// GKR layer index
42    pub layer_idx: T,
43
44    /// Transcript index at the start of this layer
45    pub tidx: T,
46
47    /// Sampled batching challenge
48    pub lambda: [T; D_EF],
49
50    /// Layer claims
51    pub p_xi_0: [T; D_EF],
52    pub q_xi_0: [T; D_EF],
53    pub p_xi_1: [T; D_EF],
54    pub q_xi_1: [T; D_EF],
55
56    // (p_xi_1 - p_xi_0) * mu + p_xi_0
57    pub numer_claim: [T; D_EF],
58    // (q_xi_1 - q_xi_0) * mu + q_xi_0
59    pub denom_claim: [T; D_EF],
60
61    // Sumcheck claim input
62    pub sumcheck_claim_in: [T; D_EF],
63
64    /// Received from GkrLayerSumcheckAir
65    pub eq_at_r_prime: [T; D_EF],
66
67    /// Corresponds to `mu` - reduction point
68    pub mu: [T; D_EF],
69}
70
71/// The GkrLayerAir handles layer-to-layer transitions in the GKR protocol
72#[derive(ColumnsAir)]
73#[columns_via(GkrLayerCols<u8>)]
74pub struct GkrLayerAir {
75    // External buses
76    pub xi_randomness_bus: XiRandomnessBus,
77    pub transcript_bus: TranscriptBus,
78    // Internal buses
79    pub layer_input_bus: GkrLayerInputBus,
80    pub layer_output_bus: GkrLayerOutputBus,
81    pub sumcheck_input_bus: GkrSumcheckInputBus,
82    pub sumcheck_output_bus: GkrSumcheckOutputBus,
83    pub sumcheck_challenge_bus: GkrSumcheckChallengeBus,
84}
85
86impl<F: Field> BaseAir<F> for GkrLayerAir {
87    fn width(&self) -> usize {
88        GkrLayerCols::<F>::width()
89    }
90}
91
92impl<F: Field> BaseAirWithPublicValues<F> for GkrLayerAir {}
93impl<F: Field> PartitionedBaseAir<F> for GkrLayerAir {}
94
95impl<AB: AirBuilder + InteractionBuilder> Air<AB> for GkrLayerAir
96where
97    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
98{
99    fn eval(&self, builder: &mut AB) {
100        let main = builder.main();
101        let (local, next) = (
102            main.row_slice(0).expect("window should have two elements"),
103            main.row_slice(1).expect("window should have two elements"),
104        );
105        let local: &GkrLayerCols<AB::Var> = (*local).borrow();
106        let next: &GkrLayerCols<AB::Var> = (*next).borrow();
107
108        ///////////////////////////////////////////////////////////////////////
109        // Boolean Constraints
110        ///////////////////////////////////////////////////////////////////////
111
112        builder.assert_bool(local.is_dummy);
113
114        ///////////////////////////////////////////////////////////////////////
115        // Proof Index and Loop Constraints
116        ///////////////////////////////////////////////////////////////////////
117
118        type LoopSubAir = NestedForLoopSubAir<1>;
119
120        // This subair has the following constraints:
121        // 1. Boolean enabled flag
122        // 2. Disabled rows are followed by disabled rows
123        // 3. Proof index increments by exactly one between enabled rows
124        LoopSubAir {}.eval(
125            builder,
126            (
127                NestedForLoopIoCols {
128                    is_enabled: local.is_enabled,
129                    counter: [local.proof_idx],
130                    is_first: [local.is_first],
131                }
132                .map_into(),
133                NestedForLoopIoCols {
134                    is_enabled: next.is_enabled,
135                    counter: [next.proof_idx],
136                    is_first: [next.is_first],
137                }
138                .map_into(),
139            ),
140        );
141
142        let is_transition = LoopSubAir::local_is_transition(next.is_enabled, next.is_first);
143        let is_last = LoopSubAir::local_is_last(local.is_enabled, next.is_enabled, next.is_first);
144
145        // Layer index starts from 0
146        builder.when(local.is_first).assert_zero(local.layer_idx);
147        // Layer index increments by 1
148        builder
149            .when(is_transition.clone())
150            .assert_eq(next.layer_idx, local.layer_idx + AB::Expr::ONE);
151
152        ///////////////////////////////////////////////////////////////////////
153        // Dummy Row Constraints
154        ///////////////////////////////////////////////////////////////////////
155
156        // A proof can't contribute both dummy and non-dummy rows
157        builder
158            .when(is_transition.clone())
159            .assert_eq(next.is_dummy, local.is_dummy);
160        // Any proof segment with more than one layer row must be non-dummy
161        builder
162            .when(is_transition.clone())
163            .assert_zero(local.is_dummy);
164
165        // Dummy rows are only allowed as a singleton placeholder row
166        builder.when(local.is_dummy).assert_one(local.is_first);
167        builder.when(local.is_dummy).assert_one(is_last.clone());
168
169        ///////////////////////////////////////////////////////////////////////
170        // Root Layer Constraints
171        ///////////////////////////////////////////////////////////////////////
172
173        // Compute cross terms: p_cross = p_xi_0 * q_xi_1 + p_xi_1 * q_xi_0
174        //                       q_cross = q_xi_0 * q_xi_1
175        let (p_cross_term, q_cross_term) =
176            compute_recursive_relations(local.p_xi_0, local.q_xi_0, local.p_xi_1, local.q_xi_1);
177
178        // Zero-check: verify p_cross = 0 at root layer
179        assert_zeros(&mut builder.when(local.is_first), p_cross_term.clone());
180
181        // Root consistency check: verify q_cross = q0_claim
182        assert_array_eq(
183            &mut builder.when(local.is_first),
184            q_cross_term.clone(),
185            local.sumcheck_claim_in,
186        );
187
188        ///////////////////////////////////////////////////////////////////////
189        // Layer Constraints
190        ///////////////////////////////////////////////////////////////////////
191
192        // Reduce to single evaluation
193        // `numer_claim = (p_xi_1 - p_xi_0) * mu + p_xi_0`
194        // `denom_claim = (q_xi_1 - q_xi_0) * mu + q_xi_0`
195        let (numer_claim, denom_claim) = reduce_to_single_evaluation(
196            local.p_xi_0,
197            local.p_xi_1,
198            local.q_xi_0,
199            local.q_xi_1,
200            local.mu,
201        );
202        assert_array_eq(builder, local.numer_claim, numer_claim);
203        assert_array_eq(builder, local.denom_claim, denom_claim);
204
205        ///////////////////////////////////////////////////////////////////////
206        // Inter-Layer Constraints
207        ///////////////////////////////////////////////////////////////////////
208
209        // Next layer claim is RLC of previous layer numer_claim and denom_claim
210        assert_array_eq(
211            &mut builder.when(is_transition.clone()),
212            next.sumcheck_claim_in,
213            ext_field_add::<AB::Expr>(
214                local.numer_claim,
215                ext_field_multiply::<AB::Expr>(next.lambda, local.denom_claim),
216            ),
217        );
218
219        // Transcript index increment
220        let tidx_after_sumcheck = local.tidx
221            // Sample lambda on non-root layer
222            + (AB::Expr::ONE - local.is_first) * AB::Expr::from_usize(D_EF)
223            + local.layer_idx * AB::Expr::from_usize(4 * D_EF);
224        let tidx_end = tidx_after_sumcheck.clone() + AB::Expr::from_usize(5 * D_EF);
225        builder
226            .when(is_transition.clone())
227            .assert_eq(next.tidx, tidx_end.clone());
228
229        ///////////////////////////////////////////////////////////////////////
230        // Module Interactions
231        ///////////////////////////////////////////////////////////////////////
232
233        let is_not_dummy = AB::Expr::ONE - local.is_dummy;
234        let is_non_root_layer = local.is_enabled * (AB::Expr::ONE - local.is_first);
235
236        // 1. GkrLayerInputBus
237        // 1a. Receive GKR layers input
238        self.layer_input_bus.receive(
239            builder,
240            local.proof_idx,
241            GkrLayerInputMessage {
242                tidx: local.tidx,
243                q0_claim: local.sumcheck_claim_in,
244            },
245            local.is_first * is_not_dummy.clone(),
246        );
247        // 2. GkrLayerOutputBus
248        // 2a. Send GKR input layer claims back
249        self.layer_output_bus.send(
250            builder,
251            local.proof_idx,
252            GkrLayerOutputMessage {
253                tidx: tidx_end,
254                layer_idx_end: local.layer_idx.into(),
255                input_layer_claim: [
256                    local.numer_claim.map(Into::into),
257                    local.denom_claim.map(Into::into),
258                ],
259            },
260            is_last.clone() * is_not_dummy.clone(),
261        );
262        // 3. GkrSumcheckInputBus
263        // 3a. Send claim to sumcheck
264        self.sumcheck_input_bus.send(
265            builder,
266            local.proof_idx,
267            GkrSumcheckInputMessage {
268                layer_idx: local.layer_idx.into(),
269                is_last_layer: is_last.clone(),
270                tidx: local.tidx + AB::Expr::from_usize(D_EF),
271                claim: local.sumcheck_claim_in.map(Into::into),
272            },
273            is_non_root_layer.clone() * is_not_dummy.clone(),
274        );
275        // 3. GkrSumcheckOutputBus
276        // 3a. Receive sumcheck results
277        let sumcheck_claim_out = ext_field_multiply::<AB::Expr>(
278            ext_field_add::<AB::Expr>(
279                p_cross_term,
280                ext_field_multiply::<AB::Expr>(local.lambda, q_cross_term),
281            ),
282            local.eq_at_r_prime,
283        );
284        self.sumcheck_output_bus.receive(
285            builder,
286            local.proof_idx,
287            GkrSumcheckOutputMessage {
288                layer_idx: local.layer_idx.into(),
289                tidx: tidx_after_sumcheck.clone(),
290                claim_out: sumcheck_claim_out.map(Into::into),
291                eq_at_r_prime: local.eq_at_r_prime.map(Into::into),
292            },
293            is_non_root_layer.clone() * is_not_dummy.clone(),
294        );
295        // 4. GkrSumcheckChallengeBus
296        // 4a. Send challenge mu
297        self.sumcheck_challenge_bus.send(
298            builder,
299            local.proof_idx,
300            GkrSumcheckChallengeMessage {
301                layer_idx: local.layer_idx.into(),
302                sumcheck_round: AB::Expr::ZERO,
303                challenge: local.mu.map(Into::into),
304            },
305            is_transition.clone() * is_not_dummy.clone(),
306        );
307
308        ///////////////////////////////////////////////////////////////////////
309        // External Interactions
310        ///////////////////////////////////////////////////////////////////////
311
312        // 1. TranscriptBus
313        // 1a. Sample `lambda`
314        self.transcript_bus.sample_ext(
315            builder,
316            local.proof_idx,
317            local.tidx,
318            local.lambda,
319            is_non_root_layer.clone() * is_not_dummy.clone(),
320        );
321        // 1b. Observe layer claims
322        let mut tidx = tidx_after_sumcheck;
323        for claim in [local.p_xi_0, local.q_xi_0, local.p_xi_1, local.q_xi_1].into_iter() {
324            self.transcript_bus.observe_ext(
325                builder,
326                local.proof_idx,
327                tidx.clone(),
328                claim,
329                local.is_enabled * is_not_dummy.clone(),
330            );
331            tidx += AB::Expr::from_usize(D_EF);
332        }
333        // 1c. Sample `mu`
334        self.transcript_bus.sample_ext(
335            builder,
336            local.proof_idx,
337            tidx,
338            local.mu,
339            local.is_enabled * is_not_dummy.clone(),
340        );
341
342        // 2. XiRandomnessBus
343        // 2a. Send shared randomness
344        self.xi_randomness_bus.send(
345            builder,
346            local.proof_idx,
347            XiRandomnessMessage {
348                idx: AB::Expr::ZERO,
349                xi: local.mu.map(Into::into),
350            },
351            is_last * is_not_dummy.clone(),
352        );
353    }
354}
355
356/// Computes recursive relations from layer claims.
357///
358/// Returns `(p_cross_term, q_cross_term)` where:
359/// - `p_cross_term = p_xi_0 * q_xi_1 + p_xi_1 * q_xi_0`
360/// - `q_cross_term = q_xi_0 * q_xi_1`
361fn compute_recursive_relations<F, FA>(
362    p_xi_0: [F; D_EF],
363    q_xi_0: [F; D_EF],
364    p_xi_1: [F; D_EF],
365    q_xi_1: [F; D_EF],
366) -> ([FA; D_EF], [FA; D_EF])
367where
368    F: Into<FA> + Copy,
369    FA: PrimeCharacteristicRing,
370    FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
371{
372    let p_cross_term = ext_field_add::<FA>(
373        ext_field_multiply::<FA>(p_xi_0, q_xi_1),
374        ext_field_multiply::<FA>(p_xi_1, q_xi_0),
375    );
376    let q_cross_term = ext_field_multiply::<FA>(q_xi_0, q_xi_1);
377    (p_cross_term, q_cross_term)
378}
379
380/// Linearly interpolates between two points at 0 and 1.
381fn interpolate_linear_at_01<F, FA>(evals: [[F; D_EF]; 2], x: [F; D_EF]) -> [FA; D_EF]
382where
383    F: Into<FA> + Copy,
384    FA: PrimeCharacteristicRing,
385    FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
386{
387    let p: [FA; D_EF] = ext_field_subtract(evals[1], evals[0]);
388    ext_field_add(ext_field_multiply::<FA>(p, x), evals[0])
389}
390
391/// Reduces claims to a single evaluation point using linear interpolation.
392///
393/// Returns `(numer, denom)` where:
394/// - `numer = (p_xi_1 - p_xi_0) * mu + p_xi_0`
395/// - `denom = (q_xi_1 - q_xi_0) * mu + q_xi_0`
396pub(super) fn reduce_to_single_evaluation<F, FA>(
397    p_xi_0: [F; D_EF],
398    p_xi_1: [F; D_EF],
399    q_xi_0: [F; D_EF],
400    q_xi_1: [F; D_EF],
401    mu: [F; D_EF],
402) -> ([FA; D_EF], [FA; D_EF])
403where
404    F: Into<FA> + Copy,
405    FA: PrimeCharacteristicRing,
406    FA::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
407{
408    let numer = interpolate_linear_at_01([p_xi_0, p_xi_1], mu);
409    let denom = interpolate_linear_at_01([q_xi_0, q_xi_1], mu);
410    (numer, denom)
411}
openvm_recursion_circuit/gkr/layer/air.rs

openvm_recursion_circuit/gkr/layer/
air.rs
