Skip to main content

o1_utils/
evaluations.rs

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