Add sparse scalar mul and sub (#259)

* Add/Sub impl for Sparse and Dense polys, Mul<F> for SparsePolynomial * Add test and update changelog * Fmt * Update CHANGELOG.md * add add_assign and sub_assign * fmt Co-authored-by: Weikeng Chen <w.k@berkeley.edu>
arkworks-rs · May 26, 2021 · 846c90b · 846c90b
1 parent f64db16
commit 846c90b
Show file tree

Hide file tree

Showing 3 changed files with 250 additions and 17 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -9,6 +9,7 @@
 - [\#230](https://github.com/arkworks-rs/algebra/pull/230) (ark-ec) Add `wnaf_mul` implementation for `ProjectiveCurve`.
 - [\#245](https://github.com/arkworks-rs/algebra/pull/245) (ark-poly) Speedup the sequential and parallel radix-2 FFT and IFFT significantly by making the method in which it accesses roots more cache-friendly.
 - [\#258](https://github.com/arkworks-rs/algebra/pull/258) (ark-poly) Add `Mul<F>` implementation for `DensePolynomial`.
+- [\#259](https://github.com/arkworks-rs/algebra/pull/259) (ark-poly) Add `Mul<F>` implementation for `SparsePolynomial` and `Add<SparsePolynomial<F>>/Sub<SparsePolynomial<F>>` for `DensePolynomial`.
 - [\#261](https://github.com/arkworks-rs/algebra/pull/261) (ark-ff) Add support for 448-bit integers and fields.
 - [\#263](https://github.com/arkworks-rs/algebra/pull/263) (ark-ff) Add `From<iXXX>` implementations to fields.
 - [\#265](https://github.com/arkworks-rs/algebra/pull/265) (ark-serialize) Add hashing as an extension trait of `CanonicalSerialize`.

diff --git a/poly/src/polynomial/univariate/dense.rs b/poly/src/polynomial/univariate/dense.rs
@@ -1,17 +1,16 @@
 //! A dense univariate polynomial represented in coefficient form.
 use crate::univariate::DenseOrSparsePolynomial;
+use crate::{univariate::SparsePolynomial, Polynomial, UVPolynomial};
 use crate::{EvaluationDomain, Evaluations, GeneralEvaluationDomain};
-use crate::{Polynomial, UVPolynomial};
+use ark_ff::{FftField, Field, Zero};
 use ark_serialize::*;
+use ark_std::rand::Rng;
 use ark_std::{
     fmt,
     ops::{Add, AddAssign, Deref, DerefMut, Div, Mul, Neg, Sub, SubAssign},
     vec::Vec,
 };
 
-use ark_ff::{FftField, Field, Zero};
-use ark_std::rand::Rng;
-
 #[cfg(feature = "parallel")]
 use ark_std::cmp::max;
 #[cfg(feature = "parallel")]
@@ -262,6 +261,37 @@ impl<'a, 'b, F: Field> Add<&'a DensePolynomial<F>> for &'b DensePolynomial<F> {
     }
 }
 
+impl<'a, 'b, F: Field> Add<&'a SparsePolynomial<F>> for &'b DensePolynomial<F> {
+    type Output = DensePolynomial<F>;
+
+    #[inline]
+    fn add(self, other: &'a SparsePolynomial<F>) -> DensePolynomial<F> {
+        let result = if self.is_zero() {
+            other.clone().into()
+        } else if other.is_zero() {
+            self.clone()
+        } else {
+            let mut result = self.clone();
+            // If `other` has higher degree than `self`, create a dense vector
+            // storing the upper coefficients of the addition
+            let mut upper_coeffs = match other.degree() > result.degree() {
+                true => vec![F::zero(); other.degree() - result.degree()],
+                false => Vec::new(),
+            };
+            for (pow, coeff) in other.iter() {
+                if *pow <= result.degree() {
+                    result.coeffs[*pow] += coeff;
+                } else {
+                    upper_coeffs[*pow - result.degree() - 1] = *coeff;
+                }
+            }
+            result.coeffs.extend(upper_coeffs);
+            result
+        };
+        result
+    }
+}
+
 impl<'a, 'b, F: Field> AddAssign<&'a DensePolynomial<F>> for DensePolynomial<F> {
     fn add_assign(&mut self, other: &'a DensePolynomial<F>) {
         if self.is_zero() {
@@ -316,6 +346,38 @@ impl<'a, 'b, F: Field> AddAssign<(F, &'a DensePolynomial<F>)> for DensePolynomia
     }
 }
 
+impl<'a, F: Field> AddAssign<&'a SparsePolynomial<F>> for DensePolynomial<F> {
+    #[inline]
+    fn add_assign(&mut self, other: &'a SparsePolynomial<F>) {
+        if self.is_zero() {
+            self.coeffs.truncate(0);
+            self.coeffs.resize(other.degree() + 1, F::zero());
+
+            for (i, coeff) in other.iter() {
+                self.coeffs[*i] = *coeff;
+            }
+            return;
+        } else if other.is_zero() {
+            return;
+        } else {
+            // If `other` has higher degree than `self`, create a dense vector
+            // storing the upper coefficients of the addition
+            let mut upper_coeffs = match other.degree() > self.degree() {
+                true => vec![F::zero(); other.degree() - self.degree()],
+                false => Vec::new(),
+            };
+            for (pow, coeff) in other.iter() {
+                if *pow <= self.degree() {
+                    self.coeffs[*pow] += coeff;
+                } else {
+                    upper_coeffs[*pow - self.degree() - 1] = *coeff;
+                }
+            }
+            self.coeffs.extend(upper_coeffs);
+        }
+    }
+}
+
 impl<F: Field> Neg for DensePolynomial<F> {
     type Output = DensePolynomial<F>;
 
@@ -362,6 +424,38 @@ impl<'a, 'b, F: Field> Sub<&'a DensePolynomial<F>> for &'b DensePolynomial<F> {
     }
 }
 
+impl<'a, 'b, F: Field> Sub<&'a SparsePolynomial<F>> for &'b DensePolynomial<F> {
+    type Output = DensePolynomial<F>;
+
+    #[inline]
+    fn sub(self, other: &'a SparsePolynomial<F>) -> DensePolynomial<F> {
+        let result = if self.is_zero() {
+            let result = other.clone();
+            result.neg().into()
+        } else if other.is_zero() {
+            self.clone()
+        } else {
+            let mut result = self.clone();
+            // If `other` has higher degree than `self`, create a dense vector
+            // storing the upper coefficients of the subtraction
+            let mut upper_coeffs = match other.degree() > result.degree() {
+                true => vec![F::zero(); other.degree() - result.degree()],
+                false => Vec::new(),
+            };
+            for (pow, coeff) in other.iter() {
+                if *pow <= result.degree() {
+                    result.coeffs[*pow] -= coeff;
+                } else {
+                    upper_coeffs[*pow - result.degree() - 1] = -*coeff;
+                }
+            }
+            result.coeffs.extend(upper_coeffs);
+            result
+        };
+        result
+    }
+}
+
 impl<'a, 'b, F: Field> SubAssign<&'a DensePolynomial<F>> for DensePolynomial<F> {
     #[inline]
     fn sub_assign(&mut self, other: &'a DensePolynomial<F>) {
@@ -387,6 +481,38 @@ impl<'a, 'b, F: Field> SubAssign<&'a DensePolynomial<F>> for DensePolynomial<F>
     }
 }
 
+impl<'a, F: Field> SubAssign<&'a SparsePolynomial<F>> for DensePolynomial<F> {
+    #[inline]
+    fn sub_assign(&mut self, other: &'a SparsePolynomial<F>) {
+        if self.is_zero() {
+            self.coeffs.truncate(0);
+            self.coeffs.resize(other.degree() + 1, F::zero());
+
+            for (i, coeff) in other.iter() {
+                self.coeffs[*i] = (*coeff).neg();
+            }
+            return;
+        } else if other.is_zero() {
+            return;
+        } else {
+            // If `other` has higher degree than `self`, create a dense vector
+            // storing the upper coefficients of the subtraction
+            let mut upper_coeffs = match other.degree() > self.degree() {
+                true => vec![F::zero(); other.degree() - self.degree()],
+                false => Vec::new(),
+            };
+            for (pow, coeff) in other.iter() {
+                if *pow <= self.degree() {
+                    self.coeffs[*pow] -= coeff;
+                } else {
+                    upper_coeffs[*pow - self.degree() - 1] = -*coeff;
+                }
+            }
+            self.coeffs.extend(upper_coeffs);
+        }
+    }
+}
+
 impl<'a, 'b, F: Field> Div<&'a DensePolynomial<F>> for &'b DensePolynomial<F> {
     type Output = DensePolynomial<F>;
 
@@ -451,9 +577,20 @@ mod tests {
     use crate::polynomial::univariate::*;
     use crate::{EvaluationDomain, GeneralEvaluationDomain};
     use ark_ff::{Field, One, UniformRand, Zero};
-    use ark_std::test_rng;
+    use ark_std::{rand::Rng, test_rng};
     use ark_test_curves::bls12_381::Fr;
 
+    fn rand_sparse_poly<R: Rng>(degree: usize, rng: &mut R) -> SparsePolynomial<Fr> {
+        // Initialize coeffs so that its guaranteed to have a x^{degree} term
+        let mut coeffs = vec![(degree, Fr::rand(rng))];
+        for i in 0..degree {
+            if !rng.gen_bool(0.8) {
+                coeffs.push((i, Fr::rand(rng)));
+            }
+        }
+        SparsePolynomial::from_coefficients_vec(coeffs)
+    }
+
     #[test]
     fn double_polynomials_random() {
         let rng = &mut test_rng();
@@ -479,6 +616,34 @@ mod tests {
         }
     }
 
+    #[test]
+    fn add_sparse_polynomials() {
+        let rng = &mut test_rng();
+        for a_degree in 0..70 {
+            for b_degree in 0..70 {
+                let p1 = DensePolynomial::<Fr>::rand(a_degree, rng);
+                let p2 = rand_sparse_poly(b_degree, rng);
+                let res = &p1 + &p2;
+                assert_eq!(res, &p1 + &Into::<DensePolynomial<Fr>>::into(p2));
+            }
+        }
+    }
+
+    #[test]
+    fn add_assign_sparse_polynomials() {
+        let rng = &mut test_rng();
+        for a_degree in 0..70 {
+            for b_degree in 0..70 {
+                let p1 = DensePolynomial::<Fr>::rand(a_degree, rng);
+                let p2 = rand_sparse_poly(b_degree, rng);
+
+                let mut res = p1.clone();
+                res += &p2;
+                assert_eq!(res, &p1 + &Into::<DensePolynomial<Fr>>::into(p2));
+            }
+        }
+    }
+
     #[test]
     fn add_polynomials_with_mul() {
         let rng = &mut test_rng();
@@ -501,16 +666,44 @@ mod tests {
     #[test]
     fn sub_polynomials() {
         let rng = &mut test_rng();
-        let p1 = DensePolynomial::<Fr>::rand(5, rng);
-        let p2 = DensePolynomial::<Fr>::rand(3, rng);
-        let res1 = &p1 - &p2;
-        let res2 = &p2 - &p1;
-        assert_eq!(
-            &res1 + &p2,
-            p1,
-            "Subtraction should be inverse of addition!"
-        );
-        assert_eq!(res1, -res2, "p2 - p1 = -(p1 - p2)");
+        for a_degree in 0..70 {
+            for b_degree in 0..70 {
+                let p1 = DensePolynomial::<Fr>::rand(a_degree, rng);
+                let p2 = DensePolynomial::<Fr>::rand(b_degree, rng);
+                let res1 = &p1 - &p2;
+                let res2 = &p2 - &p1;
+                assert_eq!(&res1 + &p2, p1);
+                assert_eq!(res1, -res2);
+            }
+        }
+    }
+
+    #[test]
+    fn sub_sparse_polynomials() {
+        let rng = &mut test_rng();
+        for a_degree in 0..70 {
+            for b_degree in 0..70 {
+                let p1 = DensePolynomial::<Fr>::rand(a_degree, rng);
+                let p2 = rand_sparse_poly(b_degree, rng);
+                let res = &p1 - &p2;
+                assert_eq!(res, &p1 - &Into::<DensePolynomial<Fr>>::into(p2));
+            }
+        }
+    }
+
+    #[test]
+    fn sub_assign_sparse_polynomials() {
+        let rng = &mut test_rng();
+        for a_degree in 0..70 {
+            for b_degree in 0..70 {
+                let p1 = DensePolynomial::<Fr>::rand(a_degree, rng);
+                let p2 = rand_sparse_poly(b_degree, rng);
+
+                let mut res = p1.clone();
+                res -= &p2;
+                assert_eq!(res, &p1 - &Into::<DensePolynomial<Fr>>::into(p2));
+            }
+        }
     }
 
     #[test]

diff --git a/poly/src/polynomial/univariate/sparse.rs b/poly/src/polynomial/univariate/sparse.rs
@@ -8,10 +8,13 @@ use ark_std::{
     collections::BTreeMap,
     fmt,
     io::{Read, Write},
-    ops::{Add, AddAssign, Neg, SubAssign},
+    ops::{Add, AddAssign, Deref, DerefMut, Mul, Neg, SubAssign},
     vec::Vec,
 };
 
+#[cfg(feature = "parallel")]
+use rayon::prelude::*;
+
 /// Stores a sparse polynomial in coefficient form.
 #[derive(Clone, PartialEq, Eq, Hash, Default, CanonicalSerialize, CanonicalDeserialize)]
 pub struct SparsePolynomial<F: Field> {
@@ -36,14 +39,20 @@ impl<F: Field> fmt::Debug for SparsePolynomial<F> {
     }
 }
 
-impl<F: Field> core::ops::Deref for SparsePolynomial<F> {
+impl<F: Field> Deref for SparsePolynomial<F> {
     type Target = [(usize, F)];
 
     fn deref(&self) -> &[(usize, F)] {
         &self.coeffs
     }
 }
 
+impl<F: Field> DerefMut for SparsePolynomial<F> {
+    fn deref_mut(&mut self) -> &mut [(usize, F)] {
+        &mut self.coeffs
+    }
+}
+
 impl<F: Field> Polynomial<F> for SparsePolynomial<F> {
     type Point = F;
 
@@ -186,6 +195,23 @@ impl<'a, 'b, F: Field> SubAssign<&'a SparsePolynomial<F>> for SparsePolynomial<F
     }
 }
 
+impl<'a, 'b, F: Field> Mul<F> for &'b SparsePolynomial<F> {
+    type Output = SparsePolynomial<F>;
+
+    #[inline]
+    fn mul(self, elem: F) -> SparsePolynomial<F> {
+        if self.is_zero() || elem.is_zero() {
+            SparsePolynomial::zero()
+        } else {
+            let mut result = self.clone();
+            cfg_iter_mut!(result).for_each(|e| {
+                (*e).1 *= elem;
+            });
+            result
+        }
+    }
+}
+
 impl<F: Field> Zero for SparsePolynomial<F> {
     /// Returns the zero polynomial.
     fn zero() -> Self {
@@ -384,6 +410,19 @@ mod tests {
         }
     }
 
+    #[test]
+    fn mul_random_element() {
+        let rng = &mut test_rng();
+        for degree in 0..20 {
+            let a = rand_sparse_poly(degree, rng);
+            let e = Fr::rand(rng);
+            assert_eq!(
+                &a * e,
+                a.mul(&SparsePolynomial::from_coefficients_slice(&[(0, e)]))
+            )
+        }
+    }
+
     #[test]
     fn mul_polynomial() {
         // Test multiplying polynomials over their domains, and over the native representation.