openvm_recursion_circuit/batch_constraint/eq_airs/eq_ns/
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, EqNOuterBus, EqNOuterMessage, EqZeroNBus, EqZeroNMessage,
20    },
21    bus::{
22        EqNsNLogupMaxBus, EqNsNLogupMaxMessage, SelHypercubeBus, SelHypercubeBusMessage,
23        XiRandomnessBus, XiRandomnessMessage,
24    },
25    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
26    utils::{
27        base_to_ext, ext_field_add, ext_field_multiply, ext_field_multiply_scalar,
28        ext_field_one_minus, ext_field_subtract,
29    },
30};
31
32#[derive(AlignedBorrow, Clone, Copy, StructReflection)]
33#[repr(C)]
34pub struct EqNsColumns<T> {
35    pub is_valid: T,
36    pub is_first: T,
37    pub proof_idx: T,
38
39    pub n: T,
40    pub n_logup: T,
41    pub n_max: T,
42    pub n_less_than_n_logup: T,
43    pub n_less_than_n_max: T,
44    pub is_transition_and_n_less_than_n_max: T,
45    pub xi_n: [T; D_EF],
46    pub r_n: [T; D_EF],
47    pub r_product: [T; D_EF],
48    pub r_pref_product: [T; D_EF],
49    pub one_minus_r_pref_prod: [T; D_EF],
50    pub eq: [T; D_EF],
51    pub eq_sharp: [T; D_EF],
52
53    /// The number of traces whose `n_lift` equals `local.n`.
54    /// Note that it cannot be derived from `xi_mult` because
55    /// `xi_mult` counts interactions, not AIRs with interactions.
56    pub num_traces: T,
57    pub xi_mult: T,
58    pub sel_first_count: T,
59    pub sel_last_and_trans_count: T,
60}
61
62#[derive(ColumnsAir)]
63#[columns_via(EqNsColumns<u8>)]
64pub struct EqNsAir {
65    pub zero_n_bus: EqZeroNBus,
66    pub xi_bus: XiRandomnessBus,
67    pub r_xi_bus: BatchConstraintConductorBus,
68    pub sel_hypercube_bus: SelHypercubeBus,
69    pub eq_n_outer_bus: EqNOuterBus,
70    pub eq_n_logup_n_max_bus: EqNsNLogupMaxBus,
71
72    pub l_skip: usize,
73}
74
75impl<F> BaseAirWithPublicValues<F> for EqNsAir {}
76impl<F> PartitionedBaseAir<F> for EqNsAir {}
77
78impl<F> BaseAir<F> for EqNsAir {
79    fn width(&self) -> usize {
80        EqNsColumns::<F>::width()
81    }
82}
83
84impl<AB: AirBuilder + InteractionBuilder> Air<AB> for EqNsAir
85where
86    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
87{
88    fn eval(&self, builder: &mut AB) {
89        let main = builder.main();
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
95        let local: &EqNsColumns<AB::Var> = (*local).borrow();
96        let next: &EqNsColumns<AB::Var> = (*next).borrow();
97
98        // Summary:
99        // - n consistency: treat `n_less_than_n_logup` as boolean, set `n = 0` on the first row,
100        //   increment `n` while continuing within the proof, propagate the “less than n_logup” flag
101        //   forward, and clear it whenever a row is invalid.
102        // - r consistency: initialize `r_product` to one when the next row begins the proof and,
103        //   otherwise, multiply by the current `r_n` values to update the running product.
104        // - eq consistency: update both `eq` and `eq_sharp` by multiplying with the shared factor
105        //   `1 - (xi + r - 2 * xi * r)` whenever advancing beyond the first row.
106
107        type LoopSubAir = NestedForLoopSubAir<1>;
108        LoopSubAir {}.eval(
109            builder,
110            (
111                NestedForLoopIoCols {
112                    is_enabled: local.is_valid,
113                    counter: [local.proof_idx],
114                    is_first: [local.is_first],
115                }
116                .map_into(),
117                NestedForLoopIoCols {
118                    is_enabled: next.is_valid,
119                    counter: [next.proof_idx],
120                    is_first: [next.is_first],
121                }
122                .map_into(),
123            ),
124        );
125
126        let is_transition = next.is_valid - next.is_first;
127        let is_last = local.is_valid - is_transition.clone();
128
129        // ========================= n consistency ==============================
130        builder.assert_bool(local.n_less_than_n_max);
131        builder.assert_bool(local.n_less_than_n_logup);
132        builder.when(local.is_first).assert_zero(local.n);
133        builder
134            .when(is_transition.clone())
135            .assert_one(next.n - local.n);
136        builder
137            .when(is_transition.clone())
138            .when(next.n_less_than_n_logup)
139            .assert_one(local.n_less_than_n_logup);
140        builder
141            .when(not(local.is_valid))
142            .assert_zero(local.n_less_than_n_logup);
143        builder
144            .when(is_transition.clone())
145            .when(next.n_less_than_n_max)
146            .assert_one(local.n_less_than_n_max);
147        builder
148            .when(not(local.is_valid))
149            .assert_zero(local.n_less_than_n_max);
150        builder.when(not(local.is_valid)).assert_zero(local.n_logup);
151        builder.when(not(local.is_valid)).assert_zero(local.n_max);
152        builder
153            .when(is_transition.clone())
154            .assert_eq(local.n_logup, next.n_logup);
155        builder
156            .when(is_transition.clone())
157            .assert_eq(local.n_max, next.n_max);
158
159        // When n_less_than_n_(logup|max) changes, it's because n was n_logup/n_max
160        builder
161            .when(local.n_less_than_n_logup)
162            .when(not(next.n_less_than_n_logup))
163            .assert_eq(next.n, next.n_logup);
164        builder
165            .when(local.n_less_than_n_max)
166            .when(not(next.n_less_than_n_max))
167            .assert_eq(next.n, next.n_max);
168
169        // These flags are zero on the last row.
170        // If n_logup/n_max is positive, the corresponding flag is one on the first row.
171        // Therefore, it changes at some point, which we constrained to be at the proper `n`.
172        // If n_logup/n_max is zero, then the corresponding flag is forced to be zero,
173        //   because if it was one at the first row, it wouldn't have a proper row to change.
174        builder
175            .when(local.n_logup)
176            .when(local.is_first)
177            .assert_one(local.n_less_than_n_logup);
178        builder
179            .when(local.n_max)
180            .when(local.is_first)
181            .assert_one(local.n_less_than_n_max);
182        builder
183            .when(is_last.clone())
184            .assert_zero(local.n_less_than_n_logup);
185        builder
186            .when(is_last.clone())
187            .assert_zero(local.n_less_than_n_max);
188        // ========================= r consistency ==============================
189        assert_array_eq(
190            &mut builder.when(local.is_valid - local.n_less_than_n_max),
191            local.r_n,
192            base_to_ext::<AB::Expr>(AB::Expr::ONE),
193        );
194        assert_array_eq(
195            &mut builder.when(is_last.clone()),
196            local.r_product,
197            base_to_ext::<AB::Expr>(AB::Expr::ONE),
198        );
199        assert_array_eq(
200            &mut builder.when(is_transition.clone()),
201            local.r_product,
202            ext_field_multiply(next.r_product, local.r_n),
203        );
204        assert_array_eq(
205            &mut builder.when(local.is_first),
206            local.r_pref_product,
207            base_to_ext::<AB::Expr>(AB::Expr::ONE),
208        );
209        assert_array_eq(
210            &mut builder.when(is_transition.clone() * local.n_less_than_n_max),
211            next.r_pref_product,
212            ext_field_multiply(local.r_pref_product, local.r_n),
213        );
214        assert_array_eq(
215            &mut builder.when(local.is_first),
216            local.one_minus_r_pref_prod,
217            base_to_ext::<AB::Expr>(AB::Expr::ONE),
218        );
219        assert_array_eq(
220            &mut builder.when(is_transition.clone() * local.n_less_than_n_max),
221            next.one_minus_r_pref_prod,
222            ext_field_multiply(
223                local.one_minus_r_pref_prod,
224                ext_field_subtract(base_to_ext::<AB::Expr>(AB::F::ONE), local.r_n),
225            ),
226        );
227        self.sel_hypercube_bus.add_key_with_lookups(
228            builder,
229            local.proof_idx,
230            SelHypercubeBusMessage {
231                n: local.n.into(),
232                is_first: AB::Expr::ZERO,
233                value: local.r_pref_product.map(Into::into),
234            },
235            local.is_valid * local.sel_last_and_trans_count,
236        );
237        self.sel_hypercube_bus.add_key_with_lookups(
238            builder,
239            local.proof_idx,
240            SelHypercubeBusMessage {
241                n: local.n.into(),
242                is_first: AB::Expr::ONE,
243                value: local.one_minus_r_pref_prod.map(Into::into),
244            },
245            local.is_valid * local.sel_first_count,
246        );
247        // ========================= eq consistency ===============================
248        let mult = ext_field_one_minus::<AB::Expr>(ext_field_subtract::<AB::Expr>(
249            ext_field_add(local.xi_n, local.r_n),
250            ext_field_multiply_scalar::<AB::Expr>(
251                ext_field_multiply(local.xi_n, local.r_n),
252                AB::Expr::TWO,
253            ),
254        ));
255        builder.assert_eq(
256            local.is_transition_and_n_less_than_n_max,
257            is_transition.clone() * local.n_less_than_n_max,
258        );
259        assert_array_eq(
260            &mut builder.when(local.is_transition_and_n_less_than_n_max),
261            next.eq,
262            ext_field_multiply(local.eq, mult.clone()),
263        );
264        assert_array_eq(
265            &mut builder.when(local.is_transition_and_n_less_than_n_max),
266            next.eq_sharp,
267            ext_field_multiply(local.eq_sharp, mult.clone()),
268        );
269
270        self.xi_bus.receive(
271            builder,
272            local.proof_idx,
273            XiRandomnessMessage {
274                idx: local.n + AB::Expr::from_usize(self.l_skip),
275                xi: local.xi_n.map(|x| x.into()),
276            },
277            is_transition.clone(),
278        );
279        // Here idx >= l_skip and all idx are different within one proof_idx
280        self.r_xi_bus.add_key_with_lookups(
281            builder,
282            local.proof_idx,
283            BatchConstraintConductorMessage {
284                msg_type: BatchConstraintInnerMessageType::Xi.to_field(),
285                idx: local.n + AB::Expr::from_usize(self.l_skip),
286                value: local.xi_n.map(|x| x.into()),
287            },
288            local.n_less_than_n_logup * local.xi_mult,
289        );
290
291        self.zero_n_bus.receive(
292            builder,
293            local.proof_idx,
294            EqZeroNMessage {
295                is_sharp: AB::Expr::ZERO,
296                value: local.eq.map(|x| x.into()),
297            },
298            local.is_first,
299        );
300        self.zero_n_bus.receive(
301            builder,
302            local.proof_idx,
303            EqZeroNMessage {
304                is_sharp: AB::Expr::ONE,
305                value: local.eq_sharp.map(|x| x.into()),
306            },
307            local.is_first,
308        );
309
310        self.r_xi_bus.lookup_key(
311            builder,
312            local.proof_idx,
313            BatchConstraintConductorMessage {
314                msg_type: BatchConstraintInnerMessageType::R.to_field(),
315                idx: local.n + AB::Expr::ONE,
316                value: local.r_n.map(|x| x.into()),
317            },
318            local.n_less_than_n_max,
319        );
320
321        self.eq_n_outer_bus.add_key_with_lookups(
322            builder,
323            next.proof_idx,
324            EqNOuterMessage {
325                is_sharp: AB::Expr::ZERO,
326                n: next.n.into(),
327                value: ext_field_multiply(next.eq, next.r_product),
328            },
329            next.is_valid * next.num_traces,
330        );
331        self.eq_n_outer_bus.add_key_with_lookups(
332            builder,
333            next.proof_idx,
334            EqNOuterMessage {
335                is_sharp: AB::Expr::ONE,
336                n: next.n.into(),
337                value: ext_field_multiply(next.eq_sharp, next.r_product),
338            },
339            next.is_valid * next.num_traces * AB::Expr::TWO, // two because num+denom per trace
340        );
341        self.eq_n_logup_n_max_bus.lookup_key(
342            builder,
343            local.proof_idx,
344            EqNsNLogupMaxMessage {
345                n_logup: local.n_logup,
346                n_max: local.n_max,
347            },
348            local.is_first,
349        );
350    }
351}
openvm_recursion_circuit/batch_constraint/eq_airs/eq_ns/air.rs

openvm_recursion_circuit/batch_constraint/eq_airs/eq_ns/
air.rs
