openvm_recursion_circuit/batch_constraint/eq_airs/eq_3b/
air.rs

1use std::borrow::Borrow;
2
3use openvm_circuit_primitives::{
4    utils::{assert_array_eq, not},
5    ColumnsAir, StructReflection, StructReflectionHelper, SubAir,
6};
7use openvm_recursion_circuit_derive::AlignedBorrow;
8use openvm_stark_backend::{
9    interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
10};
11use openvm_stark_sdk::config::baby_bear_poseidon2::D_EF;
12use p3_air::{Air, AirBuilder, BaseAir};
13use p3_field::{extension::BinomiallyExtendable, PrimeCharacteristicRing};
14use p3_matrix::Matrix;
15
16use crate::{
17    batch_constraint::bus::{
18        BatchConstraintConductorBus, BatchConstraintConductorMessage,
19        BatchConstraintInnerMessageType, Eq3bBus, Eq3bMessage,
20    },
21    bus::{Eq3bShapeBus, Eq3bShapeMessage},
22    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
23    utils::{
24        base_to_ext, ext_field_add, ext_field_multiply, ext_field_multiply_scalar,
25        ext_field_one_minus,
26    },
27};
28
29#[derive(AlignedBorrow, Clone, Copy, StructReflection)]
30#[repr(C)]
31pub struct Eq3bColumns<T> {
32    pub is_valid: T,
33    pub is_first: T,
34    pub proof_idx: T,
35
36    pub sort_idx: T,
37    pub interaction_idx: T,
38
39    pub n_lift: T,
40    pub n_logup: T,
41    pub two_to_the_n_lift: T,
42    pub n: T,
43    pub hypercube_volume: T, // 2^n
44    pub n_at_least_n_lift: T,
45
46    pub has_no_interactions: T,
47
48    pub is_first_in_air: T,
49    pub is_first_in_interaction: T,
50
51    pub idx: T,         // stacked_idx >> l_skip, restored bit by bit
52    pub running_idx: T, // the current stacked_idx >> l_skip
53    pub nth_bit: T,
54
55    pub xi: [T; D_EF],
56    pub eq: [T; D_EF],
57}
58
59#[derive(ColumnsAir)]
60#[columns_via(Eq3bColumns<u8>)]
61pub struct Eq3bAir {
62    pub eq_3b_bus: Eq3bBus,
63    pub eq_3b_shape_bus: Eq3bShapeBus,
64    pub batch_constraint_conductor_bus: BatchConstraintConductorBus,
65
66    pub l_skip: usize,
67}
68
69impl<F> BaseAirWithPublicValues<F> for Eq3bAir {}
70impl<F> PartitionedBaseAir<F> for Eq3bAir {}
71
72impl<F> BaseAir<F> for Eq3bAir {
73    fn width(&self) -> usize {
74        Eq3bColumns::<F>::width()
75    }
76}
77
78impl<AB: AirBuilder + InteractionBuilder> Air<AB> for Eq3bAir
79where
80    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
81{
82    fn eval(&self, builder: &mut AB) {
83        let main = builder.main();
84        let (local, next) = (
85            main.row_slice(0).expect("window should have two elements"),
86            main.row_slice(1).expect("window should have two elements"),
87        );
88
89        let local: &Eq3bColumns<AB::Var> = (*local).borrow();
90        let next: &Eq3bColumns<AB::Var> = (*next).borrow();
91
92        type LoopSubAir = NestedForLoopSubAir<3>;
93        LoopSubAir {}.eval(
94            builder,
95            (
96                NestedForLoopIoCols {
97                    is_enabled: local.is_valid,
98                    counter: [local.proof_idx, local.sort_idx, local.interaction_idx],
99                    is_first: [
100                        local.is_first,
101                        local.is_first_in_air,
102                        local.is_first_in_interaction,
103                    ],
104                }
105                .map_into(),
106                NestedForLoopIoCols {
107                    is_enabled: next.is_valid,
108                    counter: [next.proof_idx, next.sort_idx, next.interaction_idx],
109                    is_first: [
110                        next.is_first,
111                        next.is_first_in_air,
112                        next.is_first_in_interaction,
113                    ],
114                }
115                .map_into(),
116            ),
117        );
118
119        // ======================= incides boundary conditions =========================
120        builder.when(local.is_first).assert_zero(local.sort_idx);
121        builder
122            .when(local.is_first_in_air)
123            .assert_zero(local.interaction_idx);
124
125        builder.assert_bool(local.n_at_least_n_lift);
126        builder.assert_bool(local.nth_bit);
127        builder.assert_bool(local.has_no_interactions);
128
129        let within_one_air = next.is_valid - next.is_first_in_air;
130        let within_one_interaction = next.is_valid - next.is_first_in_interaction;
131        let is_last_in_interaction = local.is_valid - within_one_interaction.clone();
132
133        // =============================== n consistency ==================================
134        builder
135            .when(local.is_first_in_interaction)
136            .assert_zero(local.n);
137        builder
138            .when(local.is_first_in_interaction)
139            .when(local.is_valid)
140            .assert_one(local.hypercube_volume);
141        builder
142            .when(within_one_air.clone())
143            .assert_eq(next.n_lift, local.n_lift);
144        builder
145            .when(LoopSubAir::local_is_transition(
146                next.is_valid,
147                next.is_first,
148            ))
149            .assert_eq(next.n_logup, local.n_logup);
150        builder
151            .when(within_one_interaction.clone())
152            .assert_eq(next.two_to_the_n_lift, local.two_to_the_n_lift);
153        builder
154            .when(within_one_interaction.clone())
155            .assert_eq(next.n, local.n + AB::Expr::ONE);
156        builder.when(within_one_interaction.clone()).assert_eq(
157            next.hypercube_volume,
158            local.hypercube_volume * AB::Expr::TWO,
159        );
160        // n_at_least_n_lift is nondecreasing
161        builder
162            .when(within_one_interaction.clone())
163            .when(local.n_at_least_n_lift)
164            .assert_one(next.n_at_least_n_lift);
165        // it's always 1 in the end
166        builder
167            .when(not(local.has_no_interactions))
168            .when(is_last_in_interaction.clone())
169            .when(local.is_valid)
170            .assert_one(local.n_at_least_n_lift);
171
172        // Either there is a moment where it switches from 0 to 1, then it's when n = n_lift
173        builder
174            .when(not(local.n_at_least_n_lift))
175            .when(next.n_at_least_n_lift)
176            .assert_eq(next.n, next.n_lift);
177        builder
178            .when(not(local.n_at_least_n_lift))
179            .when(next.n_at_least_n_lift)
180            .assert_eq(next.hypercube_volume, next.two_to_the_n_lift);
181        // Or it's 1 from the beginning, in which case n_lift = 0
182        builder
183            .when(local.is_first_in_interaction)
184            .when(local.n_at_least_n_lift)
185            .assert_zero(local.n_lift);
186        builder
187            .when(local.is_first_in_interaction)
188            .when(local.n_at_least_n_lift)
189            .assert_one(local.two_to_the_n_lift);
190        builder
191            .when(is_last_in_interaction.clone())
192            .when(not(local.has_no_interactions))
193            .assert_eq(local.n, local.n_logup);
194
195        builder.when(next.is_valid - next.is_first).assert_eq(
196            next.running_idx,
197            local.running_idx
198                + (next.is_first_in_interaction - local.has_no_interactions)
199                    * local.two_to_the_n_lift,
200        );
201
202        // =========================== Xi and product consistency =============================
203        // Boundary conditions
204        assert_array_eq(
205            &mut builder.when(local.is_valid * local.is_first_in_interaction),
206            local.eq,
207            base_to_ext::<AB::Expr>(AB::Expr::ONE),
208        );
209        builder
210            .when(local.is_first_in_interaction)
211            .assert_zero(local.idx);
212        builder.when(local.is_first).assert_zero(local.running_idx);
213        builder
214            .when(is_last_in_interaction.clone())
215            .when(not(local.has_no_interactions))
216            .assert_eq(local.idx, local.running_idx);
217        builder
218            .when(next.is_valid)
219            .when(local.has_no_interactions)
220            .assert_eq(local.running_idx, next.running_idx);
221
222        // If n is less than n_lift, assert that eq doesn't change
223        assert_array_eq(
224            &mut builder
225                .when(local.is_valid - local.has_no_interactions)
226                .when(not(local.n_at_least_n_lift)),
227            local.eq,
228            next.eq,
229        );
230        // Within transition, idx increases by nth_bit * hypercube_volume
231        builder
232            .when(within_one_interaction.clone())
233            .assert_eq(next.idx, local.idx + local.nth_bit * local.hypercube_volume);
234        // It can't increase if n < n_lift
235        builder
236            .when(not(local.n_at_least_n_lift))
237            .assert_zero(local.nth_bit);
238        // When transition, eq multiplies correspondingly
239        assert_array_eq(
240            &mut builder.when(within_one_interaction.clone()),
241            next.eq,
242            ext_field_multiply(
243                local.eq,
244                ext_field_add::<AB::Expr>(
245                    ext_field_multiply_scalar(local.xi, local.nth_bit),
246                    ext_field_multiply_scalar::<AB::Expr>(
247                        ext_field_one_minus(local.xi),
248                        AB::Expr::ONE - local.nth_bit,
249                    ),
250                ),
251            ),
252        );
253
254        // ==================== no interactions consistency ==========================
255        builder
256            .when(local.has_no_interactions)
257            .assert_one(local.is_valid);
258        builder
259            .when(local.has_no_interactions)
260            .assert_one(local.is_first_in_air);
261        builder
262            .when(local.has_no_interactions)
263            .assert_zero(within_one_air);
264
265        self.batch_constraint_conductor_bus.lookup_key(
266            builder,
267            local.proof_idx,
268            BatchConstraintConductorMessage {
269                msg_type: BatchConstraintInnerMessageType::Xi.to_field(),
270                idx: local.n + AB::Expr::from_usize(self.l_skip),
271                value: local.xi.map(|x| x.into()),
272            },
273            local.n_at_least_n_lift * within_one_interaction,
274        );
275
276        self.eq_3b_shape_bus.lookup_key(
277            builder,
278            local.proof_idx,
279            Eq3bShapeMessage {
280                sort_idx: local.sort_idx.into(),
281                n_lift: local.n_lift.into(),
282                n_logup: local.n_logup.into(),
283                num_interactions: local.interaction_idx + not(local.has_no_interactions),
284            },
285            LoopSubAir::local_is_last(local.is_valid, next.is_valid, next.is_first_in_air),
286        );
287
288        self.eq_3b_bus.send(
289            builder,
290            local.proof_idx,
291            Eq3bMessage {
292                sort_idx: local.sort_idx,
293                interaction_idx: local.interaction_idx,
294                eq_3b: local.eq,
295            },
296            is_last_in_interaction * (local.is_valid - local.has_no_interactions),
297        );
298    }
299}
openvm_recursion_circuit/batch_constraint/eq_airs/eq_3b/air.rs

openvm_recursion_circuit/batch_constraint/eq_airs/eq_3b/
air.rs
