openvm_recursion_circuit/gkr/input/
air.rs

1use core::borrow::Borrow;
2
3use openvm_circuit_primitives::{
4    is_zero::{IsZeroAuxCols, IsZeroIo, IsZeroSubAir},
5    utils::{assert_array_eq, not, or},
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::{Field, PrimeCharacteristicRing};
15use p3_matrix::Matrix;
16
17use crate::{
18    bus::{
19        BatchConstraintModuleBus, BatchConstraintModuleMessage, ConstraintsFoldingInputBus,
20        ConstraintsFoldingInputMessage, GkrModuleBus, GkrModuleMessage,
21        InteractionsFoldingInputBus, InteractionsFoldingInputMessage, TranscriptBus,
22    },
23    gkr::bus::{
24        GkrLayerInputBus, GkrLayerInputMessage, GkrLayerOutputBus, GkrLayerOutputMessage,
25        GkrXiSamplerBus, GkrXiSamplerMessage,
26    },
27    primitives::bus::{ExpBitsLenBus, ExpBitsLenMessage},
28    subairs::proof_idx::{ProofIdxIoCols, ProofIdxSubAir},
29    utils::{assert_zeros, ext_field_subtract, pow_tidx_count},
30};
31
32#[repr(C)]
33#[derive(AlignedBorrow, Debug, StructReflection)]
34pub struct GkrInputCols<T> {
35    /// Whether the current row is enabled (i.e. not padding)
36    pub is_enabled: T,
37
38    pub proof_idx: T,
39
40    pub n_logup: T,
41    pub n_max: T,
42
43    /// Flag indicating whether there are any interactions
44    /// n_logup = 0 <=> total_interactions = 0
45    pub is_n_logup_zero: T,
46    pub is_n_logup_zero_aux: IsZeroAuxCols<T>,
47
48    pub is_n_max_greater_than_n_logup: T,
49
50    /// Transcript index
51    pub tidx: T,
52
53    /// Root denominator claim
54    pub q0_claim: [T; D_EF],
55
56    pub alpha_logup: [T; D_EF],
57
58    pub input_layer_claim: [[T; D_EF]; 2],
59
60    // Grinding
61    pub logup_pow_witness: T,
62    pub logup_pow_sample: T,
63}
64
65/// The GkrInputAir handles reading and passing the GkrInput
66#[derive(ColumnsAir)]
67#[columns_via(GkrInputCols<u8>)]
68pub struct GkrInputAir {
69    // System Params
70    pub l_skip: usize,
71    pub logup_pow_bits: usize,
72    // Buses
73    pub gkr_module_bus: GkrModuleBus,
74    pub bc_module_bus: BatchConstraintModuleBus,
75    pub transcript_bus: TranscriptBus,
76    pub exp_bits_len_bus: ExpBitsLenBus,
77    pub layer_input_bus: GkrLayerInputBus,
78    pub layer_output_bus: GkrLayerOutputBus,
79    pub xi_sampler_bus: GkrXiSamplerBus,
80    pub constraints_folding_input_bus: ConstraintsFoldingInputBus,
81    pub interactions_folding_input_bus: InteractionsFoldingInputBus,
82}
83
84impl<F: Field> BaseAir<F> for GkrInputAir {
85    fn width(&self) -> usize {
86        GkrInputCols::<F>::width()
87    }
88}
89
90impl<F: Field> BaseAirWithPublicValues<F> for GkrInputAir {}
91impl<F: Field> PartitionedBaseAir<F> for GkrInputAir {}
92
93impl<AB: AirBuilder + InteractionBuilder> Air<AB> for GkrInputAir {
94    fn eval(&self, builder: &mut AB) {
95        let main = builder.main();
96        let (local, next) = (
97            main.row_slice(0).expect("window should have two elements"),
98            main.row_slice(1).expect("window should have two elements"),
99        );
100        let local: &GkrInputCols<AB::Var> = (*local).borrow();
101        let next: &GkrInputCols<AB::Var> = (*next).borrow();
102
103        ///////////////////////////////////////////////////////////////////////
104        // Boolean Constraints
105        ///////////////////////////////////////////////////////////////////////
106
107        builder.assert_bool(local.is_n_max_greater_than_n_logup);
108
109        ///////////////////////////////////////////////////////////////////////
110        // Proof Index Constraints
111        ///////////////////////////////////////////////////////////////////////
112
113        // This subair has the following constraints:
114        // 1. Boolean enabled flag
115        // 2. Disabled rows are followed by disabled rows
116        // 3. Proof index increments by exactly one between enabled rows
117        ProofIdxSubAir.eval(
118            builder,
119            (
120                ProofIdxIoCols {
121                    is_enabled: local.is_enabled,
122                    proof_idx: local.proof_idx,
123                }
124                .map_into(),
125                ProofIdxIoCols {
126                    is_enabled: next.is_enabled,
127                    proof_idx: next.proof_idx,
128                }
129                .map_into(),
130            ),
131        );
132
133        ///////////////////////////////////////////////////////////////////////
134        // Base Constraints
135        ///////////////////////////////////////////////////////////////////////
136
137        // 1. Check if n_logup is zero (no logup constraints needed)
138        IsZeroSubAir.eval(
139            builder,
140            (
141                IsZeroIo::new(
142                    local.n_logup.into(),
143                    local.is_n_logup_zero.into(),
144                    local.is_enabled.into(),
145                ),
146                local.is_n_logup_zero_aux.inv,
147            ),
148        );
149
150        ///////////////////////////////////////////////////////////////////////
151        // Output Constraints
152        ///////////////////////////////////////////////////////////////////////
153
154        let has_interactions = AB::Expr::ONE - local.is_n_logup_zero;
155        // Input layer claim is [0, \alpha] when no interactions
156        assert_zeros(
157            &mut builder.when(not::<AB::Expr>(has_interactions.clone())),
158            local.input_layer_claim[0],
159        );
160        assert_array_eq(
161            &mut builder.when(not::<AB::Expr>(has_interactions.clone())),
162            local.input_layer_claim[1],
163            local.alpha_logup,
164        );
165        // q0_claim must be ONE when no interactions to prevent proof malleability
166        let ef_one = [AB::F::ONE, AB::F::ZERO, AB::F::ZERO, AB::F::ZERO];
167        assert_array_eq(
168            &mut builder.when(not::<AB::Expr>(has_interactions.clone())),
169            local.q0_claim,
170            ef_one,
171        );
172
173        ///////////////////////////////////////////////////////////////////////
174        // Module Interactions
175        ///////////////////////////////////////////////////////////////////////
176
177        let num_layers = local.n_logup + AB::Expr::from_usize(self.l_skip);
178
179        let needs_challenges = or(local.is_n_max_greater_than_n_logup, local.is_n_logup_zero);
180        let num_challenges = local.n_max + AB::Expr::from_usize(self.l_skip)
181            - has_interactions.clone() * num_layers.clone();
182
183        // Add PoW (if any) and alpha, beta
184        let logup_pow_offset = pow_tidx_count(self.logup_pow_bits);
185        let tidx_after_pow_and_alpha_beta =
186            local.tidx + AB::Expr::from_usize(logup_pow_offset + 2 * D_EF);
187        // Add GKR layers + Sumcheck
188        let tidx_after_gkr_layers = tidx_after_pow_and_alpha_beta.clone()
189            + has_interactions.clone()
190                * num_layers.clone()
191                * (num_layers.clone() + AB::Expr::TWO)
192                * AB::Expr::from_usize(2 * D_EF);
193        // Add separately sampled challenges
194        let tidx_end = tidx_after_gkr_layers.clone()
195            + needs_challenges.clone() * num_challenges.clone() * AB::Expr::from_usize(D_EF);
196
197        // 1. GkrLayerInputBus
198        // 1a. Send input to GkrLayerAir
199        self.layer_input_bus.send(
200            builder,
201            local.proof_idx,
202            GkrLayerInputMessage {
203                // Skip q0_claim
204                tidx: (tidx_after_pow_and_alpha_beta + AB::Expr::from_usize(D_EF))
205                    * has_interactions.clone(),
206                q0_claim: local.q0_claim.map(Into::into),
207            },
208            local.is_enabled * has_interactions.clone(),
209        );
210        // 2. GkrLayerOutputBus
211        // 2a. Receive input layer claim from GkrLayerAir
212        self.layer_output_bus.receive(
213            builder,
214            local.proof_idx,
215            GkrLayerOutputMessage {
216                tidx: tidx_after_gkr_layers.clone(),
217                layer_idx_end: num_layers.clone() - AB::Expr::ONE,
218                input_layer_claim: local.input_layer_claim.map(|claim| claim.map(Into::into)),
219            },
220            local.is_enabled * has_interactions.clone(),
221        );
222        // 3. GkrXiSamplerBus
223        // 3a. Send input to GkrXiSamplerAir
224        self.xi_sampler_bus.send(
225            builder,
226            local.proof_idx,
227            GkrXiSamplerMessage {
228                idx: has_interactions.clone() * num_layers,
229                tidx: tidx_after_gkr_layers,
230            },
231            local.is_enabled * needs_challenges.clone(),
232        );
233        // 3b. Receive output from GkrXiSamplerAir
234        self.xi_sampler_bus.receive(
235            builder,
236            local.proof_idx,
237            GkrXiSamplerMessage {
238                idx: local.n_max + AB::Expr::from_usize(self.l_skip - 1),
239                tidx: tidx_end.clone(),
240            },
241            local.is_enabled * needs_challenges,
242        );
243
244        ///////////////////////////////////////////////////////////////////////
245        // External Interactions
246        ///////////////////////////////////////////////////////////////////////
247
248        // 1. GkrModuleBus
249        // 1a. Receive initial GKR module message on first layer
250        self.gkr_module_bus.receive(
251            builder,
252            local.proof_idx,
253            GkrModuleMessage {
254                tidx: local.tidx,
255                n_logup: local.n_logup,
256                n_max: local.n_max,
257                is_n_max_greater: local.is_n_max_greater_than_n_logup,
258            },
259            local.is_enabled,
260        );
261
262        // 2. TranscriptBus
263        if self.logup_pow_bits > 0 {
264            // 2a. Observe pow witness
265            self.transcript_bus.observe(
266                builder,
267                local.proof_idx,
268                local.tidx.into(),
269                local.logup_pow_witness.into(),
270                local.is_enabled,
271            );
272            // 2b. Sample pow challenge
273            self.transcript_bus.sample(
274                builder,
275                local.proof_idx,
276                local.tidx.into() + AB::Expr::ONE,
277                local.logup_pow_sample.into(),
278                local.is_enabled,
279            );
280        }
281        // 2c. Sample alpha_logup challenge
282        self.transcript_bus.sample_ext(
283            builder,
284            local.proof_idx,
285            local.tidx.into() + AB::Expr::from_usize(logup_pow_offset),
286            local.alpha_logup.map(Into::into),
287            local.is_enabled,
288        );
289        // 2d. Observe `q0_claim` claim
290        self.transcript_bus.observe_ext(
291            builder,
292            local.proof_idx,
293            local.tidx + AB::Expr::from_usize(logup_pow_offset + 2 * D_EF),
294            local.q0_claim,
295            local.is_enabled * has_interactions,
296        );
297
298        // 3. BatchConstraintModuleBus
299        // 3a. Send input layer claims for further verification
300        self.bc_module_bus.send(
301            builder,
302            local.proof_idx,
303            BatchConstraintModuleMessage {
304                tidx: tidx_end.clone() + AB::Expr::from_usize(D_EF),
305                gkr_input_layer_claim: [
306                    local.input_layer_claim[0].map(Into::into),
307                    ext_field_subtract(local.input_layer_claim[1], local.alpha_logup),
308                ],
309            },
310            local.is_enabled,
311        );
312
313        self.constraints_folding_input_bus.send(
314            builder,
315            local.proof_idx,
316            ConstraintsFoldingInputMessage { tidx: tidx_end },
317            local.is_enabled,
318        );
319        self.interactions_folding_input_bus.send(
320            builder,
321            local.proof_idx,
322            InteractionsFoldingInputMessage {
323                tidx: local.tidx + AB::Expr::from_usize(logup_pow_offset + D_EF), // skip alpha
324            },
325            local.is_enabled,
326        );
327
328        // 4. ExpBitsLenBus
329        // 4a. Check proof-of-work using `ExpBitsLenBus`.
330        if self.logup_pow_bits > 0 {
331            self.exp_bits_len_bus.lookup_key(
332                builder,
333                ExpBitsLenMessage {
334                    base: AB::Expr::from_prime_subfield(
335                        <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield::GENERATOR,
336                    ),
337                    bit_src: local.logup_pow_sample.into(),
338                    num_bits: AB::Expr::from_usize(self.logup_pow_bits),
339                    result: AB::Expr::ONE,
340                },
341                local.is_enabled,
342            );
343        }
344    }
345}
openvm_recursion_circuit/gkr/input/air.rs

openvm_recursion_circuit/gkr/input/
air.rs
