1use 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#[derive(Educe, CanonicalSerialize, CanonicalDeserialize)]
25#[educe(Clone, PartialEq, Eq, Hash, Default)]
26pub struct SparsePolynomial<F: Field, T: Term> {
27 #[educe(PartialEq(ignore))]
29 pub num_vars: usize,
30 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 fn degree(&self) -> usize {
60 self.terms
61 .iter()
62 .map(|(_, term)| term.degree())
63 .max()
64 .unwrap_or_default()
65 }
66
67 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 fn num_vars(&self) -> usize {
99 self.num_vars
100 }
101
102 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 fn from_coefficients_vec(num_vars: usize, mut terms: Vec<(F, SparseTerm)>) -> Self {
139 terms.sort_by(|(_, t1), (_, t2)| t1.cmp(t2));
141 let mut terms_dedup: Vec<(F, SparseTerm)> = Vec::new();
143 for (coeff, term) in terms {
144 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 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 result.remove_zeros();
166 result
167 }
168
169 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 loop {
193 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 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 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 *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 fn zero() -> Self {
288 Self::default()
289 }
290
291 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 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 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 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 let eval1 = p1.evaluate(&point);
411 let eval2 = p2.evaluate(&point);
412 let sum = &p1 + &p2;
414 let eval3 = sum.evaluate(&point);
416 assert_eq!(eval1 + eval2, eval3);
417 }
418 }
419 }
420
421 #[test]
422 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)])), (Fr::from(2), SparseTerm::new(vec![(1, 1)])),
476 ],
477 );
478
479 let sum = &p1 + &p2;
480
481 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)); }
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 let poly1 = SparsePolynomial::<Fr, SparseTerm>::from_coefficients_vec(
535 3,
536 vec![
537 (Fr::from(2), SparseTerm::new(vec![(0, 3)])), (Fr::from(1), SparseTerm::new(vec![(1, 1), (2, 1)])), ],
540 );
541
542 let poly2 = SparsePolynomial::<Fr, SparseTerm>::from_coefficients_vec(
544 3,
545 vec![
546 (Fr::from(1), SparseTerm::new(vec![(0, 2)])), (Fr::from(1), SparseTerm::new(vec![(1, 2)])), (Fr::from(1), SparseTerm::new(vec![])), ],
550 );
551
552 let poly3 = SparsePolynomial::<Fr, SparseTerm>::from_coefficients_vec(
554 3,
555 vec![
556 (Fr::from(3), SparseTerm::new(vec![])), ],
558 );
559
560 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}