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