snarkvm_algorithms/fft/polynomial/
mod.rs

1// Copyright 2024 Aleo Network Foundation
2// This file is part of the snarkVM library.
3
4// Licensed under the Apache License, Version 2.0 (the "License");
5// you may not use this file except in compliance with the License.
6// You may obtain a copy of the License at:
7
8// http://www.apache.org/licenses/LICENSE-2.0
9
10// Unless required by applicable law or agreed to in writing, software
11// distributed under the License is distributed on an "AS IS" BASIS,
12// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13// See the License for the specific language governing permissions and
14// limitations under the License.
15
16//! Work with sparse and dense polynomials.
17
18use crate::fft::{EvaluationDomain, Evaluations};
19use Polynomial::*;
20use snarkvm_fields::{Field, PrimeField};
21use snarkvm_utilities::{SerializationError, cfg_iter_mut, serialize::*};
22
23use anyhow::{Result, ensure};
24use std::{borrow::Cow, convert::TryInto};
25
26#[cfg(not(feature = "serial"))]
27use rayon::prelude::*;
28
29mod dense;
30pub use dense::DensePolynomial;
31
32mod sparse;
33pub use sparse::SparsePolynomial;
34
35mod multiplier;
36pub use multiplier::*;
37
38/// Represents either a sparse polynomial or a dense one.
39#[derive(Clone, Debug, PartialEq, Eq)]
40pub enum Polynomial<'a, F: Field> {
41    /// Represents the case where `self` is a sparse polynomial
42    Sparse(Cow<'a, SparsePolynomial<F>>),
43    /// Represents the case where `self` is a dense polynomial
44    Dense(Cow<'a, DensePolynomial<F>>),
45}
46
47impl<F: Field> CanonicalSerialize for Polynomial<'_, F> {
48    fn serialize_with_mode<W: Write>(&self, writer: W, compress: Compress) -> Result<(), SerializationError> {
49        match self {
50            Sparse(p) => {
51                let p: DensePolynomial<F> = p.clone().into_owned().into();
52                CanonicalSerialize::serialize_with_mode(&p.coeffs, writer, compress)
53            }
54            Dense(p) => CanonicalSerialize::serialize_with_mode(&p.coeffs, writer, compress),
55        }
56    }
57
58    fn serialized_size(&self, mode: Compress) -> usize {
59        match self {
60            Sparse(p) => {
61                let p: DensePolynomial<F> = p.clone().into_owned().into();
62                p.serialized_size(mode)
63            }
64            Dense(p) => p.serialized_size(mode),
65        }
66    }
67}
68
69impl<F: Field> Valid for Polynomial<'_, F> {
70    fn check(&self) -> Result<(), SerializationError> {
71        // Check that the polynomial contains a trailing zero coefficient.
72        let has_trailing_zero = match self {
73            Sparse(p) => p.coeffs().last().map(|(_, c)| c.is_zero()),
74            Dense(p) => p.coeffs.last().map(|c| c.is_zero()),
75        };
76        // Fail if the trailing coefficient is zero.
77        match has_trailing_zero {
78            Some(true) => Err(SerializationError::InvalidData),
79            Some(false) | None => Ok(()),
80        }
81    }
82}
83
84impl<F: Field> CanonicalDeserialize for Polynomial<'_, F> {
85    fn deserialize_with_mode<R: Read>(
86        reader: R,
87        compress: Compress,
88        validate: Validate,
89    ) -> Result<Self, SerializationError> {
90        DensePolynomial::<F>::deserialize_with_mode(reader, compress, validate).map(|e| Self::Dense(Cow::Owned(e)))
91    }
92}
93
94impl<F: Field> From<DensePolynomial<F>> for Polynomial<'_, F> {
95    fn from(other: DensePolynomial<F>) -> Self {
96        Dense(Cow::Owned(other))
97    }
98}
99
100impl<'a, F: Field> From<&'a DensePolynomial<F>> for Polynomial<'a, F> {
101    fn from(other: &'a DensePolynomial<F>) -> Self {
102        Dense(Cow::Borrowed(other))
103    }
104}
105
106impl<F: Field> From<SparsePolynomial<F>> for Polynomial<'_, F> {
107    fn from(other: SparsePolynomial<F>) -> Self {
108        Sparse(Cow::Owned(other))
109    }
110}
111
112impl<'a, F: Field> From<&'a SparsePolynomial<F>> for Polynomial<'a, F> {
113    fn from(other: &'a SparsePolynomial<F>) -> Self {
114        Sparse(Cow::Borrowed(other))
115    }
116}
117
118#[allow(clippy::from_over_into)]
119impl<F: Field> Into<DensePolynomial<F>> for Polynomial<'_, F> {
120    fn into(self) -> DensePolynomial<F> {
121        match self {
122            Dense(p) => p.into_owned(),
123            Sparse(p) => p.into_owned().into(),
124        }
125    }
126}
127
128impl<F: Field> TryInto<SparsePolynomial<F>> for Polynomial<'_, F> {
129    type Error = ();
130
131    fn try_into(self) -> Result<SparsePolynomial<F>, ()> {
132        match self {
133            Sparse(p) => Ok(p.into_owned()),
134            _ => Err(()),
135        }
136    }
137}
138
139impl<'a, F: Field> Polynomial<'a, F> {
140    /// The zero polynomial.
141    pub fn zero() -> Self {
142        Sparse(Cow::Owned(SparsePolynomial::zero()))
143    }
144
145    /// Checks if the given polynomial is zero.
146    pub fn is_zero(&self) -> bool {
147        match self {
148            Sparse(s) => s.is_zero(),
149            Dense(d) => d.is_zero(),
150        }
151    }
152
153    /// Return the degree of `self.
154    pub fn degree(&self) -> usize {
155        match self {
156            Sparse(s) => s.degree(),
157            Dense(d) => d.degree(),
158        }
159    }
160
161    #[inline]
162    pub fn leading_coefficient(&self) -> Option<&F> {
163        match self {
164            Sparse(p) => p.coeffs().last().map(|(_, c)| c),
165            Dense(p) => p.last(),
166        }
167    }
168
169    #[inline]
170    pub fn as_dense(&self) -> Option<&DensePolynomial<F>> {
171        match self {
172            Dense(p) => Some(p.as_ref()),
173            _ => None,
174        }
175    }
176
177    #[inline]
178    pub fn to_dense(&self) -> Cow<'_, DensePolynomial<F>> {
179        match self {
180            Dense(p) => Cow::Borrowed(p.as_ref()),
181            Sparse(p) => Cow::Owned(p.clone().into_owned().into()),
182        }
183    }
184
185    #[inline]
186    pub fn as_dense_mut(&mut self) -> Option<&mut DensePolynomial<F>> {
187        match self {
188            Dense(p) => Some(p.to_mut()),
189            _ => None,
190        }
191    }
192
193    #[inline]
194    pub fn as_sparse(&self) -> Option<&SparsePolynomial<F>> {
195        match self {
196            Sparse(p) => Some(p.as_ref()),
197            _ => None,
198        }
199    }
200
201    #[inline]
202    pub fn into_dense(&self) -> DensePolynomial<F> {
203        self.clone().into()
204    }
205
206    #[inline]
207    pub fn evaluate(&self, point: F) -> F {
208        match self {
209            Sparse(p) => p.evaluate(point),
210            Dense(p) => p.evaluate(point),
211        }
212    }
213
214    pub fn coeffs(&'a self) -> Box<dyn Iterator<Item = (usize, &'a F)> + 'a> {
215        match self {
216            Sparse(p) => Box::new(p.coeffs().map(|(c, f)| (*c, f))),
217            Dense(p) => Box::new(p.coeffs.iter().enumerate()),
218        }
219    }
220
221    /// Divide self by another (sparse or dense) polynomial, and returns the quotient and remainder.
222    pub fn divide_with_q_and_r(&self, divisor: &Self) -> Result<(DensePolynomial<F>, DensePolynomial<F>)> {
223        ensure!(!divisor.is_zero(), "Dividing by zero polynomial is undefined");
224
225        if self.is_zero() {
226            Ok((DensePolynomial::zero(), DensePolynomial::zero()))
227        } else if self.degree() < divisor.degree() {
228            Ok((DensePolynomial::zero(), self.clone().into()))
229        } else {
230            // Now we know that self.degree() >= divisor.degree();
231            let mut quotient = vec![F::zero(); self.degree() - divisor.degree() + 1];
232            let mut remainder: DensePolynomial<F> = self.clone().into();
233            // Can unwrap here because we know self is not zero.
234            let divisor_leading_inv = divisor.leading_coefficient().unwrap().inverse().unwrap();
235            while !remainder.is_zero() && remainder.degree() >= divisor.degree() {
236                let cur_q_coeff = *remainder.coeffs.last().unwrap() * divisor_leading_inv;
237                let cur_q_degree = remainder.degree() - divisor.degree();
238                quotient[cur_q_degree] = cur_q_coeff;
239
240                if let Sparse(p) = divisor {
241                    for (i, div_coeff) in p.coeffs() {
242                        remainder[cur_q_degree + i] -= &(cur_q_coeff * div_coeff);
243                    }
244                } else if let Dense(p) = divisor {
245                    for (i, div_coeff) in p.iter().enumerate() {
246                        remainder[cur_q_degree + i] -= &(cur_q_coeff * div_coeff);
247                    }
248                }
249
250                while let Some(true) = remainder.coeffs.last().map(|c| c.is_zero()) {
251                    remainder.coeffs.pop();
252                }
253            }
254            Ok((DensePolynomial::from_coefficients_vec(quotient), remainder))
255        }
256    }
257}
258
259impl<F: PrimeField> Polynomial<'_, F> {
260    /// Construct `Evaluations` by evaluating a polynomial over the domain `domain`.
261    pub fn evaluate_over_domain(poly: impl Into<Self>, domain: EvaluationDomain<F>) -> Evaluations<F> {
262        let poly = poly.into();
263        poly.eval_over_domain_helper(domain)
264    }
265
266    fn eval_over_domain_helper(self, domain: EvaluationDomain<F>) -> Evaluations<F> {
267        match self {
268            Sparse(Cow::Borrowed(s)) => {
269                let evals = domain.elements().map(|elem| s.evaluate(elem)).collect();
270                Evaluations::from_vec_and_domain(evals, domain)
271            }
272            Sparse(Cow::Owned(s)) => {
273                let evals = domain.elements().map(|elem| s.evaluate(elem)).collect();
274                Evaluations::from_vec_and_domain(evals, domain)
275            }
276            Dense(Cow::Borrowed(d)) => {
277                if d.degree() >= domain.size() {
278                    d.coeffs
279                        .chunks(domain.size())
280                        .map(|d| Evaluations::from_vec_and_domain(domain.fft(d), domain))
281                        .fold(Evaluations::from_vec_and_domain(vec![F::zero(); domain.size()], domain), |mut acc, e| {
282                            cfg_iter_mut!(acc.evaluations).zip(e.evaluations).for_each(|(a, e)| *a += e);
283                            acc
284                        })
285                } else {
286                    Evaluations::from_vec_and_domain(domain.fft(&d.coeffs), domain)
287                }
288            }
289            Dense(Cow::Owned(mut d)) => {
290                if d.degree() >= domain.size() {
291                    d.coeffs
292                        .chunks(domain.size())
293                        .map(|d| Evaluations::from_vec_and_domain(domain.fft(d), domain))
294                        .fold(Evaluations::from_vec_and_domain(vec![F::zero(); domain.size()], domain), |mut acc, e| {
295                            cfg_iter_mut!(acc.evaluations).zip(e.evaluations).for_each(|(a, e)| *a += e);
296                            acc
297                        })
298                } else {
299                    domain.fft_in_place(&mut d.coeffs);
300                    Evaluations::from_vec_and_domain(d.coeffs, domain)
301                }
302            }
303        }
304    }
305}