snark_verifier/util/
poly.rs

1//! Polynomial.
2
3use crate::util::{arithmetic::Field, parallelize};
4use itertools::Itertools;
5use rand::Rng;
6use std::{
7    iter::{self, Sum},
8    ops::{
9        Add, Index, IndexMut, Mul, Range, RangeFrom, RangeFull, RangeInclusive, RangeTo,
10        RangeToInclusive, Sub,
11    },
12};
13
14#[derive(Clone, Debug)]
15/// Univariate polynomial.
16pub struct Polynomial<F>(Vec<F>);
17
18impl<F> Polynomial<F> {
19    /// Initialize an univariate polynomial.
20    pub fn new(inner: Vec<F>) -> Self {
21        Self(inner)
22    }
23
24    /// Returns `true` if the `Polynomial` contains no elements.
25    pub fn is_empty(&self) -> bool {
26        self.0.is_empty()
27    }
28
29    /// Returns the length of the `Polynomial`.
30    pub fn len(&self) -> usize {
31        self.0.len()
32    }
33
34    /// Returns an iterator of the `Polynomial`.
35    pub fn iter(&self) -> impl Iterator<Item = &F> {
36        self.0.iter()
37    }
38
39    /// Returns a mutable iterator of the `Polynomial`.
40    pub fn iter_mut(&mut self) -> impl Iterator<Item = &mut F> {
41        self.0.iter_mut()
42    }
43
44    /// Into vector of coefficients.
45    pub fn to_vec(self) -> Vec<F> {
46        self.0
47    }
48}
49
50impl<F: Field> Polynomial<F> {
51    pub(crate) fn rand<R: Rng>(n: usize, mut rng: R) -> Self {
52        Self::new(iter::repeat_with(|| F::random(&mut rng)).take(n).collect())
53    }
54
55    /// Returns evaluation at given `x`.
56    pub fn evaluate(&self, x: F) -> F {
57        let evaluate_serial =
58            |coeffs: &[F]| coeffs.iter().rev().fold(F::ZERO, |acc, coeff| acc * x + coeff);
59
60        #[cfg(feature = "parallel")]
61        {
62            use crate::util::{arithmetic::powers, current_num_threads, parallelize_iter};
63            use num_integer::Integer;
64
65            let num_threads = current_num_threads();
66            if self.len() * 2 < num_threads {
67                return evaluate_serial(&self.0);
68            }
69
70            let chunk_size = Integer::div_ceil(&self.len(), &num_threads);
71            let mut results = vec![F::ZERO; num_threads];
72            parallelize_iter(
73                results
74                    .iter_mut()
75                    .zip(self.0.chunks(chunk_size))
76                    .zip(powers(x.pow_vartime([chunk_size as u64]))),
77                |((result, coeffs), scalar)| *result = evaluate_serial(coeffs) * scalar,
78            );
79            results.iter().fold(F::ZERO, |acc, result| acc + result)
80        }
81        #[cfg(not(feature = "parallel"))]
82        evaluate_serial(&self.0)
83    }
84}
85
86impl<'a, F: Field> Add<&'a Polynomial<F>> for Polynomial<F> {
87    type Output = Polynomial<F>;
88
89    fn add(mut self, rhs: &'a Polynomial<F>) -> Polynomial<F> {
90        parallelize(&mut self.0, |(lhs, start)| {
91            for (lhs, rhs) in lhs.iter_mut().zip_eq(rhs.0[start..].iter()) {
92                *lhs += *rhs;
93            }
94        });
95        self
96    }
97}
98
99impl<'a, F: Field> Sub<&'a Polynomial<F>> for Polynomial<F> {
100    type Output = Polynomial<F>;
101
102    fn sub(mut self, rhs: &'a Polynomial<F>) -> Polynomial<F> {
103        parallelize(&mut self.0, |(lhs, start)| {
104            for (lhs, rhs) in lhs.iter_mut().zip_eq(rhs.0[start..].iter()) {
105                *lhs -= *rhs;
106            }
107        });
108        self
109    }
110}
111
112impl<F: Field> Sub<F> for Polynomial<F> {
113    type Output = Polynomial<F>;
114
115    fn sub(mut self, rhs: F) -> Polynomial<F> {
116        self.0[0] -= rhs;
117        self
118    }
119}
120
121impl<F: Field> Add<F> for Polynomial<F> {
122    type Output = Polynomial<F>;
123
124    fn add(mut self, rhs: F) -> Polynomial<F> {
125        self.0[0] += rhs;
126        self
127    }
128}
129
130impl<F: Field> Mul<F> for Polynomial<F> {
131    type Output = Polynomial<F>;
132
133    fn mul(mut self, rhs: F) -> Polynomial<F> {
134        if rhs == F::ZERO {
135            return Polynomial::new(vec![F::ZERO; self.len()]);
136        }
137        if rhs == F::ONE {
138            return self;
139        }
140        parallelize(&mut self.0, |(lhs, _)| {
141            for lhs in lhs.iter_mut() {
142                *lhs *= rhs;
143            }
144        });
145        self
146    }
147}
148
149impl<F: Field> Sum for Polynomial<F> {
150    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
151        iter.reduce(|acc, item| acc + &item).unwrap()
152    }
153}
154
155macro_rules! impl_index {
156    ($($range:ty => $output:ty,)*) => {
157        $(
158            impl<F> Index<$range> for Polynomial<F> {
159                type Output = $output;
160
161                fn index(&self, index: $range) -> &$output {
162                    self.0.index(index)
163                }
164            }
165            impl<F> IndexMut<$range> for Polynomial<F> {
166                fn index_mut(&mut self, index: $range) -> &mut $output {
167                    self.0.index_mut(index)
168                }
169            }
170        )*
171    };
172}
173
174impl_index!(
175    usize => F,
176    Range<usize> => [F],
177    RangeFrom<usize> => [F],
178    RangeFull => [F],
179    RangeInclusive<usize> => [F],
180    RangeTo<usize> => [F],
181    RangeToInclusive<usize> => [F],
182);