openvm_recursion_circuit/batch_constraint/expression_claim/
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, ExpressionClaimBus,
20        ExpressionClaimMessage, SumcheckClaimBus, SumcheckClaimMessage,
21    },
22    bus::{ExpressionClaimNMaxBus, ExpressionClaimNMaxMessage, HyperdimBus, HyperdimBusMessage},
23    primitives::bus::{PowerCheckerBus, PowerCheckerBusMessage},
24    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
25    utils::{base_to_ext, ext_field_add, ext_field_multiply, ext_field_multiply_scalar},
26};
27
28/// For each proof, this AIR will receive 2t interaction claims and t constraint claims.
29/// (2 interaction claims and 1 constraint claim per trace).
30/// These values are folded (algebraic batching) with mu into a single value, which
31/// should match the final sumcheck claim.
32///
33/// Rows are structured as a nested loop: for each proof, group 0 (interactions) comes first,
34/// then group 1 (constraints). Within the interaction group, each trace occupies 2 rows
35/// (numerator then denominator). Within the constraint group, each trace occupies 1 row.
36/// `NestedForLoopSubAir<4>` enforces canonical nested enumeration over
37/// `(proof_idx, group_idx, trace_idx, idx_parity)`.
38///
39/// Example for t = 2 traces (one proof):
40///
41/// ```text
42/// row | is_first | group_idx | is_first_in_group | trace_idx | idx_parity | idx
43/// ----|----------|-----------|-------------------|-----------|------------|----
44///   0 |    1     |     0     |         1         |     0     |     0      |  0   ← numerator trace 0
45///   1 |    0     |     0     |         0         |     0     |     1      |  1   ← denominator trace 0
46///   2 |    0     |     0     |         0         |     1     |     0      |  2   ← numerator trace 1
47///   3 |    0     |     0     |         0         |     1     |     1      |  3   ← denominator trace 1
48///   4 |    0     |     1     |         1         |     0     |     0      |  0   ← constraint trace 0
49///   5 |    0     |     1     |         0         |     1     |     0      |  1   ← constraint trace 1
50/// ```
51///
52/// `is_interaction` is derived as `1 - group_idx` (not a separate column).
53#[derive(AlignedBorrow, Copy, Clone, Debug, StructReflection)]
54#[repr(C)]
55pub struct ExpressionClaimCols<T> {
56    // --- Loop structure (enforced by NestedForLoopSubAir<4>) ---
57    pub is_valid: T,
58    /// First row of a proof. Marks proof boundaries.
59    pub is_first: T,
60    pub proof_idx: T,
61    /// 0 = interaction group, 1 = constraint group. Monotone within a proof.
62    pub group_idx: T,
63    /// Marks the first row of each group (set at proof start and interaction→constraint boundary).
64    pub is_first_in_group: T,
65
66    // --- Claim indexing (derived from loop counters) ---
67    /// Claim index within its group. For interactions: `2 * trace_idx + idx_parity` (0..2t).
68    /// For constraints: `trace_idx` (0..t).
69    pub idx: T,
70    /// 0 = numerator, 1 = denominator. Always 0 on constraint rows. Alternates on interaction
71    /// rows.
72    pub idx_parity: T,
73    /// Sorted trace index within the group. Monotone non-decreasing; resets at group boundaries.
74    pub trace_idx: T,
75    /// The received evaluation claim. Note that for interactions, this is without norm_factor and
76    /// eq_sharp_ns. These are interactions_evals (without norm_factor and eq_sharp_ns) and
77    /// constraint_evals in the rust verifier.
78    pub value: [T; D_EF],
79    /// Receive from eq_ns AIR
80    pub eq_sharp_ns: [T; D_EF],
81
82    /// For folding with mu.
83    pub cur_sum: [T; D_EF],
84    pub mu: [T; D_EF],
85    pub multiplier: [T; D_EF],
86
87    /// Need to know n as if n<0, we need to multiply some norm_factor.
88    pub n_abs: T,
89    pub n_abs_pow: T,
90    pub n_sign: T,
91    /// The round idx for final sumcheck claim.
92    pub num_multilinear_sumcheck_rounds: T,
93}
94
95#[derive(ColumnsAir)]
96#[columns_via(ExpressionClaimCols<u8>)]
97pub struct ExpressionClaimAir {
98    pub expression_claim_n_max_bus: ExpressionClaimNMaxBus,
99    pub expr_claim_bus: ExpressionClaimBus,
100    pub mu_bus: BatchConstraintConductorBus,
101    pub sumcheck_claim_bus: SumcheckClaimBus,
102    pub eq_n_outer_bus: EqNOuterBus,
103    pub pow_checker_bus: PowerCheckerBus,
104    pub hyperdim_bus: HyperdimBus,
105}
106
107impl<F> BaseAirWithPublicValues<F> for ExpressionClaimAir {}
108impl<F> PartitionedBaseAir<F> for ExpressionClaimAir {}
109
110impl<F> BaseAir<F> for ExpressionClaimAir {
111    fn width(&self) -> usize {
112        ExpressionClaimCols::<F>::width()
113    }
114}
115
116impl<AB: AirBuilder + InteractionBuilder> Air<AB> for ExpressionClaimAir
117where
118    <AB::Expr as PrimeCharacteristicRing>::PrimeSubfield: BinomiallyExtendable<{ D_EF }>,
119{
120    fn eval(&self, builder: &mut AB) {
121        let main = builder.main();
122        let (local, next) = (
123            main.row_slice(0).expect("window should have two elements"),
124            main.row_slice(1).expect("window should have two elements"),
125        );
126
127        let local: &ExpressionClaimCols<AB::Var> = (*local).borrow();
128        let next: &ExpressionClaimCols<AB::Var> = (*next).borrow();
129
130        // === Loop structure via NestedForLoopSubAir<4> ===
131        // Enforces canonical nested enumeration for:
132        // [proof_idx, group_idx, trace_idx, idx_parity]
133        // with first-flags:
134        // [is_first, is_first_in_group, is_valid - idx_parity, is_valid].
135        type LoopSubAir = NestedForLoopSubAir<4>;
136        let local_is_trace_start = local.is_valid - local.idx_parity;
137        let next_is_trace_start = next.is_valid - next.idx_parity;
138        LoopSubAir {}.eval(
139            builder,
140            (
141                NestedForLoopIoCols {
142                    is_enabled: local.is_valid.into(),
143                    counter: [
144                        local.proof_idx.into(),
145                        local.group_idx.into(),
146                        local.trace_idx.into(),
147                        local.idx_parity.into(),
148                    ],
149                    is_first: [
150                        local.is_first.into(),
151                        local.is_first_in_group.into(),
152                        local_is_trace_start,
153                        local.is_valid.into(),
154                    ],
155                },
156                NestedForLoopIoCols {
157                    is_enabled: next.is_valid.into(),
158                    counter: [
159                        next.proof_idx.into(),
160                        next.group_idx.into(),
161                        next.trace_idx.into(),
162                        next.idx_parity.into(),
163                    ],
164                    is_first: [
165                        next.is_first.into(),
166                        next.is_first_in_group.into(),
167                        next_is_trace_start,
168                        next.is_valid.into(),
169                    ],
170                },
171            ),
172        );
173
174        // Derived expressions:
175        // is_interaction: true for interaction group (group_idx == 0)
176        let is_interaction: AB::Expr = AB::Expr::ONE - local.group_idx.into();
177        // is_same_proof: next row is valid and within the same proof
178        let is_same_proof: AB::Expr = LoopSubAir::local_is_transition(next.is_valid, next.is_first);
179        // is_last_in_proof: current row is the last row of its proof
180        let is_last_in_proof: AB::Expr =
181            LoopSubAir::local_is_last(local.is_valid, next.is_valid, next.is_first);
182
183        // Each proof starts with group 0 (interactions) and ends with 1 (constraints).
184        // Start with group 0 is guaranteed by NestedForLoop
185        builder
186            .when(is_last_in_proof.clone())
187            .assert_one(local.group_idx);
188
189        // === Claim indexing constraints ===
190        builder.assert_bool(local.n_sign);
191        // idx_parity alternates 0/1.
192        builder
193            .when(local.is_valid)
194            .when(is_interaction.clone())
195            .assert_eq(local.idx_parity + next.idx_parity, AB::Expr::ONE);
196        // only group 0 can have idx_parity set.
197        builder.when(local.idx_parity).assert_zero(local.group_idx);
198
199        // idx binding to trace_idx / idx_parity
200        // Interaction rows: idx = 2 * trace_idx + idx_parity
201        builder.when(is_interaction.clone()).assert_eq(
202            local.idx,
203            local.trace_idx * AB::Expr::TWO + local.idx_parity,
204        );
205        // Constraint rows: idx = trace_idx
206        builder
207            .when(local.group_idx)
208            .assert_eq(local.idx, local.trace_idx);
209
210        // === mu constancy within a proof ===
211        assert_array_eq(
212            &mut builder.when(is_same_proof.clone()),
213            next.mu,
214            local.mu.map(Into::into),
215        );
216
217        // === Hyperdim metadata constancy within numerator/denominator pairs ===
218        // A numerator row (idx_parity=0, is_interaction=1) is always followed by its
219        // denominator (idx_parity=1) due to the alternation constraint. Ensure they
220        // share the same trace metadata so the hyperdim lookup on the numerator binds both.
221        builder
222            .when(local.is_valid)
223            .when(is_interaction.clone())
224            .when(not(local.idx_parity))
225            .assert_eq(next.n_abs, local.n_abs);
226        builder
227            .when(local.is_valid)
228            .when(is_interaction.clone())
229            .when(not(local.idx_parity))
230            .assert_eq(next.n_sign, local.n_sign);
231
232        // === cum sum folding ===
233        // Fold recurrence within a proof: cur_sum = value * multiplier + next_cur_sum * mu
234        assert_array_eq(
235            &mut builder.when(is_same_proof),
236            local.cur_sum,
237            ext_field_add::<AB::Expr>(
238                ext_field_multiply::<AB::Expr>(local.value, local.multiplier),
239                ext_field_multiply::<AB::Expr>(next.cur_sum, local.mu),
240            ),
241        );
242        // Terminal base case: last row of each proof's fold
243        assert_array_eq(
244            &mut builder.when(is_last_in_proof),
245            local.cur_sum,
246            ext_field_multiply::<AB::Expr>(local.value, local.multiplier),
247        );
248
249        // multiplier = 1 if not interaction
250        assert_array_eq(
251            &mut builder.when(local.group_idx).when(local.is_valid),
252            local.multiplier,
253            base_to_ext::<AB::Expr>(AB::Expr::ONE),
254        );
255
256        // IF negative n and numerator
257        assert_array_eq(
258            &mut builder.when(local.n_sign * (is_interaction.clone() - local.idx_parity)),
259            ext_field_multiply_scalar::<AB::Expr>(local.multiplier, local.n_abs_pow),
260            local.eq_sharp_ns,
261        );
262        // ELSE 1: positive n, interaction row
263        assert_array_eq(
264            &mut builder.when(is_interaction.clone() * (AB::Expr::ONE - local.n_sign)),
265            local.multiplier,
266            local.eq_sharp_ns,
267        );
268        // ELSE 2: denominator row
269        assert_array_eq(
270            &mut builder.when(local.idx_parity),
271            local.multiplier,
272            local.eq_sharp_ns,
273        );
274
275        // === bus interactions ===
276        self.expr_claim_bus.receive(
277            builder,
278            local.proof_idx,
279            ExpressionClaimMessage {
280                is_interaction: is_interaction.clone(),
281                idx: local.idx.into(),
282                value: local.value.map(Into::into),
283            },
284            local.is_valid,
285        );
286
287        self.mu_bus.lookup_key(
288            builder,
289            local.proof_idx,
290            BatchConstraintConductorMessage {
291                msg_type: BatchConstraintInnerMessageType::Mu.to_field(),
292                idx: AB::Expr::ZERO,
293                value: local.mu.map(Into::into),
294            },
295            local.is_first * local.is_valid,
296        );
297
298        // Receive n_max value from proof shape air
299        self.expression_claim_n_max_bus.receive(
300            builder,
301            local.proof_idx,
302            ExpressionClaimNMaxMessage {
303                n_max: local.num_multilinear_sumcheck_rounds,
304            },
305            local.is_first * local.is_valid,
306        );
307
308        self.sumcheck_claim_bus.receive(
309            builder,
310            local.proof_idx,
311            SumcheckClaimMessage::<AB::Expr> {
312                round: local.num_multilinear_sumcheck_rounds.into(),
313                value: local.cur_sum.map(Into::into),
314            },
315            local.is_first * local.is_valid,
316        );
317
318        self.hyperdim_bus.lookup_key(
319            builder,
320            local.proof_idx,
321            HyperdimBusMessage {
322                sort_idx: local.trace_idx.into(),
323                n_abs: local.n_abs.into(),
324                n_sign_bit: local.n_sign.into(),
325            },
326            local.is_valid * (is_interaction.clone() - local.idx_parity),
327        );
328
329        self.eq_n_outer_bus.lookup_key(
330            builder,
331            local.proof_idx,
332            EqNOuterMessage {
333                is_sharp: AB::Expr::ONE,
334                n: local.n_abs * (AB::Expr::ONE - local.n_sign),
335                value: local.eq_sharp_ns.map(Into::into),
336            },
337            local.is_valid * is_interaction.clone(),
338        );
339
340        self.pow_checker_bus.lookup_key(
341            builder,
342            PowerCheckerBusMessage {
343                log: local.n_abs.into(),
344                exp: local.n_abs_pow.into(),
345            },
346            local.is_valid * is_interaction,
347        );
348    }
349}
openvm_recursion_circuit/batch_constraint/expression_claim/air.rs

openvm_recursion_circuit/batch_constraint/expression_claim/
air.rs
