1//! This module defines basic types related to polynomials
  2use core::ops::Index;
  3use ff::PrimeField;
  4use rayon::prelude::*;
  5use serde::{Deserialize, Serialize};
  6
  7pub(crate) struct EqPolynomial<Scalar: PrimeField> {
  8  r: Vec<Scalar>,
  9}
 10
 11impl<Scalar: PrimeField> EqPolynomial<Scalar> {
 12  /// Creates a new polynomial from its succinct specification
 13  pub fn new(r: Vec<Scalar>) -> Self {
 14    EqPolynomial { r }
 15  }
 16
 17  /// Evaluates the polynomial at the specified point
 18  pub fn evaluate(&self, rx: &[Scalar]) -> Scalar {
 19    assert_eq!(self.r.len(), rx.len());
 20    (0..rx.len())
 21      .map(|i| rx[i] * self.r[i] + (Scalar::ONE - rx[i]) * (Scalar::ONE - self.r[i]))
 22      .fold(Scalar::ONE, |acc, item| acc * item)
 23  }
 24
 25  pub fn evals(&self) -> Vec<Scalar> {
 26    let ell = self.r.len();
 27    let mut evals: Vec<Scalar> = vec![Scalar::ZERO; (2_usize).pow(ell as u32)];
 28    let mut size = 1;
 29    evals[0] = Scalar::ONE;
 30
 31    for r in self.r.iter().rev() {
 32      let (evals_left, evals_right) = evals.split_at_mut(size);
 33      let (evals_right, _) = evals_right.split_at_mut(size);
 34
 35      evals_left
 36        .par_iter_mut()
 37        .zip(evals_right.par_iter_mut())
 38        .for_each(|(x, y)| {
 39          *y = *x * r;
 40          *x -= &*y;
 41        });
 42
 43      size *= 2;
 44    }
 45    evals
 46  }
 47}
 48
 49#[derive(Debug, Clone, Serialize, Deserialize)]
 50pub struct MultilinearPolynomial<Scalar: PrimeField> {
 51  num_vars: usize, // the number of variables in the multilinear polynomial
 52  Z: Vec<Scalar>,  // evaluations of the polynomial in all the 2^num_vars Boolean inputs
 53}
 54
 55impl<Scalar: PrimeField> MultilinearPolynomial<Scalar> {
 56  pub fn new(Z: Vec<Scalar>) -> Self {
 57    assert_eq!(Z.len(), (2_usize).pow((Z.len() as f64).log2() as u32));
 58    MultilinearPolynomial {
 59      num_vars: (Z.len() as f64).log2() as usize,
 60      Z,
 61    }
 62  }
 63
 64  pub fn get_num_vars(&self) -> usize {
 65    self.num_vars
 66  }
 67
 68  pub fn len(&self) -> usize {
 69    self.Z.len()
 70  }
 71
 72  pub fn bound_poly_var_top(&mut self, r: &Scalar) {
 73    let n = self.len() / 2;
 74
 75    let (left, right) = self.Z.split_at_mut(n);
 76    let (right, _) = right.split_at(n);
 77
 78    left
 79      .par_iter_mut()
 80      .zip(right.par_iter())
 81      .for_each(|(a, b)| {
 82        *a += *r * (*b - *a);
 83      });
 84
 85    self.Z.resize(n, Scalar::ZERO);
 86    self.num_vars -= 1;
 87  }
 88
 89  // returns Z(r) in O(n) time
 90  pub fn evaluate(&self, r: &[Scalar]) -> Scalar {
 91    // r must have a value for each variable
 92    assert_eq!(r.len(), self.get_num_vars());
 93    let chis = EqPolynomial::new(r.to_vec()).evals();
 94    assert_eq!(chis.len(), self.Z.len());
 95
 96    (0..chis.len())
 97      .into_par_iter()
 98      .map(|i| chis[i] * self.Z[i])
 99      .reduce(|| Scalar::ZERO, |x, y| x + y)
100  }
101
102  pub fn evaluate_with(Z: &[Scalar], r: &[Scalar]) -> Scalar {
103    EqPolynomial::new(r.to_vec())
104      .evals()
105      .into_par_iter()
106      .zip(Z.into_par_iter())
107      .map(|(a, b)| a * b)
108      .reduce(|| Scalar::ZERO, |x, y| x + y)
109  }
110}
111
112impl<Scalar: PrimeField> Index<usize> for MultilinearPolynomial<Scalar> {
113  type Output = Scalar;
114
115  #[inline(always)]
116  fn index(&self, _index: usize) -> &Scalar {
117    &(self.Z[_index])
118  }
119}
120
121pub(crate) struct SparsePolynomial<Scalar: PrimeField> {
122  num_vars: usize,
123  Z: Vec<(usize, Scalar)>,
124}
125
126impl<Scalar: PrimeField> SparsePolynomial<Scalar> {
127  pub fn new(num_vars: usize, Z: Vec<(usize, Scalar)>) -> Self {
128    SparsePolynomial { num_vars, Z }
129  }
130
131  fn compute_chi(a: &[bool], r: &[Scalar]) -> Scalar {
132    assert_eq!(a.len(), r.len());
133    let mut chi_i = Scalar::ONE;
134    for j in 0..r.len() {
135      if a[j] {
136        chi_i *= r[j];
137      } else {
138        chi_i *= Scalar::ONE - r[j];
139      }
140    }
141    chi_i
142  }
143
144  // Takes O(n log n)
145  pub fn evaluate(&self, r: &[Scalar]) -> Scalar {
146    assert_eq!(self.num_vars, r.len());
147
148    let get_bits = |num: usize, num_bits: usize| -> Vec<bool> {
149      (0..num_bits)
150        .into_par_iter()
151        .map(|shift_amount| ((num & (1 << (num_bits - shift_amount - 1))) > 0))
152        .collect::<Vec<bool>>()
153    };
154
155    (0..self.Z.len())
156      .into_par_iter()
157      .map(|i| {
158        let bits = get_bits(self.Z[i].0, r.len());
159        SparsePolynomial::compute_chi(&bits, r) * self.Z[i].1
160      })
161      .reduce(|| Scalar::ZERO, |x, y| x + y)
162  }
163}