Skip to main content

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    #[must_use]
12    fn scale(&self, elm: F) -> Self;
13
14    /// Square each evaluation
15    #[must_use]
16    fn square(&self) -> Self;
17
18    /// Raise each evaluation to some power `pow`
19    #[must_use]
20    fn pow(&self, pow: usize) -> Self;
21
22    /// Utility function for shifting poly along domain coordinate
23    #[must_use]
24    fn shift(&self, len: usize) -> Self;
25}
26
27impl<F: FftField> ExtendedEvaluations<F> for Evaluations<F, Radix2EvaluationDomain<F>> {
28    fn scale(&self, elm: F) -> Self {
29        let mut result = self.clone();
30        result.evals.par_iter_mut().for_each(|coeff| *coeff *= &elm);
31        result
32    }
33
34    fn square(&self) -> Self {
35        let mut result = self.clone();
36        result.evals.par_iter_mut().for_each(|e| {
37            let _ = e.square_in_place();
38        });
39        result
40    }
41
42    fn pow(&self, pow: usize) -> Self {
43        let mut result = self.clone();
44        result
45            .evals
46            .par_iter_mut()
47            .for_each(|e| *e = e.pow([pow as u64]));
48        result
49    }
50
51    fn shift(&self, len: usize) -> Self {
52        let len_new = len % self.evals.len();
53        let mut result = Self::from_vec_and_domain(Vec::with_capacity(len), self.domain());
54        result.evals.extend_from_slice(&self.evals[len_new..]);
55        result.evals.extend_from_slice(&self.evals[0..len_new]);
56        result
57    }
58}