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    fn scale(&self, elm: F) -> Self;
17
18    /// Shifts all the coefficients to the right.
19    fn shiftr(&self, size: usize) -> Self;
20
21    /// `eval_polynomial(coeffs, x)` evaluates a polynomial given its coefficients `coeffs` and a point `x`.
22    fn eval_polynomial(coeffs: &[F], x: F) -> F;
23
24    /// Convert a polynomial into chunks.
25    /// Implementors must ensure that the result contains exactly num_chunks.
26    fn to_chunked_polynomial(&self, num_chunks: usize, size: usize) -> ChunkedPolynomial<F>;
27}
28
29impl<F: Field> ExtendedDensePolynomial<F> for DensePolynomial<F> {
30    fn scale(&self, elm: F) -> Self {
31        let mut result = self.clone();
32        result
33            .coeffs
34            .par_iter_mut()
35            .for_each(|coeff: &mut F| *coeff *= &elm);
36        result
37    }
38
39    fn shiftr(&self, size: usize) -> Self {
40        let mut result = vec![F::zero(); size];
41        result.extend(self.coeffs.clone());
42        DensePolynomial::<F>::from_coefficients_vec(result)
43    }
44
45    fn eval_polynomial(coeffs: &[F], x: F) -> F {
46        DensePolynomial::from_coefficients_slice(coeffs).evaluate(&x)
47    }
48
49    fn to_chunked_polynomial(&self, num_chunks: usize, chunk_size: usize) -> ChunkedPolynomial<F> {
50        let mut chunk_polys: Vec<DensePolynomial<F>> = Vec::with_capacity(num_chunks);
51        for chunk in self.coeffs.chunks(chunk_size) {
52            chunk_polys.push(DensePolynomial::from_coefficients_slice(chunk));
53        }
54
55        // Pad unused chunks with zeros.
56        for _ in chunk_polys.len()..num_chunks {
57            chunk_polys.push(DensePolynomial::from_coefficients_vec(vec![]));
58        }
59
60        // Ensuring that the number of chunks is the one requested, following
61        // trait documentation
62        assert_eq!(chunk_polys.len(), num_chunks);
63
64        ChunkedPolynomial {
65            polys: chunk_polys,
66            size: chunk_size,
67        }
68    }
69}