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  pub fn split(&self, idx: usize) -> (Self, Self) {
112    assert!(idx < self.len());
113    (
114      Self::new(self.Z[..idx].to_vec()),
115      Self::new(self.Z[idx..2 * idx].to_vec()),
116    )
117  }
118}
119
120impl<Scalar: PrimeField> Index<usize> for MultilinearPolynomial<Scalar> {
121  type Output = Scalar;
122
123  #[inline(always)]
124  fn index(&self, _index: usize) -> &Scalar {
125    &(self.Z[_index])
126  }
127}
128
129pub(crate) struct SparsePolynomial<Scalar: PrimeField> {
130  num_vars: usize,
131  Z: Vec<(usize, Scalar)>,
132}
133
134impl<Scalar: PrimeField> SparsePolynomial<Scalar> {
135  pub fn new(num_vars: usize, Z: Vec<(usize, Scalar)>) -> Self {
136    SparsePolynomial { num_vars, Z }
137  }
138
139  fn compute_chi(a: &[bool], r: &[Scalar]) -> Scalar {
140    assert_eq!(a.len(), r.len());
141    let mut chi_i = Scalar::one();
142    for j in 0..r.len() {
143      if a[j] {
144        chi_i *= r[j];
145      } else {
146        chi_i *= Scalar::one() - r[j];
147      }
148    }
149    chi_i
150  }
151
152  // Takes O(n log n). TODO: do this in O(n) where n is the number of entries in Z
153  pub fn evaluate(&self, r: &[Scalar]) -> Scalar {
154    assert_eq!(self.num_vars, r.len());
155
156    let get_bits = |num: usize, num_bits: usize| -> Vec<bool> {
157      (0..num_bits)
158        .into_par_iter()
159        .map(|shift_amount| ((num & (1 << (num_bits - shift_amount - 1))) > 0))
160        .collect::<Vec<bool>>()
161    };
162
163    (0..self.Z.len())
164      .into_par_iter()
165      .map(|i| {
166        let bits = get_bits(self.Z[i].0, r.len());
167        SparsePolynomial::compute_chi(&bits, r) * self.Z[i].1
168      })
169      .reduce(Scalar::zero, |x, y| x + y)
170  }
171}