perf(verifier): factor shared zerofier + batch DEEP inversions

crypto/stark/src/constraints/zerofier.rs

@@ -115,10 +115,16 @@ pub fn zerofier_evaluations_on_extended_domain<F: IsFFTField>(
         .collect()
 }
 
-/// Evaluation of the constraint's zerofier at some point `z`, which may be in
-/// a field extension.
-pub fn evaluate_zerofier<F, E>(
-    meta: &ConstraintMeta,
+/// The end-exemptions correction `∏(z − rᵢ)` at `z`, where `rᵢ` are the roots for a
+/// constraint skipping its last `end_exemptions` rows (`1` when there are none).
+///
+/// This is the only per-constraint-varying factor of the transition zerofier at
+/// `z`: the full inverse zerofier is `1/(zᴺ − 1)` × this. Exposed separately so a
+/// caller evaluating many constraints at the same `z` computes `1/(zᴺ − 1)` once
+/// and this once per distinct `end_exemptions`, rather than a fresh `zᴺ` power and
+/// extension inversion per constraint (see the verifier's OOD zerofier sum).
+pub fn end_exemptions_correction<F, E>(
+    end_exemptions: usize,
     z: &FieldElement<E>,
     trace_primitive_root: &FieldElement<F>,
     trace_length: usize,
@@ -127,16 +133,78 @@ where
     F: IsSubFieldOf<E>,
     E: IsField,
 {
-    let roots = end_exemptions_roots(meta, trace_primitive_root, trace_length);
-    // Factor `z - rᵢ` written as `-(rᵢ - z)`: the field ops only go
-    // subfield − superfield, and `rᵢ ∈ F`, `z ∈ E`.
-    let end_exemptions_eval = roots.iter().fold(FieldElement::<E>::one(), |acc, root| {
-        acc * -(root.clone() - z.clone())
-    });
-
-    // 1/(z^N − 1), times the end-exemptions correction.
-    (-FieldElement::<F>::one() + z.pow(trace_length))
-        .inv()
-        .unwrap()
-        * &end_exemptions_eval
+    if end_exemptions == 0 {
+        return FieldElement::<E>::one();
+    }
+    // Roots are gᴺ⁻¹, g²⁽ᴺ⁻¹⁾, … (walking backward from the last row by
+    // g⁻¹ = gᴺ⁻¹). Written `-(rᵢ - z)` so the field ops only go subfield −
+    // superfield (`rᵢ ∈ F`, `z ∈ E`), matching `end_exemptions_roots`.
+    let decrement = trace_primitive_root.pow(trace_length - 1);
+    let mut current = decrement.clone();
+    let mut acc = FieldElement::<E>::one();
+    for _ in 0..end_exemptions {
+        acc *= -(current.clone() - z.clone());
+        current = &current * &decrement;
+    }
+    acc
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField;
+    use math::field::goldilocks::GoldilocksField;
+
+    type F = GoldilocksField;
+    type E = Degree3GoldilocksExtensionField;
+
+    /// The exempt-row product `∏(z − gⁱ)` over the last `end_exemptions` rows
+    /// (indices `N−e .. N−1`), derived directly from `gⁱ` — the value
+    /// [`end_exemptions_correction`] must equal. Independent of its backward
+    /// `g⁻¹` walk, so a mismatch catches an error in that root derivation.
+    fn exempt_product(
+        end_exemptions: usize,
+        z: &FieldElement<E>,
+        g: &FieldElement<F>,
+        n: usize,
+    ) -> FieldElement<E> {
+        let mut acc = FieldElement::<E>::one();
+        for i in (n - end_exemptions)..n {
+            // `-(gⁱ − z) = z − gⁱ`, keeping the ops subfield − superfield to
+            // match the production body.
+            acc *= -(g.pow(i) - *z);
+        }
+        acc
+    }
+
+    #[test]
+    fn correction_matches_direct_exempt_product_over_cubic_extension() {
+        let n = 16usize;
+        let g = F::get_primitive_root_of_unity(n.trailing_zeros() as u64).unwrap();
+        // A point off the trace domain, genuinely inside the cubic extension.
+        let z = FieldElement::<E>::new([
+            FieldElement::<F>::from(7u64),
+            FieldElement::<F>::from(3u64),
+            FieldElement::<F>::from(1u64),
+        ]);
+
+        // 0 → one() (no exemptions); 1..=3 exercise the multi-group fold the
+        // verifier drives, one distinct `end_exemptions` per group.
+        for end_exemptions in 0..=3usize {
+            let got = end_exemptions_correction::<F, E>(end_exemptions, &z, &g, n);
+            let want = exempt_product(end_exemptions, &z, &g, n);
+            assert_eq!(got, want, "mismatch for end_exemptions = {end_exemptions}");
+        }
+    }
+
+    #[test]
+    fn correction_with_no_exemptions_is_one() {
+        let n = 8usize;
+        let g = F::get_primitive_root_of_unity(n.trailing_zeros() as u64).unwrap();
+        let z = FieldElement::<E>::from(5u64);
+        assert_eq!(
+            end_exemptions_correction::<F, E>(0, &z, &g, n),
+            FieldElement::<E>::one()
+        );
+    }
 }

crypto/stark/src/verifier.rs

@@ -328,25 +328,49 @@ pub trait IsStarkVerifier<
         let transition_ood_frame_evaluations =
             air.compute_transition(&transition_evaluation_context);
 
-        let mut denominators =
-            vec![FieldElement::<FieldExtension>::zero(); air.num_transition_constraints()];
-        air.constraints_meta().iter().for_each(|m| {
-            denominators[m.constraint_idx] = crate::constraints::zerofier::evaluate_zerofier(
-                m,
-                &challenges.z,
-                &domain.trace_primitive_root,
-                trace_length,
-            );
-        });
+        // Every transition constraint's zerofier at z is 1/(zᴺ − 1) times an
+        // end-exemptions correction ∏(z − rᵢ) that depends only on the
+        // constraint's `end_exemptions`. 1/(zᴺ − 1) is shared by all of them, so
+        // compute it once, group the constraints by `end_exemptions` (a tiny set —
+        // almost always just {0}), accumulate Σ βᵢ·evalᵢ per group, then factor the
+        // shared inverse and each group's correction out of the sum. Replaces the
+        // previous per-constraint zᴺ power + extension inversion.
+        let inv_zerofier_denominator =
+            match (-FieldElement::<Field>::one() + challenges.z.pow(trace_length)).inv() {
+                Ok(inv) => inv,
+                // z lands on the trace domain (zᴺ = 1); an honest off-domain z never
+                // does. Reject rather than divide by zero.
+                Err(_) => return false,
+            };
 
-        let transition_c_i_evaluations_sum = itertools::izip!(
-            transition_ood_frame_evaluations,
-            &challenges.transition_coeffs,
-            denominators
-        )
-        .fold(FieldElement::zero(), |acc, (eval, beta, denominator)| {
-            acc + beta * eval * &denominator
-        });
+        // Σ βᵢ·evalᵢ bucketed by end_exemptions. `constraints_meta` has one entry per
+        // active transition constraint; a constraint absent from it contributed a
+        // zero denominator before and is likewise skipped here.
+        let mut grouped_numerator_sums: Vec<(usize, FieldElement<FieldExtension>)> = Vec::new();
+        for m in air.constraints_meta() {
+            let term = &challenges.transition_coeffs[m.constraint_idx]
+                * &transition_ood_frame_evaluations[m.constraint_idx];
+            match grouped_numerator_sums
+                .iter_mut()
+                .find(|(exemptions, _)| *exemptions == m.end_exemptions)
+            {
+                Some((_, acc)) => *acc += term,
+                None => grouped_numerator_sums.push((m.end_exemptions, term)),
+            }
+        }
+
+        let transition_c_i_evaluations_sum = grouped_numerator_sums.into_iter().fold(
+            FieldElement::zero(),
+            |acc, (end_exemptions, numerator_sum)| {
+                let correction = crate::constraints::zerofier::end_exemptions_correction(
+                    end_exemptions,
+                    &challenges.z,
+                    &domain.trace_primitive_root,
+                    trace_length,
+                );
+                acc + correction * numerator_sum
+            },
+        ) * inv_zerofier_denominator;
 
         let composition_poly_ood_evaluation =
             &boundary_quotient_ood_evaluation + transition_c_i_evaluations_sum;
@@ -839,12 +863,6 @@ pub trait IsStarkVerifier<
             return None;
         }
 
-        let mut deep_poly_evaluations = Vec::with_capacity(num_queries);
-        let mut deep_poly_evaluations_sym = Vec::with_capacity(num_queries);
-
-        // Build the base-field LDE evaluations as concatenated slice (precomputed + main)
-        // without lifting to the extension field. The helper now subtracts directly via
-        // the F: IsSubFieldOf<E> Sub impl, so we avoid a per-query base->extension lift.
         let primitive_root = &Field::get_primitive_root_of_unity(domain.root_order as u64)
             .expect("verifier domain root_order is a valid power of two");
 
@@ -856,8 +874,48 @@ pub trait IsStarkVerifier<
             step_size,
         )?;
 
-        for (i, iota) in challenges.iotas.iter().enumerate() {
+        // Each query's DEEP terms divide by (evaluation_point − z·gᵏ) for the OOD
+        // rows and by (evaluation_point − z^parts) for the composition part. Those
+        // z·gᵏ row points and z^parts are query-independent, so build them once and
+        // collect every query point's denominators (primary + symmetric) into ONE
+        // vector for a single whole-proof batch inverse — instead of a per-query-pair
+        // inversion (2·num_queries field inversions). One field inversion per proof.
+        let ood_height = ood_full.height;
+        let mut z_row_points = Vec::with_capacity(ood_height);
+        let mut current_z = challenges.z.clone();
+        for _ in 0..ood_height {
+            z_row_points.push(current_z.clone());
+            current_z = primitive_root * &current_z;
+        }
+        // z^parts, already computed once in `compute_query_invariant_deep_terms`.
+        let z_pow_parts = &query_invariant_terms.z_pow;
+
+        // Layout per query point: `ood_height` trace denominators then 1 composition
+        // denominator; points ordered [q0 primary, q0 sym, q1 primary, q1 sym, …].
+        let stride = ood_height + 1;
+        let mut denominators = Vec::with_capacity(2 * num_queries * stride);
+        for iota in challenges.iotas.iter() {
+            for is_sym in [false, true] {
+                let evaluation_point =
+                    Self::query_challenge_to_evaluation_point(*iota, is_sym, domain);
+                for z_row_point in &z_row_points {
+                    denominators.push(&evaluation_point - z_row_point);
+                }
+                denominators.push(&evaluation_point - z_pow_parts);
+            }
+        }
+        // One inversion for the whole proof's DEEP denominators. Fails closed if any
+        // evaluation point lands on an OOD point (malformed proof) — same rejection
+        // the previous per-pair inversions gave.
+        FieldElement::inplace_batch_inverse(&mut denominators).ok()?;
+
+        let mut deep_poly_evaluations = Vec::with_capacity(num_queries);
+        let mut deep_poly_evaluations_sym = Vec::with_capacity(num_queries);
+
+        for (i, _iota) in challenges.iotas.iter().enumerate() {
             let opening = proof.deep_poly_opening(i);
+            let primary_base = 2 * i * stride;
+            let sym_base = primary_base + stride;
 
             // Base-field portion as two borrowed slices in commit order —
             // precomputed columns FIRST, then main trace columns. The callee
@@ -885,14 +943,8 @@ pub trait IsStarkVerifier<
                 .map(|a| a.evaluations_sym())
                 .unwrap_or(&[]);
 
-            let evaluation_point = Self::query_challenge_to_evaluation_point(*iota, false, domain);
-            let evaluation_point_sym =
-                Self::query_challenge_to_evaluation_point(*iota, true, domain);
             let (evaluation, evaluation_sym) =
                 Self::reconstruct_deep_composition_poly_evaluation_pair(
-                    &evaluation_point,
-                    &evaluation_point_sym,
-                    primitive_root,
                     challenges,
                     &query_invariant_terms,
                     next_row_cols,
@@ -905,6 +957,10 @@ pub trait IsStarkVerifier<
                     lde_main_sym,
                     lde_aux_sym,
                     opening.composition_poly().evaluations_sym(),
+                    &denominators[primary_base..primary_base + ood_height],
+                    &denominators[primary_base + ood_height],
+                    &denominators[sym_base..sym_base + ood_height],
+                    &denominators[sym_base + ood_height],
                 )?;
             deep_poly_evaluations.push(evaluation);
             deep_poly_evaluations_sym.push(evaluation_sym);
@@ -917,15 +973,13 @@ pub trait IsStarkVerifier<
     /// trace term `coeff*(base-ood)*denom` as `denom*(coeff*base - coeff*ood)`
     /// isolates `coeff*ood` (identical for both points, hoisted into
     /// `query_invariant_terms`) from `coeff*base` (per-point), so both points
-    /// share the OOD walk and a single batch-inverse for their denominators.
-    /// g·z pruning restricts next rows (`row_idx >= step_size`) to the
-    /// transition-window columns `next_row_cols` — all other next-row
-    /// coefficients are zero, so those terms vanish from both sums.
+    /// share the OOD walk. Denominators are pre-inverted by the caller in one
+    /// whole-proof batch inverse and passed in as slices. g·z pruning restricts
+    /// next rows (`row_idx >= step_size`) to the transition-window columns
+    /// `next_row_cols` — all other next-row coefficients are zero, so those terms
+    /// vanish from both sums.
     #[allow(clippy::too_many_arguments)]
     fn reconstruct_deep_composition_poly_evaluation_pair<'b>(
-        evaluation_point: &FieldElement<Field>,
-        evaluation_point_sym: &FieldElement<Field>,
-        primitive_root: &FieldElement<Field>,
         challenges: &Challenges<FieldExtension>,
         query_invariant_terms: &QueryInvariantDeepTerms<FieldExtension>,
         next_row_cols: &[usize],
@@ -938,6 +992,13 @@ pub trait IsStarkVerifier<
         lde_trace_main_evaluations_sym: &'b [FieldElement<Field>],
         lde_trace_aux_evaluations_sym: &[FieldElement<FieldExtension>],
         lde_composition_poly_parts_evaluation_sym: &[FieldElement<FieldExtension>],
+        // Pre-inverted 1/(evaluation_point − z·gᵏ) per OOD row and
+        // 1/(evaluation_point − z^parts) for the composition part, primary then
+        // symmetric — all from the caller's single whole-proof batch inverse.
+        inv_denoms_trace: &[FieldElement<FieldExtension>],
+        inv_denom_composition: &FieldElement<FieldExtension>,
+        inv_denoms_trace_sym: &[FieldElement<FieldExtension>],
+        inv_denom_composition_sym: &FieldElement<FieldExtension>,
     ) -> Option<(FieldElement<FieldExtension>, FieldElement<FieldExtension>)> {
         let ood_evaluations_table_height = query_invariant_terms.ood_row_sum.len();
         let ood_evaluations_table_width = query_invariant_terms.ood_width;
@@ -977,23 +1038,12 @@ pub trait IsStarkVerifier<
         {
             return None;
         }
-
-        // Build both denominator sets (regular, then symmetric) and invert
-        // them together in a single batch.
-        let mut denoms = Vec::with_capacity(2 * ood_evaluations_table_height);
-        let mut current_z = challenges.z.clone();
-        for _ in 0..ood_evaluations_table_height {
-            denoms.push(evaluation_point - &current_z);
-            current_z = primitive_root * &current_z;
-        }
-        let mut current_z = challenges.z.clone();
-        for _ in 0..ood_evaluations_table_height {
-            denoms.push(evaluation_point_sym - &current_z);
-            current_z = primitive_root * &current_z;
+        // The caller passes one pre-inverted trace denominator per OOD row.
+        if inv_denoms_trace.len() != ood_evaluations_table_height
+            || inv_denoms_trace_sym.len() != ood_evaluations_table_height
+        {
+            return None;
         }
-        // A malformed proof can land an OOD evaluation point on the LDE coset, reject.
-        FieldElement::inplace_batch_inverse(&mut denoms).ok()?;
-        let (denoms_trace, denoms_trace_sym) = denoms.split_at(ood_evaluations_table_height);
 
         let mut trace_term = FieldElement::<FieldExtension>::zero();
         let mut trace_term_sym = FieldElement::<FieldExtension>::zero();
@@ -1032,8 +1082,8 @@ pub trait IsStarkVerifier<
                     }
                 }
             }
-            trace_term += &denoms_trace[row_idx] * &(&base_row_sum - ood_row_sum);
-            trace_term_sym += &denoms_trace_sym[row_idx] * &(&base_row_sum_sym - ood_row_sum);
+            trace_term += &inv_denoms_trace[row_idx] * &(&base_row_sum - ood_row_sum);
+            trace_term_sym += &inv_denoms_trace_sym[row_idx] * &(&base_row_sum_sym - ood_row_sum);
         }
 
         let number_of_parts = query_invariant_terms.number_of_parts;
@@ -1044,12 +1094,6 @@ pub trait IsStarkVerifier<
         {
             return None;
         }
-        let z_pow = &query_invariant_terms.z_pow;
-
-        // A malformed proof can make evaluation_point == z_pow, reject.
-        let mut denom_composition_pair = [evaluation_point - z_pow, evaluation_point_sym - z_pow];
-        FieldElement::inplace_batch_inverse(&mut denom_composition_pair).ok()?;
-        let [denom_composition, denom_composition_sym] = denom_composition_pair;
 
         let mut h_sum = FieldElement::<FieldExtension>::zero();
         let mut h_sum_sym = FieldElement::<FieldExtension>::zero();
@@ -1060,8 +1104,10 @@ pub trait IsStarkVerifier<
             h_sum += h_i_upsilon * gamma;
             h_sum_sym += h_i_upsilon_sym * gamma;
         }
-        let h_terms = (&h_sum - &query_invariant_terms.h_sum_zpow) * denom_composition;
-        let h_terms_sym = (&h_sum_sym - &query_invariant_terms.h_sum_zpow) * denom_composition_sym;
+        let h_diff = &h_sum - &query_invariant_terms.h_sum_zpow;
+        let h_diff_sym = &h_sum_sym - &query_invariant_terms.h_sum_zpow;
+        let h_terms = &h_diff * inv_denom_composition;
+        let h_terms_sym = &h_diff_sym * inv_denom_composition_sym;
 
         Some((trace_term + h_terms, trace_term_sym + h_terms_sym))
     }