openvm_recursion_circuit/batch_constraint/fractions_folder/
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, PrimeCharacteristicRing};
13use p3_matrix::Matrix;
14
15use crate::{
16    batch_constraint::bus::{
17        BatchConstraintConductorBus, BatchConstraintConductorMessage,
18        BatchConstraintInnerMessageType, SumcheckClaimBus, SumcheckClaimMessage,
19        UnivariateSumcheckInputBus, UnivariateSumcheckInputMessage,
20    },
21    bus::{
22        BatchConstraintModuleBus, BatchConstraintModuleMessage, FractionFolderInputBus,
23        FractionFolderInputMessage, TranscriptBus,
24    },
25    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
26    utils::{ext_field_add, ext_field_multiply},
27};
28
29#[derive(AlignedBorrow, Clone, Copy, Debug, StructReflection)]
30#[repr(C)]
31pub struct FractionsFolderCols<T> {
32    pub is_valid: T,
33    pub proof_idx: T,
34    pub is_first: T,
35
36    pub air_idx: T,
37
38    pub tidx: T,
39
40    pub sum_claim_p: [T; D_EF],
41    pub sum_claim_q: [T; D_EF],
42    pub cur_p_sum: [T; D_EF],
43    pub cur_q_sum: [T; D_EF],
44    pub mu: [T; D_EF],
45    pub cur_hash: [T; D_EF],
46}
47
48#[derive(ColumnsAir)]
49#[columns_via(FractionsFolderCols<u8>)]
50pub struct FractionsFolderAir {
51    pub transcript_bus: TranscriptBus,
52    pub fraction_folder_input_bus: FractionFolderInputBus,
53    pub univariate_sumcheck_input_bus: UnivariateSumcheckInputBus,
54    pub sumcheck_bus: SumcheckClaimBus,
55    pub mu_bus: BatchConstraintConductorBus,
56    pub gkr_claim_bus: BatchConstraintModuleBus,
57}
58
59impl<F> BaseAirWithPublicValues<F> for FractionsFolderAir {}
60impl<F> PartitionedBaseAir<F> for FractionsFolderAir {}
61
62impl<F> BaseAir<F> for FractionsFolderAir {
63    fn width(&self) -> usize {
64        FractionsFolderCols::<F>::width()
65    }
66}
67
68impl<AB: AirBuilder + InteractionBuilder> Air<AB> for FractionsFolderAir
69where
70    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
71{
72    fn eval(&self, builder: &mut AB) {
73        let main = builder.main();
74        let (local, next) = (
75            main.row_slice(0).expect("window should have two elements"),
76            main.row_slice(1).expect("window should have two elements"),
77        );
78
79        let local: &FractionsFolderCols<AB::Var> = (*local).borrow();
80        let next: &FractionsFolderCols<AB::Var> = (*next).borrow();
81
82        ///////////////////////////////////////////////////////////////////////
83        // Loop Constraints
84        ///////////////////////////////////////////////////////////////////////
85
86        type LoopSubAir = NestedForLoopSubAir<1>;
87
88        LoopSubAir {}.eval(
89            builder,
90            (
91                NestedForLoopIoCols {
92                    is_enabled: local.is_valid,
93                    counter: [local.proof_idx],
94                    is_first: [local.is_first],
95                }
96                .map_into(),
97                NestedForLoopIoCols {
98                    is_enabled: next.is_valid,
99                    counter: [next.proof_idx],
100                    is_first: [next.is_first],
101                }
102                .map_into(),
103            ),
104        );
105
106        let is_transition = next.is_valid - next.is_first;
107        let is_last = local.is_valid - is_transition.clone();
108
109        // Air index decrements by 1
110        builder
111            .when(is_transition.clone())
112            .assert_eq(next.air_idx, local.air_idx - AB::Expr::ONE);
113        // Air index ends at 0
114        builder.when(is_last.clone()).assert_zero(local.air_idx);
115
116        ///////////////////////////////////////////////////////////////////////
117        // Transition Constraints
118        ///////////////////////////////////////////////////////////////////////
119
120        assert_array_eq(&mut builder.when(is_transition.clone()), local.mu, next.mu);
121
122        builder
123            .when(is_transition.clone())
124            .assert_eq(next.tidx, local.tidx - AB::Expr::from_usize(2 * D_EF));
125
126        ///////////////////////////////////////////////////////////////////////
127        // Running Sum Constraints
128        ///////////////////////////////////////////////////////////////////////
129
130        // Initialize running sums
131        assert_array_eq(
132            &mut builder.when(local.is_first),
133            local.cur_p_sum,
134            local.sum_claim_p,
135        );
136        assert_array_eq(
137            &mut builder.when(local.is_first),
138            local.cur_q_sum,
139            local.sum_claim_q,
140        );
141
142        // Add air sums to running sums
143        assert_array_eq(
144            &mut builder.when(is_transition.clone()),
145            next.cur_p_sum,
146            ext_field_add(local.cur_p_sum, next.sum_claim_p),
147        );
148        assert_array_eq(
149            &mut builder.when(is_transition.clone()),
150            next.cur_q_sum,
151            ext_field_add(local.cur_q_sum, next.sum_claim_q),
152        );
153
154        ///////////////////////////////////////////////////////////////////////
155        // Polynomial Hash Evaluation (a la Horner's Method)
156        ///////////////////////////////////////////////////////////////////////
157
158        // Initialize hash
159        // h = p + mu * q
160        assert_array_eq(
161            &mut builder.when(local.is_first),
162            local.cur_hash,
163            ext_field_add(
164                local.sum_claim_p,
165                ext_field_multiply(local.mu, local.sum_claim_q),
166            ),
167        );
168
169        // Update hash
170        // h' = p + mu * (q + mu * h)
171        assert_array_eq(
172            &mut builder.when(is_transition.clone()),
173            next.cur_hash,
174            ext_field_add(
175                next.sum_claim_p,
176                ext_field_multiply(
177                    next.mu,
178                    ext_field_add(
179                        next.sum_claim_q,
180                        ext_field_multiply(next.mu, local.cur_hash),
181                    ),
182                ),
183            ),
184        );
185
186        ///////////////////////////////////////////////////////////////////////
187        // Interactions
188        ///////////////////////////////////////////////////////////////////////
189
190        // Air index starts at num_present_airs - 1
191        self.fraction_folder_input_bus.receive(
192            builder,
193            local.proof_idx,
194            FractionFolderInputMessage {
195                num_present_airs: local.air_idx + AB::Expr::ONE,
196            },
197            local.is_first,
198        );
199
200        // Sample mu
201        self.transcript_bus.sample_ext(
202            builder,
203            local.proof_idx,
204            local.tidx + AB::Expr::from_usize(2 * D_EF),
205            local.mu,
206            local.is_first,
207        );
208        self.transcript_bus.observe_ext(
209            builder,
210            local.proof_idx,
211            local.tidx + AB::Expr::from_usize(D_EF),
212            local.sum_claim_q,
213            local.is_valid,
214        );
215        self.transcript_bus.observe_ext(
216            builder,
217            local.proof_idx,
218            local.tidx,
219            local.sum_claim_p,
220            local.is_valid,
221        );
222
223        self.sumcheck_bus.send(
224            builder,
225            local.proof_idx,
226            SumcheckClaimMessage {
227                round: AB::Expr::NEG_ONE,
228                value: local.cur_hash.map(Into::into),
229            },
230            is_last.clone(),
231        );
232
233        // Receive initial tidx and input layer claim from gkr module
234        self.gkr_claim_bus.receive(
235            builder,
236            local.proof_idx,
237            BatchConstraintModuleMessage {
238                // Skip lambda
239                tidx: local.tidx,
240                gkr_input_layer_claim: [local.cur_p_sum, local.cur_q_sum],
241            },
242            is_last,
243        );
244        // Send final tidx value to univariate sumcheck
245        self.univariate_sumcheck_input_bus.send(
246            builder,
247            local.proof_idx,
248            UnivariateSumcheckInputMessage {
249                // Skip mu
250                tidx: local.tidx + AB::Expr::from_usize(3 * D_EF),
251            },
252            local.is_first,
253        );
254
255        self.mu_bus.add_key_with_lookups(
256            builder,
257            local.proof_idx,
258            BatchConstraintConductorMessage {
259                msg_type: BatchConstraintInnerMessageType::Mu.to_field(),
260                idx: AB::Expr::ZERO,
261                value: local.mu.map(Into::into),
262            },
263            local.is_first,
264        );
265    }
266}
openvm_recursion_circuit/batch_constraint/fractions_folder/air.rs

openvm_recursion_circuit/batch_constraint/fractions_folder/
air.rs
