clean up polynomial arithmetic and degree

Author: GiacomoPope

src/polynomial_ring/poly.rs

@@ -10,19 +10,7 @@ use std::{
 };
 
 /// Trait for arithmetic for univariate polynomials in Fp[X]
-pub trait Poly:
-    Index<usize>
-    + IndexMut<usize>
-    + Sized
-    + Neg<Output = Self>
-    + Add<Output = Self>
-    + AddAssign
-    + Sub<Output = Self>
-    + SubAssign
-    + Mul<Output = Self>
-    + MulAssign
-    + Display
-{
+pub trait Poly: Index<usize> + IndexMut<usize> + Sized + Display {
     // TODO
 }
 
@@ -47,14 +35,28 @@ impl<Fp: FpTrait> Polynomial<Fp> {
         self.coeffs.len()
     }
 
-    /// The degree of the polynomial. TODO: should we trim trailing zeros? If so, how often?
-    /// We return None for the zero polynomial (instead of -inf which is a bit too big for my
-    /// computer...)
-    fn degree(&self) -> Option<usize> {
-        if self.coeffs.is_empty() {
+    /// Truncate leading zeros from the polynomial
+    fn truncate_leading_zeros(&mut self) {
+        // find the index of the first non-zero element
+        let mut i = self.len() - 1;
+        while i > 0 && self.coeffs[i].is_zero() == u32::MAX {
+            i -= 1;
+        }
+        self.coeffs.truncate(i);
+    }
+
+    /// Return the degree of the polynomial.
+    // TODO: should we trim trailing zeros and mutate while computing the degree or do the following
+    // and allow trailing zeros to remain.
+    pub fn degree(&self) -> Option<usize> {
+        let mut i = self.len() - 1;
+        while i > 0 && self.coeffs[i].is_zero() == u32::MAX {
+            i -= 1;
+        }
+        if i == 0 && self.coeffs[0].is_zero() == u32::MAX {
             None
         } else {
-            Some(self.len() - 1)
+            Some(i)
         }
     }
 
@@ -103,30 +105,17 @@ impl<Fp: FpTrait> Polynomial<Fp> {
         is_zero
     }
 
-    /// Set self to it's negative.
-    fn set_neg(&mut self) {
-        for x in self.coeffs.iter_mut() {
-            x.set_neg();
-        }
-    }
-
-    /// Set self <- self + other
-    // TODO: should we assume other has length smaller or equal to
-    // self, or dynamically extend?
-    fn set_add(&mut self, other: &Self) {
-        let k = self.len().min(other.len());
-        for i in 0..k {
-            self.coeffs[i] += other[i];
+    /// Compute f <-- f + g, assumes that the length of f and g are the same.
+    fn add_into(f: &mut [Fp], g: &[Fp]) {
+        for i in 0..g.len() {
+            f[i] += g[i];
         }
     }
 
-    /// Set self <- self - other
-    // TODO: should we assume other has length smaller or equal to
-    // self, or dynamically extend?
-    fn set_sub(&mut self, other: &Self) {
-        let k = self.len().min(other.len());
-        for i in 0..k {
-            self.coeffs[i] -= other[i];
+    /// Compute f <-- f - g, assumes that the length of f and g are the same.
+    fn sub_into(f: &mut [Fp], g: &[Fp]) {
+        for i in 0..g.len() {
+            f[i] -= g[i];
         }
     }
 
@@ -143,6 +132,29 @@ impl<Fp: FpTrait> Polynomial<Fp> {
         }
     }
 
+    /// Set self to it's negative.
+    fn set_neg(&mut self) {
+        for x in self.coeffs.iter_mut() {
+            x.set_neg();
+        }
+    }
+
+    /// Set self <- self + other
+    fn set_add(&mut self, other: &Self) {
+        if self.len() < other.len() {
+            self.coeffs.resize(other.len(), Fp::ZERO);
+        }
+        Self::add_into(&mut self.coeffs, &other.coeffs);
+    }
+
+    /// Set self <- self - other
+    fn set_sub(&mut self, other: &Self) {
+        if self.len() < other.len() {
+            self.coeffs.resize(other.len(), Fp::ZERO);
+        }
+        Self::sub_into(&mut self.coeffs, &other.coeffs);
+    }
+
     /// Set self <- self * other
     fn set_mul(&mut self, other: &Self) {
         let mut fg_coeffs = vec![Fp::ZERO; self.len() + other.len() - 1];
@@ -179,6 +191,7 @@ impl<Fp: FpTrait> Polynomial<Fp> {
     }
 
     /// Evaluate a polynomial at a value `a`
+    // TODO: is this over engineered?
     pub fn evaluate(&self, a: &Fp) -> Fp {
         // Handle degree 0 and 1 cases early.
         if self.len() == 0 {
@@ -251,29 +264,18 @@ impl<Fp: FpTrait> IndexMut<usize> for Polynomial<Fp> {
     }
 }
 
-impl<Fp: FpTrait> Neg for Polynomial<Fp> {
+impl<Fp: FpTrait> Neg for &Polynomial<Fp> {
     type Output = Polynomial<Fp>;
 
     #[inline(always)]
     fn neg(self) -> Polynomial<Fp> {
-        let mut r = self;
+        let mut r = self.clone();
         r.set_neg();
         r
     }
 }
 
-impl<Fp: FpTrait> Add for Polynomial<Fp> {
-    type Output = Polynomial<Fp>;
-
-    #[inline(always)]
-    fn add(self, other: Polynomial<Fp>) -> Polynomial<Fp> {
-        let mut r = self;
-        r.set_add(&other);
-        r
-    }
-}
-
-impl<Fp: FpTrait> Add for &Polynomial<Fp> {
+impl<Fp: FpTrait> Add<&Polynomial<Fp>> for &Polynomial<Fp> {
     type Output = Polynomial<Fp>;
 
     #[inline(always)]
@@ -284,32 +286,14 @@ impl<Fp: FpTrait> Add for &Polynomial<Fp> {
     }
 }
 
-impl<Fp: FpTrait> AddAssign for Polynomial<Fp> {
-    #[inline(always)]
-    fn add_assign(&mut self, other: Polynomial<Fp>) {
-        self.set_add(&other);
-    }
-}
-
 impl<Fp: FpTrait> AddAssign<&Polynomial<Fp>> for Polynomial<Fp> {
     #[inline(always)]
     fn add_assign(&mut self, other: &Polynomial<Fp>) {
         self.set_add(other);
     }
 }
 
-impl<Fp: FpTrait> Sub for Polynomial<Fp> {
-    type Output = Polynomial<Fp>;
-
-    #[inline(always)]
-    fn sub(self, other: Polynomial<Fp>) -> Polynomial<Fp> {
-        let mut r = self;
-        r.set_sub(&other);
-        r
-    }
-}
-
-impl<Fp: FpTrait> Sub for &Polynomial<Fp> {
+impl<Fp: FpTrait> Sub<&Polynomial<Fp>> for &Polynomial<Fp> {
     type Output = Polynomial<Fp>;
 
     #[inline(always)]
@@ -320,32 +304,14 @@ impl<Fp: FpTrait> Sub for &Polynomial<Fp> {
     }
 }
 
-impl<Fp: FpTrait> SubAssign for Polynomial<Fp> {
-    #[inline(always)]
-    fn sub_assign(&mut self, other: Polynomial<Fp>) {
-        self.set_sub(&other);
-    }
-}
-
 impl<Fp: FpTrait> SubAssign<&Polynomial<Fp>> for Polynomial<Fp> {
     #[inline(always)]
     fn sub_assign(&mut self, other: &Polynomial<Fp>) {
         self.set_sub(other);
     }
 }
 
-impl<Fp: FpTrait> Mul for Polynomial<Fp> {
-    type Output = Polynomial<Fp>;
-
-    #[inline(always)]
-    fn mul(self, other: Polynomial<Fp>) -> Polynomial<Fp> {
-        let mut r = self;
-        r.set_mul(&other);
-        r
-    }
-}
-
-impl<Fp: FpTrait> Mul for &Polynomial<Fp> {
+impl<Fp: FpTrait> Mul<&Polynomial<Fp>> for &Polynomial<Fp> {
     type Output = Polynomial<Fp>;
 
     #[inline(always)]
@@ -356,13 +322,6 @@ impl<Fp: FpTrait> Mul for &Polynomial<Fp> {
     }
 }
 
-impl<Fp: FpTrait> MulAssign for Polynomial<Fp> {
-    #[inline(always)]
-    fn mul_assign(&mut self, other: Polynomial<Fp>) {
-        self.set_mul(&other);
-    }
-}
-
 impl<Fp: FpTrait> MulAssign<&Polynomial<Fp>> for Polynomial<Fp> {
     #[inline(always)]
     fn mul_assign(&mut self, other: &Polynomial<Fp>) {

tests/test_poly.rs

@@ -10,41 +10,44 @@ mod test_polynomial_arithmetic {
     fn test_addition() {
         let mut rng = DRNG::from_seed("polynomial_addition".as_bytes());
 
-        let f = PR::rand(&mut rng, 5);
+        let f = PR::rand(&mut rng, 4);
         let g = PR::rand(&mut rng, 5);
-        let h = PR::rand(&mut rng, 5);
+        let h = PR::rand(&mut rng, 6);
 
         let t1 = &(&f + &g) + &h;
         let t2 = &f + &(&g + &h);
 
         assert!(t1.equals(&t2) == u32::MAX);
 
         let mut t1 = f.clone();
-        t1 += g;
-        t1 += h;
+        t1 += &g;
+        t1 += &h;
         assert!(t1.equals(&t2) == u32::MAX);
 
         let zero = PR::new_from_ele(&Fp::ZERO);
         let t1 = &f + &zero;
         assert!(t1.equals(&f) == u32::MAX);
+
+        let f_neg = -&f;
+        assert!((&f + &f_neg).is_zero() == u32::MAX);
     }
 
     #[test]
     fn test_subtraction() {
         let mut rng = DRNG::from_seed("polynomial_subtraction".as_bytes());
 
-        let f = PR::rand(&mut rng, 5);
+        let f = PR::rand(&mut rng, 4);
         let g = PR::rand(&mut rng, 5);
-        let h = PR::rand(&mut rng, 5);
+        let h = PR::rand(&mut rng, 6);
 
         let t1 = &(&f - &g) - &h;
         let t2 = &f - &(&g + &h);
 
         assert!(t1.equals(&t2) == u32::MAX);
 
         let mut t1 = f.clone();
-        t1 -= g;
-        t1 -= h;
+        t1 -= &g;
+        t1 -= &h;
         assert!(t1.equals(&t2) == u32::MAX);
 
         let zero = &f - &f;
@@ -55,26 +58,26 @@ mod test_polynomial_arithmetic {
     fn test_multiplication() {
         let mut rng = DRNG::from_seed("polynomial_multiplication".as_bytes());
 
-        let f = PR::rand(&mut rng, 5);
+        let f = PR::rand(&mut rng, 4);
         let g = PR::rand(&mut rng, 5);
-        let h = PR::rand(&mut rng, 5);
+        let h = PR::rand(&mut rng, 6);
 
         let t1 = &(&f * &g) * &h;
         let t2 = &f * &(&g * &h);
 
         assert!(t1.equals(&t2) == u32::MAX);
 
         let mut t1 = f.clone();
-        t1 *= g;
-        t1 *= h;
+        t1 *= &g;
+        t1 *= &h;
         assert!(t1.equals(&t2) == u32::MAX);
 
         let one = PR::new_from_ele(&Fp::ONE);
-        t1 *= one;
+        t1 *= &one;
         assert!(t1.equals(&t2) == u32::MAX);
 
         let zero = PR::new_from_ele(&Fp::ZERO);
-        t1 *= zero;
+        t1 *= &zero;
         assert!(t1.is_zero() == u32::MAX);
     }
 
@@ -159,4 +162,31 @@ mod test_polynomial_arithmetic {
 
         assert!(ea.equals(&test_ea) == u32::MAX);
     }
+
+    #[test]
+    fn test_degree() {
+        let f = PR::new_from_slice(&[Fp::ZERO]);
+        assert!(f.degree() == None);
+
+        let f = PR::new_from_slice(&[Fp::ONE]);
+        assert!(f.degree().unwrap() == 0);
+
+        let f = PR::new_from_slice(&[Fp::ZERO, Fp::ONE]);
+        assert!(f.degree().unwrap() == 1);
+
+        let f = PR::new_from_slice(&[Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ONE, Fp::ZERO]);
+        assert!(f.degree().unwrap() == 4);
+
+        let f = PR::new_from_slice(&[Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ZERO]);
+        assert!(f.degree() == None);
+
+        let f = PR::new_from_slice(&[Fp::ONE, Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ZERO]);
+        assert!(f.degree().unwrap() == 0);
+
+        let f = PR::new_from_slice(&[Fp::ZERO, Fp::ONE, Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ZERO]);
+        assert!(f.degree().unwrap() == 1);
+
+        let f = PR::new_from_slice(&[Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ZERO, Fp::ONE, Fp::ZERO]);
+        assert!(f.degree().unwrap() == 4);
+    }
 }