Feat/fast division

crates/math/benches/polynomials/polynomial.rs

@@ -59,6 +59,17 @@ pub fn polynomial_benchmarks(c: &mut Criterion) {
         bench.iter(|| black_box(&x_poly) * black_box(&y_poly));
     });
 
+    let y_poly = rand_complex_mersenne_poly(big_order - 2);
+
+    group.bench_function("fast div big poly", |bench| {
+        bench
+            .iter(|| black_box(&x_poly).fast_division::<Degree2ExtensionField>(black_box(&y_poly)));
+    });
+
+    group.bench_function("slow div big poly", |bench| {
+        bench.iter(|| black_box(x_poly.clone()).long_division_with_remainder(black_box(&y_poly)));
+    });
+
     group.bench_function("div", |bench| {
         let x_poly = rand_poly(order);
         let y_poly = rand_poly(order);

crates/math/src/fft/polynomial.rs

@@ -1,5 +1,6 @@
 use crate::fft::errors::FFTError;
 
+use crate::field::errors::FieldError;
 use crate::field::traits::{IsField, IsSubFieldOf};
 use crate::{
     field::{
@@ -115,6 +116,67 @@ impl<E: IsField> Polynomial<FieldElement<E>> {
 
         Polynomial::interpolate_fft::<F>(&r)
     }
+
+    /// Divides two polynomials with remainder.
+    /// This is faster than the naive division if the degree of the divisor
+    /// is greater than the degree of the dividend and both degrees are large enough.
+    pub fn fast_division<F: IsSubFieldOf<E> + IsFFTField>(
+        &self,
+        divisor: &Self,
+    ) -> Result<(Self, Self), FFTError> {
+        let n = self.degree();
+        let m = divisor.degree();
+        if divisor.coefficients.is_empty()
+            || divisor
+                .coefficients
+                .iter()
+                .all(|c| c == &FieldElement::zero())
+        {
+            return Err(FieldError::DivisionByZero.into());
+        }
+        if n < m {
+            return Ok((Self::zero(), self.clone()));
+        }
+        let d = n - m; // Degree of the quotient
+        let a_rev = self.reverse(n);
+        let b_rev = divisor.reverse(m);
+        let inv_b_rev = Self::invert_polynomial_mod::<F>(&b_rev, d + 1)?;
+        let q = a_rev
+            .fast_fft_multiplication::<F>(&inv_b_rev)?
+            .truncate(d + 1)
+            .reverse(d);
+
+        let r = self - q.fast_fft_multiplication::<F>(divisor)?;
+        Ok((q, r))
+    }
+
+    /// Computes the inverse of polynomial P modulo x^k using Newton iteration.
+    /// P must have an invertible constant term.
+    pub fn invert_polynomial_mod<F: IsSubFieldOf<E> + IsFFTField>(
+        p: &Self,
+        k: usize,
+    ) -> Result<Self, FFTError> {
+        if p.coefficients.is_empty() || p.coefficients.iter().all(|c| c == &FieldElement::zero()) {
+            return Err(FieldError::DivisionByZero.into());
+        }
+        let mut q = Self::new(&[p.coefficients[0].inv()?]);
+        let mut current_precision = 1;
+
+        let two = Self::new(&[FieldElement::<F>::one() + FieldElement::one()]);
+        while current_precision < k {
+            current_precision *= 2;
+            let temp = p
+                .fast_fft_multiplication::<F>(&q)?
+                .truncate(current_precision);
+            let correction = &two - temp;
+            q = q
+                .fast_fft_multiplication::<F>(&correction)?
+                .truncate(current_precision);
+        }
+
+        // Final truncation to desired degree k
+        Ok(q.truncate(k))
+    }
 }
 
 pub fn compose_fft<F, E>(
@@ -273,6 +335,11 @@ mod tests {
                 vec
             }
         }
+        prop_compose! {
+            fn non_empty_field_vec(max_exp: u8)(vec in collection::vec(field_element(), 1 << max_exp)) -> Vec<FE> {
+                vec
+            }
+        }
         prop_compose! {
             fn non_power_of_two_sized_field_vec(max_exp: u8)(vec in collection::vec(field_element(), 2..1<<max_exp).prop_filter("Avoid polynomials of size power of two", |vec| !vec.len().is_power_of_two())) -> Vec<FE> {
                 vec
@@ -283,6 +350,11 @@ mod tests {
                 Polynomial::new(&coeffs)
             }
         }
+        prop_compose! {
+            fn non_zero_poly(max_exp: u8)(coeffs in non_empty_field_vec(max_exp)) -> Polynomial<FE> {
+                Polynomial::new(&coeffs)
+            }
+        }
         prop_compose! {
             fn poly_with_non_power_of_two_coeffs(max_exp: u8)(coeffs in non_power_of_two_sized_field_vec(max_exp)) -> Polynomial<FE> {
                 Polynomial::new(&coeffs)
@@ -336,6 +408,11 @@ mod tests {
             fn test_fft_multiplication_works(poly in poly(7), other in poly(7)) {
                 prop_assert_eq!(poly.fast_fft_multiplication::<F>(&other).unwrap(), poly * other);
             }
+
+            #[test]
+            fn test_fft_division_works(poly in non_zero_poly(7), other in non_zero_poly(7)) {
+                prop_assert_eq!(poly.fast_division::<F>(&other).unwrap(), poly.long_division_with_remainder(&other));
+            }
         }
 
         #[test]
@@ -371,6 +448,11 @@ mod tests {
                 vec
             }
         }
+        prop_compose! {
+            fn non_empty_field_vec(max_exp: u8)(vec in collection::vec(field_element(), 1 << max_exp)) -> Vec<FE> {
+                vec
+            }
+        }
         prop_compose! {
             fn non_power_of_two_sized_field_vec(max_exp: u8)(vec in collection::vec(field_element(), 2..1<<max_exp).prop_filter("Avoid polynomials of size power of two", |vec| !vec.len().is_power_of_two())) -> Vec<FE> {
                 vec
@@ -381,6 +463,11 @@ mod tests {
                 Polynomial::new(&coeffs)
             }
         }
+        prop_compose! {
+            fn non_zero_poly(max_exp: u8)(coeffs in non_empty_field_vec(max_exp)) -> Polynomial<FE> {
+                Polynomial::new(&coeffs)
+            }
+        }
         prop_compose! {
             fn poly_with_non_power_of_two_coeffs(max_exp: u8)(coeffs in non_power_of_two_sized_field_vec(max_exp)) -> Polynomial<FE> {
                 Polynomial::new(&coeffs)
@@ -436,6 +523,11 @@ mod tests {
             fn test_fft_multiplication_works(poly in poly(7), other in poly(7)) {
                 prop_assert_eq!(poly.fast_fft_multiplication::<F>(&other).unwrap(), poly * other);
             }
+
+            #[test]
+            fn test_fft_division_works(poly in poly(7), other in non_zero_poly(7)) {
+                prop_assert_eq!(poly.fast_division::<F>(&other).unwrap(), poly.long_division_with_remainder(&other));
+            }
         }
     }
 

crates/math/src/polynomial/mod.rs

@@ -335,6 +335,20 @@ impl<F: IsField> Polynomial<FieldElement<F>> {
                 .collect(),
         }
     }
+
+    pub fn truncate(&self, k: usize) -> Self {
+        if k == 0 {
+            Self::zero()
+        } else {
+            Self::new(&self.coefficients[0..k.min(self.coefficients.len())])
+        }
+    }
+    pub fn reverse(&self, d: usize) -> Self {
+        let mut coeffs = self.coefficients.clone();
+        coeffs.resize(d + 1, FieldElement::zero());
+        coeffs.reverse();
+        Self::new(&coeffs)
+    }
 }
 
 impl<F: IsPrimeField> Polynomial<FieldElement<F>> {