From ccab5dc1ab33f9e1cfe6355ed874c9a96efa7110 Mon Sep 17 00:00:00 2001 From: diegokingston Date: Fri, 22 May 2026 19:12:29 -0300 Subject: [PATCH] refactor(stark): migrate end_exemptions_poly to evaluation form Replaces the coefficient-form `end_exemptions_poly: Polynomial` with two eval-form helpers, removing transition.rs's last dependency on `Polynomial` arithmetic. Stacks on #604 (which removed boundary.rs's `Polynomial` dependency). With both landed, no production code in the prover uses `Polynomial` operators. - `end_exemptions_roots` returns the roots r_i of prod(x - r_i) (<= 2 in the example AIRs, 0 in every VM table). - `end_exemptions_lde_evaluations(domain)` evaluates the product over the precomputed LDE coset directly: an O(N * end_exemptions) loop replaces an O(N log N) FFT. - `evaluate_zerofier`'s OOD path computes prod(z - r_i) directly instead of `Polynomial::evaluate`. Performance: the VM's dominant path (end_exemptions == 0) is unchanged - the existing else-branch early-return is preserved bit-for-bit, so that case still returns the short cyclic vector instead of expanding it. The k > 0 path (example AIRs only) goes from FFT to direct product, strictly faster. Correctness verified by `cargo test -p stark` (125/125, includes the fibonacci and read-only-memory AIRs with end_exemptions = 1 and 2 - exactly the eval-form k > 0 path). VM prover lib test counts are identical to the #604 baseline (273 pass, 77 pre-existing failures unrelated to constraints). --- crypto/stark/src/constraints/transition.rs | 97 +++++++++++++--------- 1 file changed, 58 insertions(+), 39 deletions(-) diff --git a/crypto/stark/src/constraints/transition.rs b/crypto/stark/src/constraints/transition.rs index af86ed7e..cb2abb60 100644 --- a/crypto/stark/src/constraints/transition.rs +++ b/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, trace_length: usize, - ) -> Polynomial> { - let one_poly = Polynomial::new_monomial(FieldElement::::one(), 0); - if self.end_exemptions() == 0 { - return one_poly; + ) -> Vec> { + let end_exemptions = self.end_exemptions(); + if end_exemptions == 0 { + return Vec::new(); } - let period = self.period(); - let decrement = trace_primitive_root.pow(trace_length - period); + // FIXME: when offset != 0 the roots may need to be scaled by + // trace_root^(offset * trace_length / period) — carried over unresolved + // from the original coefficient-form end_exemptions_poly. + let decrement = trace_primitive_root.pow(trace_length - self.period()); + let mut roots = Vec::with_capacity(end_exemptions); 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::::one(), 1) - current.clone()); + for _ in 0..end_exemptions { + roots.push(current.clone()); current = ¤t * &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) -> Vec> { + 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::::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() { @@ -192,13 +215,7 @@ where // 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 +244,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 +273,14 @@ where trace_primitive_root: &FieldElement, trace_length: usize, ) -> FieldElement { - 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::::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 +299,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 } }