Skip to main content

o1_utils/
dense_polynomial.rs

1//! This adds a few utility functions for the [`DensePolynomial`] arkworks type.
2
3use ark_ff::Field;
4use ark_poly::{univariate::DensePolynomial, DenseUVPolynomial, Polynomial};
5use rayon::prelude::*;
6
7use crate::chunked_polynomial::ChunkedPolynomial;
8
9//
10// ExtendedDensePolynomial trait
11//
12
13/// An extension for the [`DensePolynomial`] type.
14pub trait ExtendedDensePolynomial<F: Field> {
15    /// This function "scales" (multiplies all the coefficients of) a polynomial with a scalar.
16    #[must_use]
17    fn scale(&self, elm: F) -> Self;
18
19    /// Shifts all the coefficients to the right.
20    #[must_use]
21    fn shiftr(&self, size: usize) -> Self;
22
23    /// Evaluates a polynomial given its coefficients `coeffs` and a point `x`.
24    fn eval_polynomial(coeffs: &[F], x: F) -> F;
25
26    /// Convert a polynomial into chunks.
27    /// Implementors must ensure that the result contains exactly `num_chunks`.
28    fn to_chunked_polynomial(&self, num_chunks: usize, size: usize) -> ChunkedPolynomial<F>;
29}
30
31impl<F: Field> ExtendedDensePolynomial<F> for DensePolynomial<F> {
32    fn scale(&self, elm: F) -> Self {
33        let mut result = self.clone();
34        result
35            .coeffs
36            .par_iter_mut()
37            .for_each(|coeff: &mut F| *coeff *= &elm);
38        result
39    }
40
41    fn shiftr(&self, size: usize) -> Self {
42        let mut result = vec![F::zero(); size];
43        result.extend(self.coeffs.clone());
44        Self::from_coefficients_vec(result)
45    }
46
47    fn eval_polynomial(coeffs: &[F], x: F) -> F {
48        Self::from_coefficients_slice(coeffs).evaluate(&x)
49    }
50
51    fn to_chunked_polynomial(&self, num_chunks: usize, chunk_size: usize) -> ChunkedPolynomial<F> {
52        let mut chunk_polys: Vec<Self> = Vec::with_capacity(num_chunks);
53        for chunk in self.coeffs.chunks(chunk_size) {
54            chunk_polys.push(Self::from_coefficients_slice(chunk));
55        }
56
57        // Pad unused chunks with zeros.
58        for _ in chunk_polys.len()..num_chunks {
59            chunk_polys.push(Self::from_coefficients_vec(vec![]));
60        }
61
62        // Ensuring that the number of chunks is the one requested, following
63        // trait documentation
64        assert_eq!(chunk_polys.len(), num_chunks);
65
66        ChunkedPolynomial {
67            polys: chunk_polys,
68            size: chunk_size,
69        }
70    }
71}