o1_utils/
evaluations.rs

1//! This adds a few utility functions for the [Evaluations] arkworks type.
2
3use ark_ff::FftField;
4use ark_poly::{Evaluations, Radix2EvaluationDomain};
5use rayon::prelude::*;
6
7/// An extension for the [Evaluations] type.
8pub trait ExtendedEvaluations<F: FftField> {
9    /// This function "scales" (multiplies) a polynomial with a scalar
10    /// It is implemented to have the desired functionality for DensePolynomial
11    fn scale(&self, elm: F) -> Self;
12
13    /// Square each evaluation
14    fn square(&self) -> Self;
15
16    /// Raise each evaluation to some power `pow`
17    fn pow(&self, pow: usize) -> Self;
18
19    /// Utility function for shifting poly along domain coordinate
20    fn shift(&self, len: usize) -> Self;
21}
22
23impl<F: FftField> ExtendedEvaluations<F> for Evaluations<F, Radix2EvaluationDomain<F>> {
24    fn scale(&self, elm: F) -> Self {
25        let mut result = self.clone();
26        result.evals.par_iter_mut().for_each(|coeff| *coeff *= &elm);
27        result
28    }
29
30    fn square(&self) -> Self {
31        let mut result = self.clone();
32        result.evals.par_iter_mut().for_each(|e| {
33            let _ = e.square_in_place();
34        });
35        result
36    }
37
38    fn pow(&self, pow: usize) -> Self {
39        let mut result = self.clone();
40        result
41            .evals
42            .par_iter_mut()
43            .for_each(|e| *e = e.pow([pow as u64]));
44        result
45    }
46
47    fn shift(&self, len: usize) -> Self {
48        let len_new = len % self.evals.len();
49        let mut result = Self::from_vec_and_domain(Vec::with_capacity(len), self.domain());
50        result.evals.extend_from_slice(&self.evals[len_new..]);
51        result.evals.extend_from_slice(&self.evals[0..len_new]);
52        result
53    }
54}