feat(stark): prune redundant g·z (next-row) trace-OOD openings

crypto/stark/src/lib.rs

@@ -23,6 +23,7 @@ pub mod instruments;
 #[cfg(feature = "cuda")]
 pub mod logup_gpu;
 pub mod lookup;
+pub mod ood;
 pub(crate) mod par;
 pub mod profile_markers;
 pub mod proof;

crypto/stark/src/lookup.rs

@@ -1003,6 +1003,19 @@ where
         self.trace_layout
     }
 
+    fn trace_ood_next_row_columns(&self) -> Vec<usize> {
+        // The only transition constraint that reads the next row is the circular
+        // LogUp accumulator, and after forward accumulation it reads only the
+        // accumulated column there (all committed terms and absorbed operands
+        // read the current row). Its full-width index is the main width plus the
+        // accumulated column's aux index. No interactions => no next-row reads.
+        if self.auxiliary_trace_build_data.interactions.is_empty() {
+            Vec::new()
+        } else {
+            vec![self.trace_layout.0 + self.logup.acc_column_idx]
+        }
+    }
+
     fn has_trace_interaction(&self) -> bool {
         !self.auxiliary_trace_build_data.interactions.is_empty()
     }

crypto/stark/src/ood.rs

@@ -0,0 +1,224 @@
+//! Shared, prover = verifier-identical helpers for out-of-domain (OOD) trace
+//! opening pruning.
+//!
+//! The frame OOD table has `num_offsets * step_size` rows (offsets `[0, 1]`,
+//! offset-major: the first `step_size` rows are the current-row block, the rest
+//! are next-row blocks) and one column per trace column. Only the columns a
+//! transition constraint actually reads at the next row — the AIR's transition
+//! window, [`crate::traits::AIR::trace_ood_next_row_columns`] — need to be
+//! opened in the next-row block(s). Every other next-row entry is redundant and
+//! is pruned from the proof.
+//!
+//! Everything here is a pure function of public AIR shape metadata (`step_size`,
+//! the column count, and the next-row column set), so the prover and verifier
+//! derive the identical layout without trusting proof dimensions (invariant I3).
+
+use crate::table::Table;
+use math::field::{element::FieldElement, traits::IsField};
+
+/// Per-column flags: `flags[c] == true` iff column `c` is opened at the next
+/// row. Indices outside `0..num_total_cols` are ignored.
+pub fn next_row_col_flags(num_total_cols: usize, next_row_cols: &[usize]) -> Vec<bool> {
+    let mut flags = vec![false; num_total_cols];
+    for &c in next_row_cols {
+        if c < num_total_cols {
+            flags[c] = true;
+        }
+    }
+    flags
+}
+
+/// Number of surviving trace openings: the current-row block opens every column
+/// (`step_size * num_total_cols`), and each next-row row opens only the masked
+/// columns (`(num_eval_points - step_size) * num_next_row_cols`).
+pub fn num_surviving_trace_openings(
+    num_total_cols: usize,
+    num_eval_points: usize,
+    step_size: usize,
+    num_next_row_cols: usize,
+) -> usize {
+    let next_rows = num_eval_points.saturating_sub(step_size);
+    step_size * num_total_cols + next_rows * num_next_row_cols
+}
+
+/// Build the rectangular `num_total_cols x num_eval_points` DEEP trace-term
+/// coefficient grid from `powers` (the `num_surviving_trace_openings` gamma
+/// powers drained for the trace terms). Surviving positions receive a power in a
+/// fixed order; pruned next-row positions receive zero. A rectangular DEEP
+/// evaluation over the full grid therefore yields the identical polynomial as
+/// summing only the survivors — which is what lets the prover keep its
+/// (GPU-friendly) rectangular DEEP unchanged.
+///
+/// Assignment order (mirrored exactly by [`num_surviving_trace_openings`]):
+///   1. current-row block — for every column `j`, rows `0..step_size`;
+///   2. next-row block — for each masked column `j`, rows `step_size..num_eval_points`.
+pub fn build_pruned_trace_term_coeffs<E: IsField>(
+    powers: &[FieldElement<E>],
+    num_total_cols: usize,
+    num_eval_points: usize,
+    step_size: usize,
+    next_row_cols: &[usize],
+) -> Vec<Vec<FieldElement<E>>> {
+    let flags = next_row_col_flags(num_total_cols, next_row_cols);
+    let mut coeffs = vec![vec![FieldElement::<E>::zero(); num_eval_points]; num_total_cols];
+    let mut p = 0usize;
+    // Current-row block: all columns, rows 0..step_size.
+    for col in coeffs.iter_mut() {
+        for slot in col.iter_mut().take(step_size) {
+            if p < powers.len() {
+                *slot = powers[p].clone();
+                p += 1;
+            }
+        }
+    }
+    // Next-row block(s): masked columns only, rows step_size..num_eval_points.
+    for (j, col) in coeffs.iter_mut().enumerate() {
+        if flags[j] {
+            for slot in col.iter_mut().take(num_eval_points).skip(step_size) {
+                if p < powers.len() {
+                    *slot = powers[p].clone();
+                    p += 1;
+                }
+            }
+        }
+    }
+    debug_assert_eq!(p, powers.len(), "power assignment must consume every power");
+    coeffs
+}
+
+/// Split the full `num_eval_points x num_total_cols` OOD table (computed by the
+/// prover) into the two blocks carried by the proof:
+///   * block 0 — the current-row block, `step_size x num_total_cols` (all columns);
+///   * block 1 — the next-row block, `next_rows x num_next_row_cols`, holding only
+///     the masked columns in `next_row_cols` order.
+///
+/// Block 1 has width 0 (an empty table) when the AIR reads no next-row columns.
+pub fn split_ood_blocks<E: IsField>(
+    full: &Table<E>,
+    step_size: usize,
+    next_row_cols: &[usize],
+) -> (Table<E>, Table<E>) {
+    let w = full.width;
+
+    let mut b0 = Vec::with_capacity(step_size * w);
+    for r in 0..step_size {
+        b0.extend_from_slice(full.get_row(r));
+    }
+    let block0 = Table::new(b0, w);
+
+    let mut b1 = Vec::with_capacity((full.height.saturating_sub(step_size)) * next_row_cols.len());
+    for r in step_size..full.height {
+        let row = full.get_row(r);
+        for &c in next_row_cols {
+            b1.push(row[c].clone());
+        }
+    }
+    let block1 = Table::new(b1, next_row_cols.len());
+
+    (block0, block1)
+}
+
+/// Rebuild the full `num_eval_points x num_total_cols` OOD table from the two
+/// pruned proof blocks. Current-row rows come straight from `block0`; each
+/// next-row row scatters the masked values from `block1` into their columns and
+/// leaves every other column zero. Those zero entries are never read — no
+/// transition constraint references a pruned column at the next row, and DEEP
+/// skips them — so the reconstruction is exact where it matters.
+pub fn reconstruct_ood_full<E: IsField>(
+    block0: &Table<E>,
+    block1: &Table<E>,
+    num_eval_points: usize,
+    step_size: usize,
+    next_row_cols: &[usize],
+) -> Table<E> {
+    let w = block0.width;
+    let mut data = Vec::with_capacity(num_eval_points * w);
+
+    for r in 0..step_size {
+        data.extend_from_slice(block0.get_row(r));
+    }
+
+    for r in step_size..num_eval_points {
+        let has_next = block1.width > 0 && (r - step_size) < block1.height;
+        for c in 0..w {
+            let mut val = FieldElement::<E>::zero();
+            if has_next {
+                for (m, &mc) in next_row_cols.iter().enumerate() {
+                    if mc == c {
+                        val = block1.get_row(r - step_size)[m].clone();
+                        break;
+                    }
+                }
+            }
+            data.push(val);
+        }
+    }
+
+    Table::new(data, w)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use math::field::goldilocks::GoldilocksField as Gl;
+
+    type Fe = FieldElement<Gl>;
+
+    fn fe(x: u64) -> Fe {
+        Fe::from(x)
+    }
+
+    #[test]
+    fn surviving_count_matches_layout() {
+        // 3 columns, 2 eval points (step_size 1), 1 next-row column:
+        // current-row opens 3, next-row opens 1 => 4.
+        assert_eq!(num_surviving_trace_openings(3, 2, 1, 1), 4);
+        // No next-row columns => only the current-row block survives.
+        assert_eq!(num_surviving_trace_openings(3, 2, 1, 0), 3);
+        // Every column open at the next row => full 2*W grid.
+        assert_eq!(num_surviving_trace_openings(3, 2, 1, 3), 6);
+    }
+
+    #[test]
+    fn split_then_reconstruct_preserves_survivors_and_zeros_pruned() {
+        // Full 2x3 OOD table: row 0 (current row), row 1 (next row).
+        let full = Table::new(vec![fe(10), fe(11), fe(12), fe(20), fe(21), fe(22)], 3);
+        let next_row_cols = [1usize]; // only column 1 opens at the next row
+        let step_size = 1;
+
+        let (b0, b1) = split_ood_blocks(&full, step_size, &next_row_cols);
+        assert_eq!((b0.width, b0.height), (3, 1));
+        assert_eq!((b1.width, b1.height), (1, 1));
+        assert_eq!(b1.get_row(0)[0], fe(21)); // full[1][1]
+
+        let recon = reconstruct_ood_full(&b0, &b1, 2, step_size, &next_row_cols);
+        assert_eq!(recon.get_row(0), full.get_row(0)); // current row is exact
+        assert_eq!(recon.get_row(1)[1], fe(21)); // survivor placed
+        assert_eq!(recon.get_row(1)[0], Fe::zero()); // pruned -> zero
+        assert_eq!(recon.get_row(1)[2], Fe::zero()); // pruned -> zero
+    }
+
+    #[test]
+    fn empty_next_row_block_reconstructs_to_zeros() {
+        let full = Table::new(vec![fe(10), fe(11), fe(20), fe(21)], 2);
+        let (b0, b1) = split_ood_blocks(&full, 1, &[]);
+        assert_eq!(b1.width, 0);
+        let recon = reconstruct_ood_full(&b0, &b1, 2, 1, &[]);
+        assert_eq!(recon.get_row(0), full.get_row(0));
+        assert_eq!(recon.get_row(1), &[Fe::zero(), Fe::zero()]);
+    }
+
+    #[test]
+    fn pruned_coeffs_are_zero_off_the_window() {
+        // 4 surviving powers for W=3, num_eval_points=2, mask={1}.
+        let powers: Vec<Fe> = (1..=4).map(fe).collect();
+        let coeffs = build_pruned_trace_term_coeffs(&powers, 3, 2, 1, &[1]);
+        // Current-row row (k=0) is fully populated; next-row row (k=1) only col 1.
+        assert_ne!(coeffs[0][0], Fe::zero());
+        assert_ne!(coeffs[1][0], Fe::zero());
+        assert_ne!(coeffs[2][0], Fe::zero());
+        assert_ne!(coeffs[1][1], Fe::zero()); // masked column, next row
+        assert_eq!(coeffs[0][1], Fe::zero()); // pruned
+        assert_eq!(coeffs[2][1], Fe::zero()); // pruned
+    }
+}

crypto/stark/src/proof/stark.rs

@@ -49,8 +49,12 @@ pub struct StarkProof<F: IsSubFieldOf<E>, E: IsField, PI> {
     // For preprocessed tables: commitment to precomputed columns only.
     // Verifier checks this matches the hardcoded commitment from AIR.
     pub lde_trace_precomputed_merkle_root: Option<Commitment>,
-    // tⱼ(zgᵏ)
+    // tⱼ(zgᵏ) for the current-row block (offset 0): every trace column at z.
     pub trace_ood_evaluations: Table<E>,
+    // tⱼ(zgᵏ) for the next-row block(s) (offset >= 1), pruned to only the columns
+    // a transition constraint reads at the next row (the AIR transition window).
+    // Empty (width 0) when the AIR reads no next-row columns.
+    pub trace_ood_next_evaluations: Table<E>,
     // Commitments to Hᵢ
     pub composition_poly_root: Commitment,
     // Hᵢ(z^N)

crypto/stark/src/prover.rs

@@ -1458,21 +1458,36 @@ pub trait IsStarkProver<
         let gamma = transcript.sample_field_element();
 
         let n_terms_composition_poly = round_2_result.lde_composition_poly_evaluations.len();
-        let num_terms_trace =
-            air.context().transition_offsets.len() * air.step_size() * air.context().trace_columns;
+        let num_eval_points = air.context().transition_offsets.len() * air.step_size();
+        let next_row_cols = air.trace_ood_next_row_columns();
+        // g·z pruning: only the current-row block (all columns) plus the masked
+        // next-row columns get an opening / DEEP coefficient.
+        let num_terms_trace = crate::ood::num_surviving_trace_openings(
+            air.context().trace_columns,
+            num_eval_points,
+            air.step_size(),
+            next_row_cols.len(),
+        );
 
         // <<<< Receive challenges: 𝛾, 𝛾'
         let mut deep_composition_coefficients: Vec<_> =
             core::iter::successors(Some(FieldElement::one()), |x| Some(x * &gamma))
                 .take(n_terms_composition_poly + num_terms_trace)
                 .collect();
 
-        let trace_term_coeffs: Vec<_> = deep_composition_coefficients
+        let trace_term_powers: Vec<_> = deep_composition_coefficients
             .drain(..num_terms_trace)
-            .collect::<Vec<_>>()
-            .chunks(air.context().transition_offsets.len() * air.step_size())
-            .map(|chunk| chunk.to_vec())
             .collect();
+        // Rectangular W×num_eval_points grid with the sampled powers at surviving
+        // positions and zeros at pruned next-row positions, so the DEEP loop
+        // below (and the GPU path) stay unchanged — zero-coefficient terms vanish.
+        let trace_term_coeffs = crate::ood::build_pruned_trace_term_coeffs(
+            &trace_term_powers,
+            air.context().trace_columns,
+            num_eval_points,
+            air.step_size(),
+            &next_row_cols,
+        );
 
         // <<<< Receive challenges: 𝛾ⱼ, 𝛾ⱼ'
         let gammas = deep_composition_coefficients;
@@ -3103,11 +3118,21 @@ pub trait IsStarkProver<
         #[cfg(feature = "instruments")]
         let round_3_dur = t_r3.elapsed();
 
-        // >>>> Send values: tⱼ(zgᵏ)
-        let trace_ood_evaluations_columns = round_3_result.trace_ood_evaluations.columns();
-        for col in trace_ood_evaluations_columns.iter() {
-            for elem in col.iter() {
-                transcript.append_field_element(elem);
+        // >>>> Send values: tⱼ(zgᵏ). g·z pruning: split the full OOD table into
+        // the current-row block (all columns) and the pruned next-row block
+        // (masked columns only), and absorb only the surviving values — the
+        // verifier absorbs the identical two blocks in the same order.
+        let ood_next_row_cols = air.trace_ood_next_row_columns();
+        let (ood_block0, ood_block1) = crate::ood::split_ood_blocks(
+            &round_3_result.trace_ood_evaluations,
+            air.step_size(),
+            &ood_next_row_cols,
+        );
+        for block in [&ood_block0, &ood_block1] {
+            for col in block.columns().iter() {
+                for elem in col.iter() {
+                    transcript.append_field_element(elem);
+                }
             }
         }
 
@@ -3161,8 +3186,9 @@ pub trait IsStarkProver<
             lde_trace_aux_merkle_root: round_1_result.aux.as_ref().map(|x| x.root),
             // For preprocessed tables: commitment to precomputed columns only
             lde_trace_precomputed_merkle_root: round_1_result.main.precomputed_root,
-            // tⱼ(zgᵏ)
-            trace_ood_evaluations: round_3_result.trace_ood_evaluations,
+            // tⱼ(zgᵏ): current-row block + pruned next-row block.
+            trace_ood_evaluations: ood_block0,
+            trace_ood_next_evaluations: ood_block1,
             // [H₁] and [H₂]
             composition_poly_root: round_2_result.composition_poly_root,
             // Hᵢ(z^N)