FFT operation wrappers and evaluation with coset (#93)

* Add wrappers for fft related operations

* Add function to scale polynomials

* Add property tests for new fft with cosets

* Fix error message for invalid polynomial order

* Fix test (cooley-tukey implementation did not always work)

* Move tests to fft operations file

* Add `blowup_factor` to increase the domain evaluations on ftt.

* Fix test and parameter type for interpolation

* Add newline in gitignore

---------

Co-authored-by: Javier Chatruc <[email protected]>
Co-authored-by: Martin Paulucci <[email protected]>

Author: 3 people

.gitignore

@@ -8,3 +8,5 @@ Cargo.lock
 
 # These are backup files generated by rustfmt
 **/*.rs.bk
+
+**/proptest-regressions

math/src/fft/fft_cooley_tukey.rs

@@ -5,6 +5,16 @@ use crate::field::{
 
 use super::{errors::FFTError, helpers::log2};
 
+pub fn fft_with_blowup<F: IsField + IsTwoAdicField>(
+    coeffs: &[FieldElement<F>],
+    blowup_factor: usize,
+) -> Result<Vec<FieldElement<F>>, FFTError> {
+    let domain_size = coeffs.len() * blowup_factor;
+    let mut padded_coeffs = coeffs.to_vec();
+    padded_coeffs.resize(domain_size, FieldElement::zero());
+    fft(&padded_coeffs[..])
+}
+
 pub fn fft<F: IsField + IsTwoAdicField>(
     coeffs: &[FieldElement<F>],
 ) -> Result<Vec<FieldElement<F>>, FFTError> {
@@ -31,15 +41,16 @@ fn cooley_tukey<F: IsField>(
     let coeffs_odd: Vec<FieldElement<F>> = coeffs.iter().skip(1).step_by(2).cloned().collect();
 
     let (y_even, y_odd) = (
-        cooley_tukey(&coeffs_even, omega),
-        cooley_tukey(&coeffs_odd, omega),
+        cooley_tukey(&coeffs_even, &omega.pow(2_usize)),
+        cooley_tukey(&coeffs_odd, &omega.pow(2_usize)),
     );
+
     let mut y = vec![FieldElement::zero(); n];
-    for i in 0..n / 2 {
-        let a = &y_even[i];
-        let b = &(omega.pow(i) * &y_odd[i]);
+
+    for i in 0..n {
+        let a = &y_even[i % (n / 2)];
+        let b = &(omega.pow(i as u64) * &y_odd[i % (n / 2)]);
         y[i] = a + b;
-        y[i + n / 2] = a - b;
     }
     y
 }
@@ -56,81 +67,3 @@ pub fn inverse_cooley_tukey<F: IsField>(
         .map(|coeff| coeff * &inverse_n)
         .collect()
 }
-
-#[cfg(test)]
-mod fft_test {
-    use crate::field::test_fields::u64_test_field::U64TestField;
-    use crate::polynomial::Polynomial;
-    use proptest::prelude::*;
-
-    use super::*;
-
-    const MODULUS: u64 = 0xFFFFFFFF00000001;
-    type F = U64TestField<MODULUS>;
-    type FE = FieldElement<F>;
-
-    prop_compose! {
-        fn powers_of_two(max_exp: u8)(exp in 1..max_exp) -> usize { 1 << exp }
-        // max_exp cannot be multiple of the bits that represent a usize, generally 64 or 32.
-        // also it can't exceed the test field's two-adicity.
-    }
-    prop_compose! {
-        fn field_element()(num in any::<u64>()) -> FE { FE::from(num) }
-    }
-    prop_compose! {
-        fn field_vec(max_exp: u8)(elem in field_element(), size in powers_of_two(max_exp)) -> Vec<FE> {
-            vec![elem; size]
-        }
-    }
-    prop_compose! {
-        // non-power-of-two sized vector
-        fn unsuitable_field_vec(max_exp: u8)(elem in field_element(), size in powers_of_two(max_exp)) -> Vec<FE> {
-            vec![elem; size + 1]
-        }
-    }
-
-    proptest! {
-        // Property-based test that ensures FFT gives same result as a naive polynomial evaluation.
-        #[test]
-        fn test_fft_matches_naive_evaluation(coeffs in field_vec(8)) {
-            let poly = Polynomial::new(&coeffs[..]);
-            let omega = F::get_root_of_unity(log2(poly.coefficients().len()).unwrap()).unwrap();
-
-            let result = fft(poly.coefficients()).unwrap();
-
-            let twiddles_iter = (0..poly.coefficients().len() as u64).map(|i| omega.pow(i));
-            let expected: Vec<FE> = twiddles_iter.map(|x| poly.evaluate(&x)).collect();
-
-            prop_assert_eq!(result, expected);
-        }
-    }
-    proptest! {
-        // Property-based test that ensures IFFT is the inverse operation of FFT.
-        #[test]
-        fn test_ifft_composed_fft_is_identity(coeffs in field_vec(8)) {
-            let result = fft(&coeffs).unwrap();
-            let recovered_poly = inverse_fft(&result).unwrap();
-
-            prop_assert_eq!(recovered_poly, coeffs);
-        }
-    }
-
-    proptest! {
-        // Property-based test that ensures FFT won't work with a non-power-of-two polynomial.
-        #[test]
-        fn test_fft_panics_on_non_power_of_two(coeffs in unsuitable_field_vec(8)) {
-            let result = fft(&coeffs);
-
-            prop_assert!(matches!(result, Err(FFTError::InvalidOrder(_))));
-        }
-    }
-    proptest! {
-        // Property-based test that ensures FFT won't work with a degree 0 polynomial.
-        #[test]
-        fn test_fft_constant_poly(elem in field_element()) {
-            let result = fft(&[elem]);
-
-            prop_assert!(matches!(result, Err(FFTError::RootOfUnityError(_, k)) if k == 0));
-        }
-    }
-}

math/src/fft/helpers.rs

@@ -3,7 +3,7 @@ use super::errors::FFTError;
 pub fn log2(n: usize) -> Result<u64, FFTError> {
     if !n.is_power_of_two() {
         return Err(FFTError::InvalidOrder(
-            "The order of polynomial should a be power of 2".to_string(),
+            "The order of polynomial + 1 should a be power of 2".to_string(),
         ));
     }
     Ok(n.trailing_zeros() as u64)

math/src/fft/mod.rs

@@ -1,3 +1,4 @@
 pub mod errors;
 pub mod fft_cooley_tukey;
 mod helpers;
+pub mod operations;

math/src/fft/operations.rs

@@ -0,0 +1,129 @@
+use crate::{
+    field::{
+        element::FieldElement,
+        traits::{IsField, IsTwoAdicField},
+    },
+    polynomial::Polynomial,
+};
+
+use super::{
+    errors::FFTError,
+    fft_cooley_tukey::{fft, fft_with_blowup, inverse_fft},
+};
+
+pub fn evaluate_poly<F: IsField + IsTwoAdicField>(
+    polynomial: &Polynomial<FieldElement<F>>,
+) -> Result<Vec<FieldElement<F>>, FFTError> {
+    fft(polynomial.coefficients())
+}
+
+pub fn interpolate_poly<F: IsField + IsTwoAdicField>(
+    evaluation_points: &[FieldElement<F>],
+) -> Result<Vec<FieldElement<F>>, FFTError> {
+    inverse_fft(evaluation_points)
+}
+
+pub fn evaluate_poly_with_offset<F: IsField + IsTwoAdicField>(
+    polynomial: &Polynomial<FieldElement<F>>,
+    offset: &FieldElement<F>,
+    blowup_factor: usize,
+) -> Result<Vec<FieldElement<F>>, FFTError> {
+    let scaled_polynomial = polynomial.scale(offset);
+    fft_with_blowup(scaled_polynomial.coefficients(), blowup_factor)
+}
+
+#[cfg(test)]
+mod fft_test {
+    use crate::fft::helpers::log2;
+    use crate::field::test_fields::u64_test_field::U64TestField;
+    use crate::polynomial::Polynomial;
+    use proptest::prelude::*;
+
+    use super::*;
+
+    const MODULUS: u64 = 0xFFFFFFFF00000001;
+    type F = U64TestField<MODULUS>;
+    type FE = FieldElement<F>;
+
+    prop_compose! {
+        fn powers_of_two(max_exp: u8)(exp in 1..max_exp) -> usize { 1 << exp }
+        // max_exp cannot be multiple of the bits that represent a usize, generally 64 or 32.
+        // also it can't exceed the test field's two-adicity.
+    }
+    prop_compose! {
+        fn field_element()(num in any::<u64>().prop_filter("Avoid null polynomial", |x| x != &0)) -> FE {
+            FE::from(num)
+        }
+    }
+    prop_compose! {
+        fn offset()(num in 1..MODULUS - 1) -> FE { FE::from(num) }
+    }
+    prop_compose! {
+        fn field_vec(max_exp: u8)(elem in field_element(), size in powers_of_two(max_exp)) -> Vec<FE> {
+            vec![elem; size]
+        }
+    }
+    prop_compose! {
+        // non-power-of-two sized vector
+        fn unsuitable_field_vec(max_exp: u8)(elem in field_element(), size in powers_of_two(max_exp)) -> Vec<FE> {
+            vec![elem; size + 1]
+        }
+    }
+
+    proptest! {
+        // Property-based test that ensures FFT gives same result as a naive polynomial evaluation.
+        #[test]
+        fn test_fft_matches_naive_evaluation(coeffs in field_vec(8)) {
+            let poly = Polynomial::new(&coeffs[..]);
+            let result = evaluate_poly(&poly).unwrap();
+
+            let omega = F::get_root_of_unity(log2(poly.coefficients().len()).unwrap()).unwrap();
+            let twiddles_iter = (0..poly.coefficients().len() as u64).map(|i| omega.pow(i));
+            let expected: Vec<FE> = twiddles_iter.map(|x| poly.evaluate(&x)).collect();
+
+            prop_assert_eq!(result, expected);
+        }
+    }
+
+    proptest! {
+        // Property-based test that ensures FFT with coset gives same result as a naive polynomial evaluation.
+        #[test]
+        fn test_fft_with_coset_matches_naive_evaluation(coeffs in field_vec(8), blowup_factor in powers_of_two(3)) {
+            let domain_size = coeffs.len() * blowup_factor;
+
+            let poly = Polynomial::new(&coeffs[..]);
+
+            let offset = FE::new(2);
+            let result = evaluate_poly_with_offset(&poly, &offset, blowup_factor).unwrap();
+
+            let omega = F::get_root_of_unity(log2(domain_size).unwrap()).unwrap();
+            let twiddles_iter = (0..domain_size as u64).map(|i| omega.pow(i) * &offset);
+            let expected: Vec<FE> = twiddles_iter.map(|x| poly.evaluate(&(x))).collect();
+
+            prop_assert_eq!(result, expected);
+        }
+    }
+
+    proptest! {
+        // Property-based test that ensures IFFT is the inverse operation of FFT.
+        #[test]
+        fn test_ifft_composed_fft_is_identity(coeffs in field_vec(8)) {
+            let poly_to_evaluate = Polynomial::new(&coeffs[..]);
+            let evaluation_result = evaluate_poly(&poly_to_evaluate).unwrap();
+
+            let recovered_coefficients = interpolate_poly(&evaluation_result).unwrap();
+            prop_assert_eq!(recovered_coefficients, coeffs);
+        }
+    }
+
+    proptest! {
+        // Property-based test that ensures FFT won't work with a degree 0 polynomial.
+        #[test]
+        fn test_fft_constant_poly(elem in field_element()) {
+            let poly_to_evaluate = Polynomial::new(&[elem]);
+            let result = evaluate_poly(&poly_to_evaluate);
+
+            prop_assert!(matches!(result, Err(FFTError::RootOfUnityError(_, k)) if k == 0));
+        }
+    }
+}

math/src/polynomial.rs

@@ -146,6 +146,18 @@ impl<F: IsField> Polynomial<FieldElement<F>> {
             Polynomial::new(&coefficients)
         }
     }
+
+    pub fn scale(&self, factor: &FieldElement<F>) -> Self {
+        let scaled_coefficients = self
+            .coefficients
+            .iter()
+            .enumerate()
+            .map(|(i, coeff)| factor.pow(i) * coeff)
+            .collect();
+        Self {
+            coefficients: scaled_coefficients,
+        }
+    }
 }
 
 impl<F: IsField> ops::Add<&Polynomial<FieldElement<F>>> for &Polynomial<FieldElement<F>> {