Skip to main content

o1_utils/
dense_polynomial.rs

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