perf(stark): migrate end_exemptions_poly to evaluation form

crypto/stark/src/constraints/transition.rs

@@ -1,11 +1,9 @@
 use core::ops::Div;
 
 use crate::domain::Domain;
-use crate::prover::evaluate_polynomial_on_lde_domain;
 use crate::traits::TransitionEvaluationContext;
 use math::field::element::FieldElement;
 use math::field::traits::{IsFFTField, IsField, IsSubFieldOf};
-use math::polynomial::Polynomial;
 
 /// TransitionConstraintEvaluator represents the behaviour that a transition constraint
 /// over the computation that wants to be proven must comply with.
@@ -105,29 +103,56 @@ where
         self.evaluate_verifier(evaluation_context, ext_evals);
     }
 
-    /// Method for calculating the end exemptions polynomial.
+    /// Roots of the end-exemptions polynomial `∏(x - rᵢ)`.
     ///
-    /// This polynomial is used to compute zerofiers of the constraint, and the default
-    /// implementation should normally not be changed.
-    fn end_exemptions_poly(
+    /// The end-exemptions polynomial vanishes on the last `end_exemptions()`
+    /// rows the constraint must skip. This returns its roots `rᵢ` so callers can
+    /// evaluate the product `∏(x - rᵢ)` directly at the points they need — the
+    /// eval-form replacement for the former coefficient-form `end_exemptions_poly`.
+    /// The default implementation should normally not be changed.
+    fn end_exemptions_roots(
         &self,
         trace_primitive_root: &FieldElement<F>,
         trace_length: usize,
-    ) -> Polynomial<FieldElement<F>> {
-        let one_poly = Polynomial::new_monomial(FieldElement::<F>::one(), 0);
-        if self.end_exemptions() == 0 {
-            return one_poly;
+    ) -> Vec<FieldElement<F>> {
+        let end_exemptions = self.end_exemptions();
+        if end_exemptions == 0 {
+            return Vec::new();
         }
+        // Last row in the constraint's evaluation domain is g^(offset + N - period);
+        // walking backward by g^period gives the remaining end-exemption roots.
         let period = self.period();
         let decrement = trace_primitive_root.pow(trace_length - period);
-        let mut current = decrement.clone();
-        // FIXME: CHECK IF WE NEED TO CHANGE THE NEW MONOMIAL'S ARGUMENTS TO trace_root^(offset * trace_length / period) INSTEAD OF ONE!!!!
-        (0..self.end_exemptions()).fold(one_poly, |acc, _| {
-            let next =
-                acc * (Polynomial::new_monomial(FieldElement::<F>::one(), 1) - current.clone());
+        let mut current = trace_primitive_root.pow(self.offset() + trace_length - period);
+        let mut roots = Vec::with_capacity(end_exemptions);
+        for _ in 0..end_exemptions {
+            roots.push(current.clone());
             current = &current * &decrement;
-            next
-        })
+        }
+        roots
+    }
+
+    /// Evaluations of the end-exemptions polynomial `∏(x - rᵢ)` over the LDE
+    /// domain.
+    ///
+    /// Eval-form replacement for FFT-evaluating the coefficient-form polynomial:
+    /// the product has degree `end_exemptions()` (≤ 2 in practice), so the direct
+    /// `O(N · end_exemptions)` product over the precomputed LDE coset is cheaper
+    /// than an `O(N log N)` FFT. With no exemptions this yields all ones.
+    fn end_exemptions_lde_evaluations(&self, domain: &Domain<F>) -> Vec<FieldElement<F>> {
+        let roots = self.end_exemptions_roots(
+            &domain.trace_primitive_root,
+            domain.trace_roots_of_unity.len(),
+        );
+        domain
+            .lde_roots_of_unity_coset
+            .iter()
+            .map(|x| {
+                roots
+                    .iter()
+                    .fold(FieldElement::<F>::one(), |acc, r| acc * (x - r))
+            })
+            .collect()
     }
 
     /// Compute evaluations of the constraints zerofier over a LDE domain.
@@ -140,8 +165,6 @@ where
         let lde_root_order = u64::from((blowup_factor * trace_length).trailing_zeros());
         let lde_root = F::get_primitive_root_of_unity(lde_root_order).unwrap();
 
-        let end_exemptions_poly = self.end_exemptions_poly(trace_primitive_root, trace_length);
-
         // If there is an exemptions period defined for this constraint, the evaluations are calculated directly
         // by computing P_exemptions(x) / Zerofier(x)
         if let Some(exemptions_period) = self.exemptions_period() {
@@ -189,16 +212,16 @@ where
                 .map(|(num, denom_inv)| num * denom_inv)
                 .collect();
 
+            // Mirror the else-branch fast path: with no end exemptions the zerofier stays
+            // cyclic, so return the short period-length vector and let the consumer cycle.
+            if self.end_exemptions() == 0 {
+                return evaluations;
+            }
+
             // FIXME: Instead of computing this evaluations for each constraint, they can be computed
             // once for every constraint with the same end exemptions (combination of end_exemptions()
             // and period).
-            let end_exemption_evaluations = evaluate_polynomial_on_lde_domain(
-                &end_exemptions_poly,
-                blowup_factor,
-                domain.interpolation_domain_size,
-                coset_offset,
-            )
-            .unwrap();
+            let end_exemption_evaluations = self.end_exemptions_lde_evaluations(domain);
 
             let cycled_evaluations = evaluations
                 .iter()
@@ -227,19 +250,14 @@ where
 
             FieldElement::inplace_batch_inverse(&mut evaluations).unwrap();
 
-            // Fast path: when end_exemptions == 0, the end_exemptions_poly is constant 1,
-            // so the multiplication is identity. Skip the expensive FFT evaluation.
+            // Fast path: when end_exemptions == 0 there are no exemption roots, so
+            // the zerofier stays cyclic — return the short period-length vector
+            // directly instead of expanding it over the full LDE domain.
             if self.end_exemptions() == 0 {
                 return evaluations;
             }
 
-            let end_exemption_evaluations = evaluate_polynomial_on_lde_domain(
-                &end_exemptions_poly,
-                blowup_factor,
-                domain.interpolation_domain_size,
-                coset_offset,
-            )
-            .unwrap();
+            let end_exemption_evaluations = self.end_exemptions_lde_evaluations(domain);
 
             let cycled_evaluations = evaluations
                 .iter()
@@ -261,7 +279,14 @@ where
         trace_primitive_root: &FieldElement<F>,
         trace_length: usize,
     ) -> FieldElement<E> {
-        let end_exemptions_poly = self.end_exemptions_poly(trace_primitive_root, trace_length);
+        let end_exemptions_roots = self.end_exemptions_roots(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 = end_exemptions_roots
+            .iter()
+            .fold(FieldElement::<E>::one(), |acc, root| {
+                acc * -(root.clone() - z.clone())
+            });
 
         if let Some(exemptions_period) = self.exemptions_period() {
             debug_assert!(exemptions_period.is_multiple_of(self.period()));
@@ -280,14 +305,14 @@ where
             return numerator
                 .div(denominator)
                 .expect("zerofier denominator is non-zero: z is sampled out-of-domain")
-                * end_exemptions_poly.evaluate(z);
+                * &end_exemptions_eval;
         }
 
         (-trace_primitive_root.pow(self.offset() * trace_length / self.period())
             + z.pow(trace_length / self.period()))
         .inv()
         .unwrap()
-            * end_exemptions_poly.evaluate(z)
+            * &end_exemptions_eval
     }
 }
 

crypto/stark/src/tests/mod.rs

@@ -6,3 +6,4 @@ pub mod proof_options_tests;
 pub mod prove_verify_roundtrip_tests;
 pub mod prover_tests;
 pub mod small_trace_tests;
+pub mod transition_tests;

crypto/stark/src/tests/transition_tests.rs

@@ -0,0 +1,85 @@
+use crate::constraints::transition::TransitionConstraintEvaluator;
+use crate::traits::TransitionEvaluationContext;
+use math::field::element::FieldElement;
+use math::field::goldilocks::GoldilocksField;
+use math::field::traits::IsFFTField;
+use std::marker::PhantomData;
+
+/// Dummy evaluator that only exposes the trait knobs we need (`period`, `offset`,
+/// `end_exemptions`) to exercise `end_exemptions_roots`.
+struct DummyConstraint<F: IsFFTField + Send + Sync> {
+    period: usize,
+    offset: usize,
+    end_exemptions: usize,
+    phantom: PhantomData<F>,
+}
+
+impl<F: IsFFTField + Send + Sync> TransitionConstraintEvaluator<F, F> for DummyConstraint<F> {
+    fn degree(&self) -> usize {
+        1
+    }
+    fn constraint_idx(&self) -> usize {
+        0
+    }
+    fn period(&self) -> usize {
+        self.period
+    }
+    fn offset(&self) -> usize {
+        self.offset
+    }
+    fn end_exemptions(&self) -> usize {
+        self.end_exemptions
+    }
+    fn evaluate_verifier(&self, _: &TransitionEvaluationContext<F, F>, _: &mut [FieldElement<F>]) {}
+}
+
+#[test]
+fn end_exemptions_roots_default_offset_matches_last_rows() {
+    let trace_length = 8usize;
+    let g =
+        GoldilocksField::get_primitive_root_of_unity(trace_length.trailing_zeros() as u64).unwrap();
+    let c = DummyConstraint::<GoldilocksField> {
+        period: 1,
+        offset: 0,
+        end_exemptions: 2,
+        phantom: PhantomData,
+    };
+
+    let roots = c.end_exemptions_roots(&g, trace_length);
+
+    // Constraint applies on rows 0..8; last two rows are 6 and 7.
+    assert_eq!(roots, vec![g.pow(7u64), g.pow(6u64)]);
+}
+
+#[test]
+fn end_exemptions_roots_nonzero_offset_walks_the_offset_domain() {
+    let trace_length = 8usize;
+    let g =
+        GoldilocksField::get_primitive_root_of_unity(trace_length.trailing_zeros() as u64).unwrap();
+    let c = DummyConstraint::<GoldilocksField> {
+        period: 2,
+        offset: 1,
+        end_exemptions: 2,
+        phantom: PhantomData,
+    };
+
+    let roots = c.end_exemptions_roots(&g, trace_length);
+
+    // Constraint applies on rows {1, 3, 5, 7}; last two are 5 and 7.
+    assert_eq!(roots, vec![g.pow(7u64), g.pow(5u64)]);
+}
+
+#[test]
+fn end_exemptions_roots_zero_exemptions_is_empty() {
+    let trace_length = 8usize;
+    let g =
+        GoldilocksField::get_primitive_root_of_unity(trace_length.trailing_zeros() as u64).unwrap();
+    let c = DummyConstraint::<GoldilocksField> {
+        period: 1,
+        offset: 0,
+        end_exemptions: 0,
+        phantom: PhantomData,
+    };
+
+    assert!(c.end_exemptions_roots(&g, trace_length).is_empty());
+}