diff --git a/crypto/math/src/fft/polynomial.rs b/crypto/math/src/fft/polynomial.rs index 9157903fd..ccc1ac391 100644 --- a/crypto/math/src/fft/polynomial.rs +++ b/crypto/math/src/fft/polynomial.rs @@ -80,6 +80,33 @@ impl Polynomial> { evaluate_fft_cpu::(&coeffs) } + /// Same as `evaluate_fft` but returns the evaluations in bit-reversed order, + /// skipping the final natural-order permutation. Use when the consumer expects + /// bit-reversed input (e.g. FRI commit phase, which pairs consecutive values as + /// {f(x), f(-x)}). + pub fn evaluate_fft_bit_reversed>( + poly: &Polynomial>, + blowup_factor: usize, + domain_size: Option, + ) -> Result>, FFTError> + where + E: Send + Sync, + { + let domain_size = domain_size.unwrap_or(0); + let len = core::cmp::max(poly.coeff_len(), domain_size).next_power_of_two() * blowup_factor; + if len.trailing_zeros() as u64 > F::TWO_ADICITY { + return Err(FFTError::DomainSizeError(len.trailing_zeros() as usize)); + } + if poly.coefficients().is_empty() { + return Ok(vec![FieldElement::zero(); len]); + } + + let mut coeffs = poly.coefficients().to_vec(); + coeffs.resize(len, FieldElement::zero()); + + evaluate_fft_cpu_raw::(&coeffs, false) + } + /// Returns `N` evaluations with an offset of this polynomial using FFT over a domain in a subfield F of E /// (so the results are P(w^i), with w being a primitive root of unity). /// `N = max(self.coeff_len(), domain_size).next_power_of_two() * blowup_factor`. @@ -279,6 +306,17 @@ where } pub fn evaluate_fft_cpu(coeffs: &[FieldElement]) -> Result>, FFTError> +where + F: IsFFTField + IsSubFieldOf, + E: IsField + Send + Sync, +{ + evaluate_fft_cpu_raw::(coeffs, true) +} + +fn evaluate_fft_cpu_raw( + coeffs: &[FieldElement], + permute_to_natural: bool, +) -> Result>, FFTError> where F: IsFFTField + IsSubFieldOf, E: IsField + Send + Sync, @@ -293,7 +331,9 @@ where let mut result = coeffs.to_vec(); dispatch_fft(&mut result, &layer_twiddles)?; - in_place_bit_reverse_permute(&mut result); + if permute_to_natural { + in_place_bit_reverse_permute(&mut result); + } Ok(result) } diff --git a/crypto/math/src/tests/fft_tests.rs b/crypto/math/src/tests/fft_tests.rs index 989fbf70b..767507ad0 100644 --- a/crypto/math/src/tests/fft_tests.rs +++ b/crypto/math/src/tests/fft_tests.rs @@ -48,6 +48,7 @@ mod fft_helpers_test { mod fft_polynomial_tests { use crate::field::traits::IsField; + use crate::fft::cpu::bit_reversing::in_place_bit_reverse_permute; use crate::fft::cpu::roots_of_unity::{ get_powers_of_primitive_root, get_powers_of_primitive_root_coset, }; @@ -227,6 +228,17 @@ mod fft_polynomial_tests { prop_assert_eq!(poly, new_poly); } + // Property-based test that ensures evaluate_fft_bit_reversed returns the same + // values as evaluate_fft followed by an in-place bit-reverse permutation, + // across varying blowup factors. + #[test] + fn test_fft_bit_reversed_matches_evaluate_fft_then_permute(poly in poly(6), blowup_factor in powers_of_two(4)) { + let mut expected = Polynomial::evaluate_fft::(&poly, blowup_factor, None).unwrap(); + in_place_bit_reverse_permute(&mut expected); + let got = Polynomial::evaluate_fft_bit_reversed::(&poly, blowup_factor, None).unwrap(); + prop_assert_eq!(got, expected); + } + #[test] fn test_fft_multiplication_works(poly in poly(7), other in poly(7)) { prop_assert_eq!(poly.fast_fft_multiplication::(&other).unwrap(), poly * other); @@ -253,6 +265,24 @@ mod fft_polynomial_tests { Polynomial::new(&[FE::new(0), FE::new(0), FE::new(0), FE::new(2)]) ); } + + #[test] + fn fft_bit_reversed_handles_domain_size_greater_than_coeff_len() { + let poly = Polynomial::new(&[FE::new(1), FE::new(2), FE::new(3)]); + let domain_size = 32; + let mut expected = Polynomial::evaluate_fft::(&poly, 1, Some(domain_size)).unwrap(); + in_place_bit_reverse_permute(&mut expected); + let got = + Polynomial::evaluate_fft_bit_reversed::(&poly, 1, Some(domain_size)).unwrap(); + assert_eq!(got, expected); + } + + #[test] + fn fft_bit_reversed_returns_zeros_for_empty_polynomial() { + let poly: Polynomial = Polynomial::new(&[]); + let got = Polynomial::evaluate_fft_bit_reversed::(&poly, 1, Some(8)).unwrap(); + assert_eq!(got, vec![FE::zero(); 8]); + } } #[test] diff --git a/crypto/stark/src/prover.rs b/crypto/stark/src/prover.rs index 43086d4fa..9924f3efb 100644 --- a/crypto/stark/src/prover.rs +++ b/crypto/stark/src/prover.rs @@ -4,7 +4,7 @@ use std::sync::Arc; use std::time::Instant; use crypto::fiat_shamir::is_transcript::IsStarkTranscript; -use math::fft::cpu::bit_reversing::{in_place_bit_reverse_permute, reverse_index}; +use math::fft::cpu::bit_reversing::reverse_index; use math::fft::cpu::bowers_fft::LayerTwiddles; use math::fft::errors::FFTError; @@ -1109,9 +1109,12 @@ pub trait IsStarkProver< let t_sub = Instant::now(); let deep_poly = Polynomial::interpolate_fft::(&deep_evals).expect("iFFT should succeed"); - let mut lde_evals = Polynomial::evaluate_fft::(&deep_poly, 1, Some(domain_size)) - .expect("FFT should succeed"); - in_place_bit_reverse_permute(&mut lde_evals); + // FRI commit_phase consumes bit-reversed evaluations natively. Request them + // directly from evaluate_fft_bit_reversed to avoid a pair of redundant permutes + // (evaluate_fft's internal natural-order permute + an external re-bit-reverse). + let lde_evals = + Polynomial::evaluate_fft_bit_reversed::(&deep_poly, 1, Some(domain_size)) + .expect("FFT should succeed"); #[cfg(feature = "instruments")] let r4_fft_dur = t_sub.elapsed();