From 88977526e08fd6e9f33bfd8fb9025704d5727a4c Mon Sep 17 00:00:00 2001 From: oskarth Date: Sun, 6 Aug 2023 10:25:36 +0100 Subject: [PATCH 1/2] feat(spartan/polynomial): Update polynomial with docs and new methods - Add documentation to spartan/polynomial.rs file - Make polynomial crate public to easily generate docs and encourage API use - Add Add trait and scalar_mul method - Add tests for polynomial functions Co-authored-by: wangtsiao Co-authored-by: CPerezz --- src/spartan/mod.rs | 6 +- src/spartan/polynomial.rs | 242 +++++++++++++++++++++++++++++++++++--- 2 files changed, 230 insertions(+), 18 deletions(-) diff --git a/src/spartan/mod.rs b/src/spartan/mod.rs index a8ef0223..282926a6 100644 --- a/src/spartan/mod.rs +++ b/src/spartan/mod.rs @@ -3,9 +3,11 @@ //! We provide two implementations, one in snark.rs (which does not use any preprocessing) //! and another in ppsnark.rs (which uses preprocessing to keep the verifier's state small if the PCS scheme provides a succinct verifier) //! We also provide direct.rs that allows proving a step circuit directly with either of the two SNARKs. +//! +//! In polynomial.rs we also provide foundational types and functions for manipulating multilinear polynomials. pub mod direct; -mod math; -pub(crate) mod polynomial; +pub(crate) mod math; +pub mod polynomial; pub mod ppsnark; pub mod snark; mod sumcheck; diff --git a/src/spartan/polynomial.rs b/src/spartan/polynomial.rs index 3aaa075d..dd23a455 100644 --- a/src/spartan/polynomial.rs +++ b/src/spartan/polynomial.rs @@ -1,20 +1,46 @@ -//! This module defines basic types related to polynomials +//! This module provides foundational types and functions for manipulating multilinear polynomials in the context of cryptographic computations. +//! +//! Main components: +//! - `EqPolynomial`: Represents multilinear extension of equality polynomials, evaluated based on binary input values. +//! - `MultilinearPolynomial`: Dense representation of multilinear polynomials, represented by evaluations over all possible binary inputs. +//! - `SparsePolynomial`: Efficient representation of sparse multilinear polynomials, storing only non-zero evaluations. use core::ops::Index; use ff::PrimeField; use rayon::prelude::*; use serde::{Deserialize, Serialize}; +use std::ops::Add; -pub(crate) struct EqPolynomial { +use crate::spartan::math::Math; + +/// Represents the multilinear extension polynomial (MLE) of the equality polynomial $eq(x,e)$, denoted as $\tilde{eq}(x, e)$. +/// +/// The polynomial is defined by the formula: +/// $$ +/// \tilde{eq}(x, e) = \prod_{i=0}^m(e_i * x_i + (1 - e_i) * (1 - x_i)) +/// $$ +/// +/// Each element in the vector `r` corresponds to a component $e_i$, representing a bit from the binary representation of an input value $e$. +/// This polynomial evaluates to 1 if every component $x_i$ equals its corresponding $e_i$, and 0 otherwise. +/// +/// For instance, for e = 6 (with a binary representation of 0b110), the vector r would be [1, 1, 0]. +pub struct EqPolynomial { r: Vec, } impl EqPolynomial { - /// Creates a new polynomial from its succinct specification + /// Creates a new `EqPolynomial` from a vector of Scalars `r`. + /// + /// Each Scalar in `r` corresponds to a bit from the binary representation of an input value `e`. pub const fn new(r: Vec) -> Self { EqPolynomial { r } } - /// Evaluates the polynomial at the specified point + /// Evaluates the `EqPolynomial` at a given point `rx`. + /// + /// This function computes the value of the polynomial at the point specified by `rx`. + /// It expects `rx` to have the same length as the internal vector `r`. + /// + /// Panics if `rx` and `r` have different lengths. pub fn evaluate(&self, rx: &[Scalar]) -> Scalar { assert_eq!(self.r.len(), rx.len()); (0..rx.len()) @@ -22,6 +48,9 @@ impl EqPolynomial { .fold(Scalar::ONE, |acc, item| acc * item) } + /// Evaluates the `EqPolynomial` at all the `2^|r|` points in its domain. + /// + /// Returns a vector of Scalars, each corresponding to the polynomial evaluation at a specific point. pub fn evals(&self) -> Vec { let ell = self.r.len(); let mut evals: Vec = vec![Scalar::ZERO; (2_usize).pow(ell as u32)]; @@ -42,17 +71,37 @@ impl EqPolynomial { size *= 2; } + evals } } -#[derive(Debug, Clone, Serialize, Deserialize)] +/// A multilinear extension of a polynomial $Z(\cdot)$, denote it as $\tilde{Z}(x_1, ..., x_m)$ +/// where the degree of each variable is at most one. +/// +/// This is the dense representation of a multilinear poynomial. +/// Let it be $\mathbb{G}(\cdot): \mathbb{F}^m \rightarrow \mathbb{F}$, it can be represented uniquely by the list of +/// evaluations of $\mathbb{G}(\cdot)$ over the Boolean hypercube $\{0, 1\}^m$. +/// +/// For example, a 3 variables multilinear polynomial can be represented by evaluation +/// at points $[0, 2^3-1]$. +/// +/// The implementation follows +/// $$ +/// \tilde{Z}(x_1, ..., x_m) = \sum_{e\in {0,1}^m}Z(e)\cdot \prod_{i=0}^m(x_i\cdot e_i)\cdot (1-e_i) +/// $$ +/// +/// Vector $Z$ indicates $Z(e)$ where $e$ ranges from $0$ to $2^m-1$. +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)] pub struct MultilinearPolynomial { - num_vars: usize, // the number of variables in the multilinear polynomial - Z: Vec, // evaluations of the polynomial in all the 2^num_vars Boolean inputs + num_vars: usize, // the number of variables in the multilinear polynomial + pub(crate) Z: Vec, // evaluations of the polynomial in all the 2^num_vars Boolean inputs } impl MultilinearPolynomial { + /// Creates a new MultilinearPolynomial from the given evaluations. + /// + /// The number of evaluations must be a power of two. pub fn new(Z: Vec) -> Self { assert_eq!(Z.len(), (2_usize).pow((Z.len() as f64).log2() as u32)); MultilinearPolynomial { @@ -61,14 +110,26 @@ impl MultilinearPolynomial { } } + /// Returns the number of variables in the multilinear polynomial pub const fn get_num_vars(&self) -> usize { self.num_vars } + /// Returns the total number of evaluations. pub fn len(&self) -> usize { self.Z.len() } + /// Checks if the multilinear polynomial is empty. + /// + /// This method returns true if the polynomial has no evaluations, and false otherwise. + pub fn is_empty(&self) -> bool { + self.Z.is_empty() + } + + /// Bounds the polynomial's top variable using the given scalar. + /// + /// This operation modifies the polynomial in-place. pub fn bound_poly_var_top(&mut self, r: &Scalar) { let n = self.len() / 2; @@ -86,7 +147,10 @@ impl MultilinearPolynomial { self.num_vars -= 1; } - // returns Z(r) in O(n) time + /// Evaluates the polynomial at the given point. + /// Returns Z(r) in O(n) time. + /// + /// The point must have a value for each variable. pub fn evaluate(&self, r: &[Scalar]) -> Scalar { // r must have a value for each variable assert_eq!(r.len(), self.get_num_vars()); @@ -99,6 +163,7 @@ impl MultilinearPolynomial { .sum() } + /// Evaluates the polynomial with the given evaluations and point. pub fn evaluate_with(Z: &[Scalar], r: &[Scalar]) -> Scalar { EqPolynomial::new(r.to_vec()) .evals() @@ -107,6 +172,15 @@ impl MultilinearPolynomial { .map(|(a, b)| a * b) .sum() } + + /// Multiplies the polynomial by a scalar. + pub fn scalar_mul(&self, scalar: &Scalar) -> Self { + let mut new_poly = self.clone(); + for z in &mut new_poly.Z { + *z *= scalar; + } + new_poly + } } impl Index for MultilinearPolynomial { @@ -118,6 +192,12 @@ impl Index for MultilinearPolynomial { } } +/// Sparse multilinear polynomial, which means the $Z(\cdot)$ is zero at most points. +/// So we do not have to store every evaluations of $Z(\cdot)$, only store the non-zero points. +/// +/// For example, the evaluations are [0, 0, 0, 1, 0, 1, 0, 2]. +/// The sparse polynomial only store the non-zero values, [(3, 1), (5, 1), (7, 2)]. +/// In the tuple, the first is index, the second is value. pub(crate) struct SparsePolynomial { num_vars: usize, Z: Vec<(usize, Scalar)>, @@ -128,6 +208,8 @@ impl SparsePolynomial { SparsePolynomial { num_vars, Z } } + /// Computes the $\tilde{eq}$ extension polynomial. + /// return 1 when a == r, otherwise return 0. fn compute_chi(a: &[bool], r: &[Scalar]) -> Scalar { assert_eq!(a.len(), r.len()); let mut chi_i = Scalar::ONE; @@ -145,19 +227,147 @@ impl SparsePolynomial { pub fn evaluate(&self, r: &[Scalar]) -> Scalar { assert_eq!(self.num_vars, r.len()); - let get_bits = |num: usize, num_bits: usize| -> Vec { - (0..num_bits) - .into_par_iter() - .map(|shift_amount| ((num & (1 << (num_bits - shift_amount - 1))) > 0)) - .collect::>() - }; - (0..self.Z.len()) .into_par_iter() .map(|i| { - let bits = get_bits(self.Z[i].0, r.len()); + let bits = (self.Z[i].0).get_bits(r.len()); SparsePolynomial::compute_chi(&bits, r) * self.Z[i].1 }) .sum() } } + +/// Adds another multilinear polynomial to `self`. +/// Assumes the two polynomials have the same number of variables. +impl Add for MultilinearPolynomial { + type Output = Result; + + fn add(self, other: Self) -> Self::Output { + if self.get_num_vars() != other.get_num_vars() { + return Err("The two polynomials must have the same number of variables"); + } + + let sum: Vec = self + .Z + .iter() + .zip(other.Z.iter()) + .map(|(a, b)| *a + *b) + .collect(); + + Ok(MultilinearPolynomial::new(sum)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use pasta_curves::Fp; + + fn make_mlp(len: usize, value: F) -> MultilinearPolynomial { + MultilinearPolynomial { + num_vars: len.count_ones() as usize, + Z: vec![value; len], + } + } + + fn test_eq_polynomial_with() { + let eq_poly = EqPolynomial::::new(vec![F::ONE, F::ZERO, F::ONE]); + let y = eq_poly.evaluate(vec![F::ONE, F::ONE, F::ONE].as_slice()); + assert_eq!(y, F::ZERO); + + let y = eq_poly.evaluate(vec![F::ONE, F::ZERO, F::ONE].as_slice()); + assert_eq!(y, F::ONE); + + let eval_list = eq_poly.evals(); + for (i, &coeff) in eval_list.iter().enumerate().take((2_usize).pow(3)) { + if i == 5 { + assert_eq!(coeff, F::ONE); + } else { + assert_eq!(coeff, F::ZERO); + } + } + } + + fn test_multilinear_polynomial_with() { + // Let the polynomial has 3 variables, p(x_1, x_2, x_3) = (x_1 + x_2) * x_3 + // Evaluations of the polynomial at boolean cube are [0, 0, 0, 1, 0, 1, 0, 2]. + + let TWO = F::from(2); + + let Z = vec![ + F::ZERO, + F::ZERO, + F::ZERO, + F::ONE, + F::ZERO, + F::ONE, + F::ZERO, + TWO, + ]; + let m_poly = MultilinearPolynomial::::new(Z.clone()); + assert_eq!(m_poly.get_num_vars(), 3); + + let x = vec![F::ONE, F::ONE, F::ONE]; + assert_eq!(m_poly.evaluate(x.as_slice()), TWO); + + let y = MultilinearPolynomial::::evaluate_with(Z.as_slice(), x.as_slice()); + assert_eq!(y, TWO); + } + + fn test_sparse_polynomial_with() { + // Let the polynomial have 3 variables, p(x_1, x_2, x_3) = (x_1 + x_2) * x_3 + // Evaluations of the polynomial at boolean cube are [0, 0, 0, 1, 0, 1, 0, 2]. + + let TWO = F::from(2); + let Z = vec![(3, F::ONE), (5, F::ONE), (7, TWO)]; + let m_poly = SparsePolynomial::::new(3, Z); + + let x = vec![F::ONE, F::ONE, F::ONE]; + assert_eq!(m_poly.evaluate(x.as_slice()), TWO); + + let x = vec![F::ONE, F::ZERO, F::ONE]; + assert_eq!(m_poly.evaluate(x.as_slice()), F::ONE); + } + + #[test] + fn test_eq_polynomial() { + test_eq_polynomial_with::(); + } + + #[test] + fn test_multilinear_polynomial() { + test_multilinear_polynomial_with::(); + } + + #[test] + fn test_sparse_polynomial() { + test_sparse_polynomial_with::(); + } + + fn test_mlp_add_with() { + let mlp1 = make_mlp(4, F::from(3)); + let mlp2 = make_mlp(4, F::from(7)); + + let mlp3 = mlp1.add(mlp2).unwrap(); + + assert_eq!(mlp3.Z, vec![F::from(10); 4]); + } + + fn test_mlp_scalar_mul_with() { + let mlp = make_mlp(4, F::from(3)); + + let mlp2 = mlp.scalar_mul(&F::from(2)); + + assert_eq!(mlp2.Z, vec![F::from(6); 4]); + } + + #[test] + fn test_mlp_add() { + test_mlp_add_with::(); + } + + #[test] + fn test_mlp_scalar_mul() { + test_mlp_scalar_mul_with::(); + } +} From db47a30815bd38955116bf4526660f51ed8e1e39 Mon Sep 17 00:00:00 2001 From: oskarth Date: Sun, 6 Aug 2023 12:03:18 +0100 Subject: [PATCH 2/2] docs: Wrap link in <> Otherwise `cargo doc` doesn't run. --- src/gadgets/ecc.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/gadgets/ecc.rs b/src/gadgets/ecc.rs index b1eb2767..f6c2d1dd 100644 --- a/src/gadgets/ecc.rs +++ b/src/gadgets/ecc.rs @@ -424,7 +424,7 @@ where } /// A gadget for scalar multiplication, optimized to use incomplete addition law. - /// The optimization here is analogous to https://github.com/arkworks-rs/r1cs-std/blob/6d64f379a27011b3629cf4c9cb38b7b7b695d5a0/src/groups/curves/short_weierstrass/mod.rs#L295, + /// The optimization here is analogous to , /// except we use complete addition law over affine coordinates instead of projective coordinates for the tail bits pub fn scalar_mul>( &self,