Skip to main content

ark_poly/polynomial/multivariate/
sparse.rs

1//! A sparse multivariate polynomial represented in coefficient form.
2use crate::{
3    multivariate::{SparseTerm, Term},
4    DenseMVPolynomial, Polynomial,
5};
6use ark_ff::{Field, Zero};
7use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
8use ark_std::{
9    cfg_into_iter,
10    cmp::Ordering,
11    fmt,
12    ops::{Add, AddAssign, Neg, Sub, SubAssign},
13    rand::Rng,
14    vec,
15    vec::*,
16};
17
18use educe::Educe;
19
20#[cfg(feature = "parallel")]
21use rayon::prelude::*;
22
23/// Stores a sparse multivariate polynomial in coefficient form.
24#[derive(Educe, CanonicalSerialize, CanonicalDeserialize)]
25#[educe(Clone, PartialEq, Eq, Hash, Default)]
26pub struct SparsePolynomial<F: Field, T: Term> {
27    /// The number of variables the polynomial supports
28    #[educe(PartialEq(ignore))]
29    pub num_vars: usize,
30    /// List of each term along with its coefficient
31    pub terms: Vec<(F, T)>,
32}
33
34impl<F: Field, T: Term> SparsePolynomial<F, T> {
35    fn remove_zeros(&mut self) {
36        self.terms.retain(|(c, _)| !c.is_zero());
37    }
38}
39
40impl<F: Field> Polynomial<F> for SparsePolynomial<F, SparseTerm> {
41    type Point = Vec<F>;
42
43    /// Returns the total degree of the polynomial
44    ///
45    /// # Examples
46    /// ```
47    /// use ark_poly::{
48    ///     polynomial::multivariate::{SparsePolynomial, SparseTerm},
49    ///     DenseMVPolynomial, Polynomial,
50    /// };
51    /// use ark_std::test_rng;
52    /// use ark_test_curves::bls12_381::Fq;
53    ///
54    /// let rng = &mut test_rng();
55    /// // Create a multivariate polynomial of degree 7
56    /// let poly: SparsePolynomial<Fq, SparseTerm> = SparsePolynomial::rand(7, 2, rng);
57    /// assert_eq!(poly.degree(), 7);
58    /// ```
59    fn degree(&self) -> usize {
60        self.terms
61            .iter()
62            .map(|(_, term)| term.degree())
63            .max()
64            .unwrap_or_default()
65    }
66
67    /// Evaluates `self` at the given `point` in `Self::Point`.
68    ///
69    /// # Examples
70    /// ```
71    /// use ark_ff::UniformRand;
72    /// use ark_poly::{
73    ///     polynomial::multivariate::{SparsePolynomial, SparseTerm, Term},
74    ///     DenseMVPolynomial, Polynomial,
75    /// };
76    /// use ark_std::test_rng;
77    /// use ark_test_curves::bls12_381::Fq;
78    ///
79    /// let rng = &mut test_rng();
80    /// let poly = SparsePolynomial::rand(4, 3, rng);
81    /// let random_point = vec![Fq::rand(rng); 3];
82    /// // The result will be a single element in the field.
83    /// let result: Fq = poly.evaluate(&random_point);
84    /// ```
85    fn evaluate(&self, point: &Vec<F>) -> F {
86        assert!(point.len() >= self.num_vars, "Invalid evaluation domain");
87        if self.is_zero() {
88            return F::zero();
89        }
90        cfg_into_iter!(&self.terms)
91            .map(|(coeff, term)| *coeff * term.evaluate(point))
92            .sum()
93    }
94}
95
96impl<F: Field> DenseMVPolynomial<F> for SparsePolynomial<F, SparseTerm> {
97    /// Returns the number of variables in `self`
98    fn num_vars(&self) -> usize {
99        self.num_vars
100    }
101
102    /// Outputs an `l`-variate polynomial which is the sum of `l` `d`-degree
103    /// univariate polynomials where each coefficient is sampled uniformly at random.
104    fn rand<R: Rng>(d: usize, l: usize, rng: &mut R) -> Self {
105        let mut random_terms = vec![(F::rand(rng), SparseTerm::new(vec![]))];
106        for var in 0..l {
107            for deg in 1..=d {
108                random_terms.push((F::rand(rng), SparseTerm::new(vec![(var, deg)])));
109            }
110        }
111        Self::from_coefficients_vec(l, random_terms)
112    }
113
114    type Term = SparseTerm;
115
116    /// Constructs a new polynomial from a list of tuples of the form `(coeff, Self::Term)`
117    ///
118    /// # Examples
119    /// ```
120    /// use ark_poly::{
121    ///     polynomial::multivariate::{SparsePolynomial, SparseTerm, Term},
122    ///     DenseMVPolynomial, Polynomial,
123    /// };
124    /// use ark_test_curves::bls12_381::Fq;
125    ///
126    /// // Create a multivariate polynomial in 3 variables, with 4 terms:
127    /// // 2*x_0^3 + x_0*x_2 + x_1*x_2 + 5
128    /// let poly = SparsePolynomial::from_coefficients_vec(
129    ///     3,
130    ///     vec![
131    ///         (Fq::from(2), SparseTerm::new(vec![(0, 3)])),
132    ///         (Fq::from(1), SparseTerm::new(vec![(0, 1), (2, 1)])),
133    ///         (Fq::from(1), SparseTerm::new(vec![(1, 1), (2, 1)])),
134    ///         (Fq::from(5), SparseTerm::new(vec![])),
135    ///     ],
136    /// );
137    /// ```
138    fn from_coefficients_vec(num_vars: usize, mut terms: Vec<(F, SparseTerm)>) -> Self {
139        // Ensure that terms are in ascending order.
140        terms.sort_by(|(_, t1), (_, t2)| t1.cmp(t2));
141        // If any terms are duplicated, add them together
142        let mut terms_dedup: Vec<(F, SparseTerm)> = Vec::new();
143        for (coeff, term) in terms {
144            // Assert correct number of indeterminates
145            assert!(
146                term.iter().all(|(var, _)| *var < num_vars),
147                "Invalid number of indeterminates"
148            );
149
150            if let Some((prev_coeff, prev_term)) = terms_dedup.last_mut() {
151                // If terms match, add the coefficients.
152                if prev_term == &term {
153                    *prev_coeff += coeff;
154                    continue;
155                }
156            }
157
158            terms_dedup.push((coeff, term));
159        }
160        let mut result = Self {
161            num_vars,
162            terms: terms_dedup,
163        };
164        // Remove any terms with zero coefficients
165        result.remove_zeros();
166        result
167    }
168
169    /// Returns the terms of a `self` as a list of tuples of the form `(coeff, Self::Term)`
170    fn terms(&self) -> &[(F, Self::Term)] {
171        self.terms.as_slice()
172    }
173}
174
175impl<F: Field, T: Term> Add for SparsePolynomial<F, T> {
176    type Output = Self;
177
178    fn add(self, other: Self) -> Self {
179        &self + &other
180    }
181}
182
183impl<'a, F: Field, T: Term> Add<&'a SparsePolynomial<F, T>> for &SparsePolynomial<F, T> {
184    type Output = SparsePolynomial<F, T>;
185
186    fn add(self, other: &'a SparsePolynomial<F, T>) -> SparsePolynomial<F, T> {
187        let mut result = Vec::new();
188        let mut cur_iter = self.terms.iter().peekable();
189        let mut other_iter = other.terms.iter().peekable();
190        // Since both polynomials are sorted, iterate over them in ascending order,
191        // combining any common terms
192        loop {
193            // Peek at iterators to decide which to take from
194            let which = match (cur_iter.peek(), other_iter.peek()) {
195                (Some(cur), Some(other)) => Some((cur.1).cmp(&other.1)),
196                (Some(_), None) => Some(Ordering::Less),
197                (None, Some(_)) => Some(Ordering::Greater),
198                (None, None) => None,
199            };
200            // Push the smallest element to the `result` coefficient vec
201            let smallest = match which {
202                Some(Ordering::Less) => cur_iter.next().unwrap().clone(),
203                Some(Ordering::Equal) => {
204                    let other = other_iter.next().unwrap();
205                    let cur = cur_iter.next().unwrap();
206                    (cur.0 + other.0, cur.1.clone())
207                },
208                Some(Ordering::Greater) => other_iter.next().unwrap().clone(),
209                None => break,
210            };
211            result.push(smallest);
212        }
213        // Remove any zero terms
214        result.retain(|(c, _)| !c.is_zero());
215        SparsePolynomial {
216            num_vars: core::cmp::max(self.num_vars, other.num_vars),
217            terms: result,
218        }
219    }
220}
221
222impl<'a, F: Field, T: Term> AddAssign<&'a Self> for SparsePolynomial<F, T> {
223    fn add_assign(&mut self, other: &'a Self) {
224        *self = &*self + other;
225    }
226}
227
228impl<'a, F: Field, T: Term> AddAssign<(F, &'a Self)> for SparsePolynomial<F, T> {
229    fn add_assign(&mut self, (f, other): (F, &'a Self)) {
230        let other = Self {
231            num_vars: other.num_vars,
232            terms: other
233                .terms
234                .iter()
235                .map(|(coeff, term)| (*coeff * f, term.clone()))
236                .collect(),
237        };
238        // Note the call to `Add` will remove also any duplicates
239        *self = &*self + &other;
240    }
241}
242
243impl<F: Field, T: Term> Neg for SparsePolynomial<F, T> {
244    type Output = Self;
245
246    #[inline]
247    fn neg(mut self) -> Self {
248        for coeff in &mut self.terms {
249            (coeff).0 = -coeff.0;
250        }
251        self
252    }
253}
254
255impl<'a, F: Field, T: Term> Sub<&'a SparsePolynomial<F, T>> for &SparsePolynomial<F, T> {
256    type Output = SparsePolynomial<F, T>;
257
258    #[inline]
259    fn sub(self, other: &'a SparsePolynomial<F, T>) -> SparsePolynomial<F, T> {
260        let neg_other = other.clone().neg();
261        self + &neg_other
262    }
263}
264
265impl<'a, F: Field, T: Term> SubAssign<&'a Self> for SparsePolynomial<F, T> {
266    #[inline]
267    fn sub_assign(&mut self, other: &'a Self) {
268        *self = &*self - other;
269    }
270}
271
272impl<F: Field, T: Term> fmt::Debug for SparsePolynomial<F, T> {
273    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> {
274        for (coeff, term) in self.terms.iter().filter(|(c, _)| !c.is_zero()) {
275            if term.is_constant() {
276                write!(f, "\n{coeff:?}")?;
277            } else {
278                write!(f, "\n{:?} {:?}", coeff, term)?;
279            }
280        }
281        Ok(())
282    }
283}
284
285impl<F: Field, T: Term> Zero for SparsePolynomial<F, T> {
286    /// Returns the zero polynomial.
287    fn zero() -> Self {
288        Self::default()
289    }
290
291    /// Checks if the given polynomial is zero.
292    fn is_zero(&self) -> bool {
293        self.terms.is_empty() || self.terms.iter().all(|(c, _)| c.is_zero())
294    }
295}
296
297#[cfg(test)]
298mod tests {
299    use super::*;
300    use ark_ff::UniformRand;
301    use ark_std::test_rng;
302    use ark_test_curves::bls12_381::Fr;
303
304    // TODO: Make tests generic over term type
305
306    /// Generate random `l`-variate polynomial of maximum individual degree `d`
307    fn rand_poly<R: Rng>(l: usize, d: usize, rng: &mut R) -> SparsePolynomial<Fr, SparseTerm> {
308        let mut random_terms = Vec::new();
309        let num_terms = rng.gen_range(1..1000);
310        // For each term, randomly select up to `l` variables with degree
311        // in [1,d] and random coefficient
312        random_terms.push((Fr::rand(rng), SparseTerm::new(vec![])));
313        for _ in 1..num_terms {
314            let term = (0..l)
315                .filter_map(|i| rng.gen_bool(0.5).then(|| (i, rng.gen_range(1..(d + 1)))))
316                .collect();
317            let coeff = Fr::rand(rng);
318            random_terms.push((coeff, SparseTerm::new(term)));
319        }
320        SparsePolynomial::from_coefficients_slice(l, &random_terms)
321    }
322
323    /// Perform a naive n^2 multiplication of `self` by `other`.
324    fn naive_mul(
325        cur: &SparsePolynomial<Fr, SparseTerm>,
326        other: &SparsePolynomial<Fr, SparseTerm>,
327    ) -> SparsePolynomial<Fr, SparseTerm> {
328        if cur.is_zero() || other.is_zero() {
329            SparsePolynomial::zero()
330        } else {
331            let mut result_terms = Vec::new();
332            for (cur_coeff, cur_term) in &cur.terms {
333                for (other_coeff, other_term) in &other.terms {
334                    let mut term = cur_term.0.clone();
335                    term.extend(other_term.0.clone());
336                    result_terms.push((*cur_coeff * *other_coeff, SparseTerm::new(term)));
337                }
338            }
339            SparsePolynomial::from_coefficients_slice(cur.num_vars, result_terms.as_slice())
340        }
341    }
342
343    #[test]
344    fn add_polynomials() {
345        let rng = &mut test_rng();
346        let max_degree = 10;
347        for a_var_count in 1..20 {
348            for b_var_count in 1..20 {
349                let p1 = rand_poly(a_var_count, max_degree, rng);
350                let p2 = rand_poly(b_var_count, max_degree, rng);
351                let res1 = &p1 + &p2;
352                let res2 = &p2 + &p1;
353                assert_eq!(res1, res2);
354            }
355        }
356    }
357
358    #[test]
359    fn sub_polynomials() {
360        let rng = &mut test_rng();
361        let max_degree = 10;
362        for a_var_count in 1..20 {
363            for b_var_count in 1..20 {
364                let p1 = rand_poly(a_var_count, max_degree, rng);
365                let p2 = rand_poly(b_var_count, max_degree, rng);
366                let res1 = &p1 - &p2;
367                let res2 = &p2 - &p1;
368                assert_eq!(&res1 + &p2, p1);
369                assert_eq!(res1, -res2);
370            }
371        }
372    }
373
374    #[test]
375    fn evaluate_polynomials() {
376        let rng = &mut test_rng();
377        let max_degree = 10;
378        for var_count in 1..20 {
379            let p = rand_poly(var_count, max_degree, rng);
380            let mut point = Vec::with_capacity(var_count);
381            for _ in 0..var_count {
382                point.push(Fr::rand(rng));
383            }
384            let mut total = Fr::zero();
385            for (coeff, term) in &p.terms {
386                let mut summand = *coeff;
387                for var in term.iter() {
388                    let eval = point.get(var.0).unwrap();
389                    summand *= eval.pow([var.1 as u64]);
390                }
391                total += summand;
392            }
393            assert_eq!(p.evaluate(&point), total);
394        }
395    }
396
397    #[test]
398    fn add_and_evaluate_polynomials() {
399        let rng = &mut test_rng();
400        let max_degree = 10;
401        for a_var_count in 1..20 {
402            for b_var_count in 1..20 {
403                let p1 = rand_poly(a_var_count, max_degree, rng);
404                let p2 = rand_poly(b_var_count, max_degree, rng);
405                let mut point = Vec::new();
406                for _ in 0..core::cmp::max(a_var_count, b_var_count) {
407                    point.push(Fr::rand(rng));
408                }
409                // Evaluate both polynomials at a given point
410                let eval1 = p1.evaluate(&point);
411                let eval2 = p2.evaluate(&point);
412                // Add polynomials
413                let sum = &p1 + &p2;
414                // Evaluate result at same point
415                let eval3 = sum.evaluate(&point);
416                assert_eq!(eval1 + eval2, eval3);
417            }
418        }
419    }
420
421    #[test]
422    /// This is just to make sure naive_mul works as expected
423    fn mul_polynomials_fixed() {
424        let a = SparsePolynomial::from_coefficients_slice(
425            4,
426            &[
427                ("2".parse().unwrap(), SparseTerm(vec![])),
428                ("4".parse().unwrap(), SparseTerm(vec![(0, 1), (1, 2)])),
429                ("8".parse().unwrap(), SparseTerm(vec![(0, 1), (0, 1)])),
430                ("1".parse().unwrap(), SparseTerm(vec![(3, 0)])),
431            ],
432        );
433        let b = SparsePolynomial::from_coefficients_slice(
434            4,
435            &[
436                ("1".parse().unwrap(), SparseTerm(vec![(0, 1), (1, 2)])),
437                ("2".parse().unwrap(), SparseTerm(vec![(2, 1)])),
438                ("1".parse().unwrap(), SparseTerm(vec![(3, 1)])),
439            ],
440        );
441        let result = naive_mul(&a, &b);
442        let expected = SparsePolynomial::from_coefficients_slice(
443            4,
444            &[
445                ("3".parse().unwrap(), SparseTerm(vec![(0, 1), (1, 2)])),
446                ("6".parse().unwrap(), SparseTerm(vec![(2, 1)])),
447                ("3".parse().unwrap(), SparseTerm(vec![(3, 1)])),
448                ("4".parse().unwrap(), SparseTerm(vec![(0, 2), (1, 4)])),
449                (
450                    "8".parse().unwrap(),
451                    SparseTerm(vec![(0, 1), (1, 2), (2, 1)]),
452                ),
453                (
454                    "4".parse().unwrap(),
455                    SparseTerm(vec![(0, 1), (1, 2), (3, 1)]),
456                ),
457                ("8".parse().unwrap(), SparseTerm(vec![(0, 3), (1, 2)])),
458                ("16".parse().unwrap(), SparseTerm(vec![(0, 2), (2, 1)])),
459                ("8".parse().unwrap(), SparseTerm(vec![(0, 2), (3, 1)])),
460            ],
461        );
462        assert_eq!(expected, result);
463    }
464
465    #[test]
466    fn test_polynomial_with_zero_coefficients() {
467        let rng = &mut test_rng();
468        let max_degree = 10;
469        let p1 = rand_poly(3, max_degree, rng);
470
471        let p2 = SparsePolynomial::from_coefficients_vec(
472            3,
473            vec![
474                (Fr::zero(), SparseTerm::new(vec![(0, 1)])), // A zero coefficient term
475                (Fr::from(2), SparseTerm::new(vec![(1, 1)])),
476            ],
477        );
478
479        let sum = &p1 + &p2;
480
481        // Ensure that the zero coefficient term is ignored in the evaluation.
482        let point = vec![Fr::from(1), Fr::from(2), Fr::from(3)];
483        let result = sum.evaluate(&point);
484
485        assert_eq!(result, p1.evaluate(&point) + p2.evaluate(&point)); // Should be the sum of the evaluations
486    }
487
488    #[test]
489    fn test_constant_polynomial() {
490        let constant_term = SparsePolynomial::from_coefficients_vec(
491            3,
492            vec![(Fr::from(5), SparseTerm::new(vec![]))],
493        );
494
495        let point = vec![Fr::from(1), Fr::from(2), Fr::from(3)];
496        assert_eq!(constant_term.evaluate(&point), Fr::from(5));
497    }
498
499    #[test]
500    fn test_polynomial_addition_with_overlapping_terms() {
501        let poly1 = SparsePolynomial::from_coefficients_vec(
502            3,
503            vec![
504                (Fr::from(2), SparseTerm::new(vec![(0, 1)])),
505                (Fr::from(3), SparseTerm::new(vec![(1, 1)])),
506            ],
507        );
508
509        let poly2 = SparsePolynomial::from_coefficients_vec(
510            3,
511            vec![
512                (Fr::from(4), SparseTerm::new(vec![(0, 1)])),
513                (Fr::from(1), SparseTerm::new(vec![(2, 1)])),
514            ],
515        );
516
517        let expected = SparsePolynomial::from_coefficients_vec(
518            3,
519            vec![
520                (Fr::from(6), SparseTerm::new(vec![(0, 1)])),
521                (Fr::from(3), SparseTerm::new(vec![(1, 1)])),
522                (Fr::from(1), SparseTerm::new(vec![(2, 1)])),
523            ],
524        );
525
526        let result = &poly1 + &poly2;
527
528        assert_eq!(expected, result);
529    }
530
531    #[test]
532    fn test_polynomial_degree() {
533        // Polynomial: 2*x_0^3 + x_1*x_2
534        let poly1 = SparsePolynomial::<Fr, SparseTerm>::from_coefficients_vec(
535            3,
536            vec![
537                (Fr::from(2), SparseTerm::new(vec![(0, 3)])), // term with degree 3
538                (Fr::from(1), SparseTerm::new(vec![(1, 1), (2, 1)])), // term with degree 2
539            ],
540        );
541
542        // Polynomial: x_0^2 + x_1^2 + 1
543        let poly2 = SparsePolynomial::<Fr, SparseTerm>::from_coefficients_vec(
544            3,
545            vec![
546                (Fr::from(1), SparseTerm::new(vec![(0, 2)])), // term with degree 2
547                (Fr::from(1), SparseTerm::new(vec![(1, 2)])), // term with degree 2
548                (Fr::from(1), SparseTerm::new(vec![])),       // constant term (degree 0)
549            ],
550        );
551
552        // Polynomial: 3
553        let poly3 = SparsePolynomial::<Fr, SparseTerm>::from_coefficients_vec(
554            3,
555            vec![
556                (Fr::from(3), SparseTerm::new(vec![])), // constant term (degree 0)
557            ],
558        );
559
560        // Test the degree method
561        assert_eq!(poly1.degree(), 3, "Degree of poly1 should be 3");
562        assert_eq!(poly2.degree(), 2, "Degree of poly2 should be 2");
563        assert_eq!(poly3.degree(), 0, "Degree of poly3 should be 0");
564    }
565}