From 846c90beafb5e5e43958a2ba476602a3d9372e66 Mon Sep 17 00:00:00 2001 From: Ryan Lehmkuhl Date: Tue, 25 May 2021 22:25:51 -0700 Subject: [PATCH] Add sparse scalar mul and sub (#259) * Add/Sub impl for Sparse and Dense polys, Mul for SparsePolynomial * Add test and update changelog * Fmt * Update CHANGELOG.md * add add_assign and sub_assign * fmt Co-authored-by: Weikeng Chen --- CHANGELOG.md | 1 + poly/src/polynomial/univariate/dense.rs | 223 +++++++++++++++++++++-- poly/src/polynomial/univariate/sparse.rs | 43 ++++- 3 files changed, 250 insertions(+), 17 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7400f4382..f4e276f15 100644 --- 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` implementation for `DensePolynomial`. +- [\#259](https://github.com/arkworks-rs/algebra/pull/259) (ark-poly) Add `Mul` implementation for `SparsePolynomial` and `Add>/Sub>` 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` 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 index a86474904..f49671745 100644 --- 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> for &'b DensePolynomial { } } +impl<'a, 'b, F: Field> Add<&'a SparsePolynomial> for &'b DensePolynomial { + type Output = DensePolynomial; + + #[inline] + fn add(self, other: &'a SparsePolynomial) -> DensePolynomial { + 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> for DensePolynomial { fn add_assign(&mut self, other: &'a DensePolynomial) { if self.is_zero() { @@ -316,6 +346,38 @@ impl<'a, 'b, F: Field> AddAssign<(F, &'a DensePolynomial)> for DensePolynomia } } +impl<'a, F: Field> AddAssign<&'a SparsePolynomial> for DensePolynomial { + #[inline] + fn add_assign(&mut self, other: &'a SparsePolynomial) { + 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 Neg for DensePolynomial { type Output = DensePolynomial; @@ -362,6 +424,38 @@ impl<'a, 'b, F: Field> Sub<&'a DensePolynomial> for &'b DensePolynomial { } } +impl<'a, 'b, F: Field> Sub<&'a SparsePolynomial> for &'b DensePolynomial { + type Output = DensePolynomial; + + #[inline] + fn sub(self, other: &'a SparsePolynomial) -> DensePolynomial { + 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> for DensePolynomial { #[inline] fn sub_assign(&mut self, other: &'a DensePolynomial) { @@ -387,6 +481,38 @@ impl<'a, 'b, F: Field> SubAssign<&'a DensePolynomial> for DensePolynomial } } +impl<'a, F: Field> SubAssign<&'a SparsePolynomial> for DensePolynomial { + #[inline] + fn sub_assign(&mut self, other: &'a SparsePolynomial) { + 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> for &'b DensePolynomial { type Output = DensePolynomial; @@ -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(degree: usize, rng: &mut R) -> SparsePolynomial { + // 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::::rand(a_degree, rng); + let p2 = rand_sparse_poly(b_degree, rng); + let res = &p1 + &p2; + assert_eq!(res, &p1 + &Into::>::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::::rand(a_degree, rng); + let p2 = rand_sparse_poly(b_degree, rng); + + let mut res = p1.clone(); + res += &p2; + assert_eq!(res, &p1 + &Into::>::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::::rand(5, rng); - let p2 = DensePolynomial::::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::::rand(a_degree, rng); + let p2 = DensePolynomial::::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::::rand(a_degree, rng); + let p2 = rand_sparse_poly(b_degree, rng); + let res = &p1 - &p2; + assert_eq!(res, &p1 - &Into::>::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::::rand(a_degree, rng); + let p2 = rand_sparse_poly(b_degree, rng); + + let mut res = p1.clone(); + res -= &p2; + assert_eq!(res, &p1 - &Into::>::into(p2)); + } + } } #[test] diff --git a/poly/src/polynomial/univariate/sparse.rs b/poly/src/polynomial/univariate/sparse.rs index 1b72bce1e..3e68dd320 100644 --- 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 { @@ -36,7 +39,7 @@ impl fmt::Debug for SparsePolynomial { } } -impl core::ops::Deref for SparsePolynomial { +impl Deref for SparsePolynomial { type Target = [(usize, F)]; fn deref(&self) -> &[(usize, F)] { @@ -44,6 +47,12 @@ impl core::ops::Deref for SparsePolynomial { } } +impl DerefMut for SparsePolynomial { + fn deref_mut(&mut self) -> &mut [(usize, F)] { + &mut self.coeffs + } +} + impl Polynomial for SparsePolynomial { type Point = F; @@ -186,6 +195,23 @@ impl<'a, 'b, F: Field> SubAssign<&'a SparsePolynomial> for SparsePolynomial Mul for &'b SparsePolynomial { + type Output = SparsePolynomial; + + #[inline] + fn mul(self, elem: F) -> SparsePolynomial { + 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 Zero for SparsePolynomial { /// 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.