fix: expression claim constraints (#2498)

luffykai · web-flow · commit 9c6ccf36821a · 2026-03-09T12:04:13.000-07:00
fix: INT-6590

fix: INT-6396, INT-6397, INT-6398, INT-6399, INT-6402, INT-6439
diff --git a/crates/recursion/src/batch_constraint/expression_claim/air.rs b/crates/recursion/src/batch_constraint/expression_claim/air.rs
@@ -1,6 +1,9 @@
 use std::borrow::Borrow;
 
-use openvm_circuit_primitives::utils::{assert_array_eq, not};
+use openvm_circuit_primitives::{
+    utils::{assert_array_eq, not},
+    SubAir,
+};
 use openvm_stark_backend::{
     interaction::InteractionBuilder, BaseAirWithPublicValues, PartitionedBaseAir,
 };
@@ -18,24 +21,56 @@ use crate::{
     },
     bus::{ExpressionClaimNMaxBus, ExpressionClaimNMaxMessage, HyperdimBus, HyperdimBusMessage},
     primitives::bus::{PowerCheckerBus, PowerCheckerBusMessage},
+    subairs::nested_for_loop::{NestedForLoopIoCols, NestedForLoopSubAir},
     utils::{base_to_ext, ext_field_add, ext_field_multiply, ext_field_multiply_scalar},
 };
 
 /// For each proof, this AIR will receive 2t interaction claims and t constraint claims.
 /// (2 interaction claims and 1 constraint claim per trace).
 /// These values are folded (algebraic batching) with mu into a single value, which
 /// should match the final sumcheck claim.
+///
+/// Rows are structured as a nested loop: for each proof, group 0 (interactions) comes first,
+/// then group 1 (constraints). Within the interaction group, each trace occupies 2 rows
+/// (numerator then denominator). Within the constraint group, each trace occupies 1 row.
+/// `NestedForLoopSubAir<4>` enforces canonical nested enumeration over
+/// `(proof_idx, group_idx, trace_idx, idx_parity)`.
+///
+/// Example for t = 2 traces (one proof):
+///
+/// ```text
+/// row | is_first | group_idx | is_first_in_group | trace_idx | idx_parity | idx
+/// ----|----------|-----------|-------------------|-----------|------------|----
+///   0 |    1     |     0     |         1         |     0     |     0      |  0   ← numerator trace 0
+///   1 |    0     |     0     |         0         |     0     |     1      |  1   ← denominator trace 0
+///   2 |    0     |     0     |         0         |     1     |     0      |  2   ← numerator trace 1
+///   3 |    0     |     0     |         0         |     1     |     1      |  3   ← denominator trace 1
+///   4 |    0     |     1     |         1         |     0     |     0      |  0   ← constraint trace 0
+///   5 |    0     |     1     |         0         |     1     |     0      |  1   ← constraint trace 1
+/// ```
+///
+/// `is_interaction` is derived as `1 - group_idx` (not a separate column).
 #[derive(AlignedBorrow, Copy, Clone, Debug)]
 #[repr(C)]
 pub struct ExpressionClaimCols<T> {
+    // --- Loop structure (enforced by NestedForLoopSubAir<4>) ---
     pub is_valid: T,
+    /// First row of a proof. Marks proof boundaries.
     pub is_first: T,
     pub proof_idx: T,
+    /// 0 = interaction group, 1 = constraint group. Monotone within a proof.
+    pub group_idx: T,
+    /// Marks the first row of each group (set at proof start and interaction→constraint boundary).
+    pub is_first_in_group: T,
 
-    pub is_interaction: T,
-    /// Index within the proof, 0 ~ 2t-1 are interaction claims, 0~t-1 are constraint claims.
+    // --- Claim indexing (derived from loop counters) ---
+    /// Claim index within its group. For interactions: `2 * trace_idx + idx_parity` (0..2t).
+    /// For constraints: `trace_idx` (0..t).
     pub idx: T,
+    /// 0 = numerator, 1 = denominator. Always 0 on constraint rows. Alternates on interaction
+    /// rows.
     pub idx_parity: T,
+    /// Sorted trace index within the group. Monotone non-decreasing; resets at group boundaries.
     pub trace_idx: T,
     /// The received evaluation claim. Note that for interactions, this is without norm_factor and
     /// eq_sharp_ns. These are interactions_evals (without norm_factor and eq_sharp_ns) and
@@ -90,66 +125,159 @@ where
         let local: &ExpressionClaimCols<AB::Var> = (*local).borrow();
         let next: &ExpressionClaimCols<AB::Var> = (*next).borrow();
 
-        builder.assert_bool(local.is_valid);
-        builder.assert_bool(local.is_first);
-        builder.assert_bool(local.is_interaction);
-        builder.assert_bool(local.idx_parity);
-        builder.assert_bool(local.n_sign);
+        // === Loop structure via NestedForLoopSubAir<4> ===
+        // Enforces canonical nested enumeration for:
+        // [proof_idx, group_idx, trace_idx, idx_parity]
+        // with first-flags:
+        // [is_first, is_first_in_group, is_valid - idx_parity, is_valid].
+        type LoopSubAir = NestedForLoopSubAir<4>;
+        let local_is_trace_start = local.is_valid - local.idx_parity;
+        let next_is_trace_start = next.is_valid - next.idx_parity;
+        LoopSubAir {}.eval(
+            builder,
+            (
+                NestedForLoopIoCols {
+                    is_enabled: local.is_valid.into(),
+                    counter: [
+                        local.proof_idx.into(),
+                        local.group_idx.into(),
+                        local.trace_idx.into(),
+                        local.idx_parity.into(),
+                    ],
+                    is_first: [
+                        local.is_first.into(),
+                        local.is_first_in_group.into(),
+                        local_is_trace_start,
+                        local.is_valid.into(),
+                    ],
+                },
+                NestedForLoopIoCols {
+                    is_enabled: next.is_valid.into(),
+                    counter: [
+                        next.proof_idx.into(),
+                        next.group_idx.into(),
+                        next.trace_idx.into(),
+                        next.idx_parity.into(),
+                    ],
+                    is_first: [
+                        next.is_first.into(),
+                        next.is_first_in_group.into(),
+                        next_is_trace_start,
+                        next.is_valid.into(),
+                    ],
+                },
+            ),
+        );
+
+        // Derived expressions:
+        // is_interaction: true for interaction group (group_idx == 0)
+        let is_interaction: AB::Expr = AB::Expr::ONE - local.group_idx.into();
+        // is_same_proof: next row is valid and within the same proof
+        let is_same_proof: AB::Expr = LoopSubAir::local_is_transition(next.is_valid, next.is_first);
+        // is_last_in_proof: current row is the last row of its proof
+        let is_last_in_proof: AB::Expr =
+            LoopSubAir::local_is_last(local.is_valid, next.is_valid, next.is_first);
+
+        // Each proof starts with group 0 (interactions) and ends with 1 (constraints).
+        // Start with group 0 is guaranteed by NestedForLoop
         builder
-            .when(local.is_first)
-            .assert_one(local.is_interaction);
-        builder.when(local.is_first).assert_zero(local.idx_parity);
+            .when(is_last_in_proof.clone())
+            .assert_one(local.group_idx);
+
+        // === Claim indexing constraints ===
+        builder.assert_bool(local.n_sign);
+        // idx_parity alternates 0/1.
         builder
-            .when(local.is_interaction)
+            .when(local.is_valid)
+            .when(is_interaction.clone())
             .assert_eq(local.idx_parity + next.idx_parity, AB::Expr::ONE);
+        // only group 0 can have idx_parity set.
+        builder.when(local.idx_parity).assert_zero(local.group_idx);
+
+        // idx binding to trace_idx / idx_parity
+        // Interaction rows: idx = 2 * trace_idx + idx_parity
+        builder.when(is_interaction.clone()).assert_eq(
+            local.idx,
+            local.trace_idx * AB::Expr::TWO + local.idx_parity,
+        );
+        // Constraint rows: idx = trace_idx
         builder
-            .when(local.idx_parity)
-            .assert_one(local.is_interaction);
+            .when(local.group_idx)
+            .assert_eq(local.idx, local.trace_idx);
+
+        // === mu constancy within a proof ===
+        assert_array_eq(
+            &mut builder.when(is_same_proof.clone()),
+            next.mu,
+            local.mu.map(Into::into),
+        );
+
+        // === Hyperdim metadata constancy within numerator/denominator pairs ===
+        // A numerator row (idx_parity=0, is_interaction=1) is always followed by its
+        // denominator (idx_parity=1) due to the alternation constraint. Ensure they
+        // share the same trace metadata so the hyperdim lookup on the numerator binds both.
+        builder
+            .when(local.is_valid)
+            .when(is_interaction.clone())
+            .when(not(local.idx_parity))
+            .assert_eq(next.n_abs, local.n_abs);
+        builder
+            .when(local.is_valid)
+            .when(is_interaction.clone())
+            .when(not(local.idx_parity))
+            .assert_eq(next.n_sign, local.n_sign);
 
         // === cum sum folding ===
-        // cur_sum = next_cur_sum * mu + value * multiplier
+        // Fold recurrence within a proof: cur_sum = value * multiplier + next_cur_sum * mu
         assert_array_eq(
-            &mut builder.when(local.is_valid * not(next.is_first)),
+            &mut builder.when(is_same_proof),
             local.cur_sum,
             ext_field_add::<AB::Expr>(
                 ext_field_multiply::<AB::Expr>(local.value, local.multiplier),
                 ext_field_multiply::<AB::Expr>(next.cur_sum, local.mu),
             ),
         );
+        // Terminal base case: last row of each proof's fold
+        assert_array_eq(
+            &mut builder.when(is_last_in_proof),
+            local.cur_sum,
+            ext_field_multiply::<AB::Expr>(local.value, local.multiplier),
+        );
+
         // multiplier = 1 if not interaction
         assert_array_eq(
-            &mut builder.when(not(local.is_interaction)).when(local.is_valid),
+            &mut builder.when(local.group_idx).when(local.is_valid),
             local.multiplier,
             base_to_ext::<AB::Expr>(AB::Expr::ONE),
         );
 
         // IF negative n and numerator
         assert_array_eq(
-            &mut builder.when(local.n_sign * (local.is_interaction - local.idx_parity)),
+            &mut builder.when(local.n_sign * (is_interaction.clone() - local.idx_parity)),
             ext_field_multiply_scalar::<AB::Expr>(local.multiplier, local.n_abs_pow),
             local.eq_sharp_ns,
         );
-        // ELSE 1
+        // ELSE 1: positive n, interaction row
         assert_array_eq(
-            &mut builder.when(local.is_interaction * (AB::Expr::ONE - local.n_sign)),
+            &mut builder.when(is_interaction.clone() * (AB::Expr::ONE - local.n_sign)),
             local.multiplier,
             local.eq_sharp_ns,
         );
-        // ELSE 2
+        // ELSE 2: denominator row
         assert_array_eq(
             &mut builder.when(local.idx_parity),
             local.multiplier,
             local.eq_sharp_ns,
         );
 
-        // === interactions ===
+        // === bus interactions ===
         self.expr_claim_bus.receive(
             builder,
             local.proof_idx,
             ExpressionClaimMessage {
-                is_interaction: local.is_interaction,
-                idx: local.idx,
-                value: local.value,
+                is_interaction: is_interaction.clone(),
+                idx: local.idx.into(),
+                value: local.value.map(Into::into),
             },
             local.is_valid,
         );
@@ -193,7 +321,7 @@ where
                 n_abs: local.n_abs.into(),
                 n_sign_bit: local.n_sign.into(),
             },
-            local.is_valid * (local.is_interaction - local.idx_parity),
+            local.is_valid * (is_interaction.clone() - local.idx_parity),
         );
 
         self.eq_n_outer_bus.lookup_key(
@@ -204,7 +332,7 @@ where
                 n: local.n_abs * (AB::Expr::ONE - local.n_sign),
                 value: local.eq_sharp_ns.map(Into::into),
             },
-            local.is_valid * local.is_interaction,
+            local.is_valid * is_interaction.clone(),
         );
 
         self.pow_checker_bus.lookup_key(
@@ -213,7 +341,7 @@ where
                 log: local.n_abs.into(),
                 exp: local.n_abs_pow.into(),
             },
-            local.is_valid * local.is_interaction,
+            local.is_valid * is_interaction,
         );
     }
 }
diff --git a/crates/recursion/src/batch_constraint/expression_claim/trace.rs b/crates/recursion/src/batch_constraint/expression_claim/trace.rs
@@ -92,7 +92,11 @@ impl RowMajorChip<F> for ExpressionClaimTraceGenerator {
                     cols.is_first = F::from_bool(i == 0);
                     cols.is_valid = F::ONE;
                     cols.proof_idx = F::from_usize(pidx);
-                    cols.is_interaction = F::from_bool(is_interaction);
+                    // group_idx: 0 for interactions, 1 for constraints
+                    cols.group_idx = F::from_bool(!is_interaction);
+                    // is_first_in_group: true at start of proof (i==0) and at start of
+                    // constraint group (i == 2*num_present)
+                    cols.is_first_in_group = F::from_bool(i == 0 || i == 2 * num_present);
                     cols.num_multilinear_sumcheck_rounds = F::from_usize(num_rounds);
                     cols.idx = F::from_usize(if i < 2 * num_present {
                         i
@@ -133,7 +137,7 @@ impl RowMajorChip<F> for ExpressionClaimTraceGenerator {
                 .for_each(|chunk| {
                     let cols: &mut ExpressionClaimCols<_> = chunk.borrow_mut();
                     // if it's interaction, we need to multiply by eq_sharp_ns and norm_factor
-                    let multiplier = if cols.is_interaction == F::ONE {
+                    let multiplier = if cols.group_idx == F::ZERO {
                         let mut mult =
                             EF::from_basis_coefficients_slice(&cols.eq_sharp_ns).unwrap();
                         if cols.n_sign == F::ONE && cols.idx.as_canonical_u32() % 2 == 0 {
@@ -153,13 +157,6 @@ impl RowMajorChip<F> for ExpressionClaimTraceGenerator {
 
             cur_height += claims.len();
         }
-        trace[cur_height * width..]
-            .par_chunks_mut(width)
-            .enumerate()
-            .for_each(|(i, chunk)| {
-                let cols: &mut ExpressionClaimCols<F> = chunk.borrow_mut();
-                cols.proof_idx = F::from_usize(preflights.len() + i);
-            });
         Some(RowMajorMatrix::new(trace, width))
     }
 }
