diff --git a/CHANGELOG.md b/CHANGELOG.md index 0cb7f6e..5a2f82c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,6 +7,8 @@ You can update downstream usage with `grep -rl 'AllocatedBit' . | xargs env LANG ### Features +- [\#53](https://github.com/arkworks-rs/r1cs-std/pull/53) Add univariate evaluation domain and lagrange interpolation. + ### Improvements ### Bug Fixes diff --git a/Cargo.toml b/Cargo.toml index 0d7206f..37ca6a0 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -36,4 +36,4 @@ ark-poly = { version = "^0.2.0", default-features = false } [features] default = ["std"] std = [ "ark-ff/std", "ark-relations/std", "ark-std/std", "num-bigint/std" ] -parallel = [ "std", "ark-ff/parallel" ] +parallel = [ "std", "ark-ff/parallel", "ark-std/parallel"] diff --git a/src/poly/domain/mod.rs b/src/poly/domain/mod.rs new file mode 100644 index 0000000..35e0189 --- /dev/null +++ b/src/poly/domain/mod.rs @@ -0,0 +1,77 @@ +use crate::boolean::Boolean; +use crate::fields::fp::FpVar; +use crate::fields::FieldVar; +use ark_ff::PrimeField; +use ark_relations::r1cs::SynthesisError; +use ark_std::vec::Vec; + +pub mod vanishing_poly; + +#[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)] +/// Defines an evaluation domain over a prime field. The domain is a coset of size `1< { + /// generator of subgroup g + pub gen: F, + /// index of the quotient group (i.e. the `offset`) + pub offset: F, + /// dimension of evaluation domain + pub dim: u64, +} + +impl Radix2Domain { + /// order of the domain + pub fn order(&self) -> usize { + 1 << self.dim + } + + /// Returns g, g^2, ..., g^{dim} + fn powers_of_gen(&self, dim: usize) -> Vec { + let mut result = Vec::new(); + let mut cur = self.gen; + for _ in 0..dim { + result.push(cur); + cur = cur * cur; + } + result + } + + /// For domain `h` with dimension `n`, `position` represented by `query_pos` in big endian form, + /// returns `h*g^{position}` + pub fn query_position_to_coset( + &self, + query_pos: &[Boolean], + coset_dim: u64, + ) -> Result>, SynthesisError> { + let mut coset_index = query_pos; + assert!( + query_pos.len() == self.dim as usize + || query_pos.len() == (self.dim - coset_dim) as usize + ); + if query_pos.len() == self.dim as usize { + coset_index = &coset_index[0..(coset_index.len() - coset_dim as usize)]; + } + let mut coset = Vec::new(); + let powers_of_g = &self.powers_of_gen(self.dim as usize)[(coset_dim as usize)..]; + + let mut first_point_in_coset: FpVar = FpVar::zero(); + for i in 0..coset_index.len() { + let term = coset_index[i].select(&FpVar::constant(powers_of_g[i]), &FpVar::zero())?; + first_point_in_coset += &term; + } + + first_point_in_coset *= &FpVar::Constant(self.offset); + + coset.push(first_point_in_coset); + for i in 1..(1 << (coset_dim as usize)) { + let new_elem = &coset[i - 1] * &FpVar::Constant(self.gen); + coset.push(new_elem); + } + + Ok(coset) + } +} diff --git a/src/poly/domain/vanishing_poly.rs b/src/poly/domain/vanishing_poly.rs new file mode 100644 index 0000000..8adce83 --- /dev/null +++ b/src/poly/domain/vanishing_poly.rs @@ -0,0 +1,79 @@ +use crate::fields::fp::FpVar; +use crate::fields::FieldVar; +use ark_ff::{Field, PrimeField}; +use ark_relations::r1cs::SynthesisError; +use ark_std::ops::Sub; + +/// Struct describing vanishing polynomial for a multiplicative coset H where |H| is a power of 2. +/// As H is a coset, every element can be described as h*g^i and therefore +/// has vanishing polynomial Z_H(x) = x^|H| - h^|H| +#[derive(Clone)] +pub struct VanishingPolynomial { + /// h^|H| + pub constant_term: F, + /// log_2(|H|) + pub dim_h: u64, + /// |H| + pub order_h: u64, +} + +impl VanishingPolynomial { + /// returns a VanishingPolynomial of coset `H = h`. + pub fn new(offset: F, dim_h: u64) -> Self { + let order_h = 1 << dim_h; + let vp = VanishingPolynomial { + constant_term: offset.pow([order_h]), + dim_h, + order_h, + }; + vp + } + + /// Evaluates the vanishing polynomial without generating the constraints. + pub fn evaluate(&self, x: &F) -> F { + let mut result = x.pow([self.order_h]); + result -= &self.constant_term; + result + } + + /// Evaluates the constraints and just gives you the gadget for the result. + /// Caution for use in holographic lincheck: The output has 2 entries in one matrix + pub fn evaluate_constraints(&self, x: &FpVar) -> Result, SynthesisError> { + if self.dim_h == 1 { + let result = x.sub(x); + return Ok(result); + } + + let mut cur = x.square()?; + for _ in 1..self.dim_h { + cur.square_in_place()?; + } + cur -= &FpVar::Constant(self.constant_term); + Ok(cur) + } +} + +#[cfg(test)] +mod tests { + use crate::alloc::AllocVar; + use crate::fields::fp::FpVar; + use crate::poly::domain::vanishing_poly::VanishingPolynomial; + use crate::R1CSVar; + use ark_relations::r1cs::ConstraintSystem; + use ark_std::{test_rng, UniformRand}; + use ark_test_curves::bls12_381::Fr; + + #[test] + fn constraints_test() { + let mut rng = test_rng(); + let offset = Fr::rand(&mut rng); + let cs = ConstraintSystem::new_ref(); + let x = Fr::rand(&mut rng); + let x_var = FpVar::new_witness(ns!(cs, "x_var"), || Ok(x)).unwrap(); + let vp = VanishingPolynomial::new(offset, 12); + let native = vp.evaluate(&x); + let result_var = vp.evaluate_constraints(&x_var).unwrap(); + assert!(cs.is_satisfied().unwrap()); + assert_eq!(result_var.value().unwrap(), native); + } +} diff --git a/src/poly/evaluations/mod.rs b/src/poly/evaluations/mod.rs new file mode 100644 index 0000000..becd6c4 --- /dev/null +++ b/src/poly/evaluations/mod.rs @@ -0,0 +1 @@ +pub mod univariate; diff --git a/src/poly/evaluations/univariate/lagrange_interpolator.rs b/src/poly/evaluations/univariate/lagrange_interpolator.rs new file mode 100644 index 0000000..ebfb9c4 --- /dev/null +++ b/src/poly/evaluations/univariate/lagrange_interpolator.rs @@ -0,0 +1,142 @@ +use crate::poly::domain::vanishing_poly::VanishingPolynomial; +use ark_ff::{batch_inversion, PrimeField}; +use ark_std::vec::Vec; +/// Struct describing Lagrange interpolation for a multiplicative coset I, +/// with |I| a power of 2. +/// TODO: Pull in lagrange poly explanation from libiop +#[derive(Clone)] +pub struct LagrangeInterpolator { + pub(crate) domain_order: usize, + pub(crate) all_domain_elems: Vec, + pub(crate) v_inv_elems: Vec, + pub(crate) domain_vp: VanishingPolynomial, + poly_evaluations: Vec, +} + +impl LagrangeInterpolator { + /// Returns a lagrange interpolator, given the domain specification. + pub fn new( + domain_offset: F, + domain_generator: F, + domain_dim: u64, + poly_evaluations: Vec, + ) -> Self { + let domain_order = 1 << domain_dim; + assert_eq!(poly_evaluations.len(), domain_order); + let mut cur_elem = domain_offset; + let mut all_domain_elems = vec![domain_offset]; + let mut v_inv_elems: Vec = Vec::new(); + // Cache all elements in the domain + for _ in 1..domain_order { + cur_elem *= domain_generator; + all_domain_elems.push(cur_elem); + } + /* + By computing the following elements as constants, + we can further reduce the interpolation costs. + + m = order of the interpolation domain + v_inv[i] = prod_{j != i} h(g^i - g^j) + We use the following facts to compute this: + v_inv[0] = m*h^{m-1} + v_inv[i] = g^{-1} * v_inv[i-1] + */ + // TODO: Include proof of the above two points + let g_inv = domain_generator.inverse().unwrap(); + let m = F::from((1 << domain_dim) as u128); + let mut v_inv_i = m * domain_offset.pow([(domain_order - 1) as u64]); + for _ in 0..domain_order { + v_inv_elems.push(v_inv_i); + v_inv_i *= g_inv; + } + + // TODO: Cache the intermediate terms with Z_H(x) evaluations. + let vp = VanishingPolynomial::new(domain_offset, domain_dim); + + let lagrange_interpolation: LagrangeInterpolator = LagrangeInterpolator { + domain_order, + all_domain_elems, + v_inv_elems, + domain_vp: vp, + poly_evaluations, + }; + lagrange_interpolation + } + + pub(crate) fn compute_lagrange_coefficients(&self, interpolation_point: F) -> Vec { + /* + * Let t be the interpolation point, H be the multiplicative coset, with elements of the form h*g^i. + Compute each L_{i,H}(t) as Z_{H}(t) * v_i / (t- h g^i) + where: + - Z_{H}(t) = \prod_{j} (t-h*g^j) = (t^m-h^m), and + - v_{i} = 1 / \prod_{j \neq i} h(g^i-g^j). + Below we use the fact that v_{0} = 1/(m * h^(m-1)) and v_{i+1} = g * v_{i}. + We compute the inverse of each coefficient, and then batch invert the entire result. + */ + let vp_t_inv = self + .domain_vp + .evaluate(&interpolation_point) + .inverse() + .unwrap(); + let mut inverted_lagrange_coeffs: Vec = Vec::with_capacity(self.all_domain_elems.len()); + for i in 0..self.domain_order { + let l = vp_t_inv * self.v_inv_elems[i]; + let r = self.all_domain_elems[i]; + inverted_lagrange_coeffs.push(l * (interpolation_point - r)); + } + let lagrange_coeffs = inverted_lagrange_coeffs.as_mut_slice(); + batch_inversion::(lagrange_coeffs); + lagrange_coeffs.iter().cloned().collect() + } + + pub fn interpolate(&self, interpolation_point: F) -> F { + let lagrange_coeffs = self.compute_lagrange_coefficients(interpolation_point); + let mut interpolation = F::zero(); + for i in 0..self.domain_order { + interpolation += lagrange_coeffs[i] * self.poly_evaluations[i]; + } + interpolation + } +} + +#[cfg(test)] +mod tests { + use crate::poly::domain::Radix2Domain; + use crate::poly::evaluations::univariate::lagrange_interpolator::LagrangeInterpolator; + use ark_ff::{FftField, Field, One}; + use ark_poly::univariate::DensePolynomial; + use ark_poly::{Polynomial, UVPolynomial}; + use ark_std::{test_rng, UniformRand}; + use ark_test_curves::bls12_381::Fr; + + #[test] + pub fn test_native_interpolate() { + let mut rng = test_rng(); + let poly = DensePolynomial::rand(15, &mut rng); + let gen = Fr::get_root_of_unity(1 << 4).unwrap(); + assert_eq!(gen.pow(&[1 << 4]), Fr::one()); + let domain = Radix2Domain { + gen, + offset: Fr::multiplicative_generator(), + dim: 4, // 2^4 = 16 + }; + // generate evaluations of `poly` on this domain + let mut coset_point = domain.offset; + let mut oracle_evals = Vec::new(); + for _ in 0..(1 << 4) { + oracle_evals.push(poly.evaluate(&coset_point)); + coset_point *= gen; + } + + let interpolator = + LagrangeInterpolator::new(domain.offset, domain.gen, domain.dim, oracle_evals); + + // the point to evaluate at + let interpolate_point = Fr::rand(&mut rng); + + let expected = poly.evaluate(&interpolate_point); + let actual = interpolator.interpolate(interpolate_point); + + assert_eq!(actual, expected) + } +} diff --git a/src/poly/evaluations/univariate/mod.rs b/src/poly/evaluations/univariate/mod.rs new file mode 100644 index 0000000..32f67bc --- /dev/null +++ b/src/poly/evaluations/univariate/mod.rs @@ -0,0 +1,320 @@ +pub mod lagrange_interpolator; + +use crate::alloc::AllocVar; +use crate::fields::fp::FpVar; +use crate::fields::FieldVar; +use crate::poly::domain::Radix2Domain; +use crate::poly::evaluations::univariate::lagrange_interpolator::LagrangeInterpolator; +use crate::R1CSVar; +use ark_ff::{batch_inversion, PrimeField}; +use ark_relations::r1cs::SynthesisError; +use ark_std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Sub, SubAssign}; +use ark_std::vec::Vec; + +#[derive(Clone)] +/// Stores a UV polynomial in evaluation form. +pub struct EvaluationsVar { + /// Evaluations of univariate polynomial over domain + pub evals: Vec>, + /// Optional Lagrange Interpolator. Useful for lagrange interpolation. + pub lagrange_interpolator: Option>, + domain: Radix2Domain, +} + +impl EvaluationsVar { + /// Construct `Self` from evaluations and a domain. + /// `interpolate` indicates if user wants to interpolate this polynomial + /// using lagrange interpolation. + pub fn from_vec_and_domain( + evaluations: Vec>, + domain: Radix2Domain, + interpolate: bool, + ) -> Self { + assert_eq!( + evaluations.len(), + 1 << domain.dim, + "evaluations and domain has different dimensions" + ); + + let mut ev = Self { + evals: evaluations, + lagrange_interpolator: None, + domain, + }; + if interpolate { + ev.generate_lagrange_interpolator(); + } + ev + } + + /// Generate lagrange interpolator and mark it ready to interpolate + pub fn generate_lagrange_interpolator(&mut self) { + let poly_evaluations_val: Vec<_> = self.evals.iter().map(|v| v.value().unwrap()).collect(); + let domain = &self.domain; + let lagrange_interpolator = + LagrangeInterpolator::new(domain.offset, domain.gen, domain.dim, poly_evaluations_val); + self.lagrange_interpolator = Some(lagrange_interpolator) + } + + fn compute_lagrange_coefficients( + &self, + interpolation_point: &FpVar, + ) -> Result>, SynthesisError> { + // ref: https://github.com/alexchmit/perfect-constraints/blob/79692f2652a95a57f2c7187f5b5276345e680230/fractal/src/algebra/lagrange_interpolation.rs#L159 + let cs = interpolation_point.cs(); + let t = interpolation_point; + let lagrange_interpolator = self + .lagrange_interpolator + .as_ref() + .expect("lagrange interpolator has not been initialized. \ + Call `self.generate_lagrange_interpolator` first or set `interpolate` to true in constructor. "); + let lagrange_coeffs = + lagrange_interpolator.compute_lagrange_coefficients(t.value().unwrap()); + let mut lagrange_coeffs_fg = Vec::new(); + // Now we convert these lagrange coefficients to gadgets, and then constrain them. + // The i-th lagrange coefficients constraint is: + // (v_inv[i] * t - v_inv[i] * domain_elem[i]) * (coeff) = 1/Z_I(t) + let vp_t = lagrange_interpolator.domain_vp.evaluate_constraints(t)?; + // let inv_vp_t = vp_t.inverse()?; + for i in 0..lagrange_interpolator.domain_order { + let constant: F = + (-lagrange_interpolator.all_domain_elems[i]) * lagrange_interpolator.v_inv_elems[i]; + let mut a_element: FpVar = + t * &FpVar::constant(lagrange_interpolator.v_inv_elems[i]); + a_element += FpVar::constant(constant); + + let lag_coeff: FpVar = + FpVar::new_witness(ns!(cs, "generate lagrange coefficient"), || { + Ok(lagrange_coeffs[i]) + })?; + // Enforce the actual constraint (A_element) * (lagrange_coeff) = 1/Z_I(t) + assert_eq!( + (lagrange_interpolator.v_inv_elems[i] * t.value().unwrap() + - lagrange_interpolator.v_inv_elems[i] + * lagrange_interpolator.all_domain_elems[i]) + * lagrange_coeffs[i], + vp_t.value().unwrap() + ); + a_element.mul_equals(&lag_coeff, &vp_t)?; + lagrange_coeffs_fg.push(lag_coeff); + } + Ok(lagrange_coeffs_fg) + } + + /// Returns constraints for Interpolating and evaluating at `interpolation_point` + pub fn interpolate_and_evaluate( + &self, + interpolation_point: &FpVar, + ) -> Result, SynthesisError> { + let lagrange_interpolator = self + .lagrange_interpolator + .as_ref() + .expect("lagrange interpolator has not been initialized. "); + let lagrange_coeffs = self.compute_lagrange_coefficients(interpolation_point)?; + let mut interpolation: FpVar = FpVar::zero(); + for i in 0..lagrange_interpolator.domain_order { + let intermediate = &lagrange_coeffs[i] * &self.evals[i]; + interpolation += &intermediate + } + + Ok(interpolation) + } +} + +impl<'a, 'b, F: PrimeField> Add<&'a EvaluationsVar> for &'b EvaluationsVar { + type Output = EvaluationsVar; + + fn add(self, rhs: &'a EvaluationsVar) -> Self::Output { + let mut result = self.clone(); + result += rhs; + result + } +} + +impl<'a, F: PrimeField> AddAssign<&'a EvaluationsVar> for EvaluationsVar { + fn add_assign(&mut self, other: &'a EvaluationsVar) { + assert_eq!(self.domain, other.domain, "domains are unequal"); + + self.lagrange_interpolator = None; + self.evals + .iter_mut() + .zip(&other.evals) + .for_each(|(a, b)| *a = &*a + b) + } +} + +impl<'a, 'b, F: PrimeField> Sub<&'a EvaluationsVar> for &'b EvaluationsVar { + type Output = EvaluationsVar; + + fn sub(self, rhs: &'a EvaluationsVar) -> Self::Output { + let mut result = self.clone(); + result -= rhs; + result + } +} + +impl<'a, F: PrimeField> SubAssign<&'a EvaluationsVar> for EvaluationsVar { + fn sub_assign(&mut self, other: &'a EvaluationsVar) { + assert_eq!(self.domain, other.domain, "domains are unequal"); + + self.lagrange_interpolator = None; + self.evals + .iter_mut() + .zip(&other.evals) + .for_each(|(a, b)| *a = &*a - b) + } +} + +impl<'a, 'b, F: PrimeField> Mul<&'a EvaluationsVar> for &'b EvaluationsVar { + type Output = EvaluationsVar; + + fn mul(self, rhs: &'a EvaluationsVar) -> Self::Output { + let mut result = self.clone(); + result *= rhs; + result + } +} + +impl<'a, F: PrimeField> MulAssign<&'a EvaluationsVar> for EvaluationsVar { + fn mul_assign(&mut self, other: &'a EvaluationsVar) { + assert_eq!(self.domain, other.domain, "domains are unequal"); + + self.lagrange_interpolator = None; + self.evals + .iter_mut() + .zip(&other.evals) + .for_each(|(a, b)| *a = &*a * b) + } +} + +impl<'a, 'b, F: PrimeField> Div<&'a EvaluationsVar> for &'b EvaluationsVar { + type Output = EvaluationsVar; + + fn div(self, rhs: &'a EvaluationsVar) -> Self::Output { + let mut result = self.clone(); + result /= rhs; + result + } +} + +impl<'a, F: PrimeField> DivAssign<&'a EvaluationsVar> for EvaluationsVar { + fn div_assign(&mut self, other: &'a EvaluationsVar) { + assert_eq!(self.domain, other.domain, "domains are unequal"); + let cs = self.evals[0].cs(); + // the prover can generate result = (1 / other) * self offline + let mut result_val: Vec<_> = other.evals.iter().map(|x| x.value().unwrap()).collect(); + batch_inversion(&mut result_val); + result_val + .iter_mut() + .zip(&self.evals) + .for_each(|(a, self_var)| *a *= self_var.value().unwrap()); + let result_var: Vec<_> = result_val + .iter() + .map(|x| FpVar::new_witness(ns!(cs, "div result"), || Ok(*x)).unwrap()) + .collect(); + // enforce constraint + for i in 0..result_var.len() { + result_var[i] + .mul_equals(&other.evals[i], &self.evals[i]) + .unwrap(); + } + + self.lagrange_interpolator = None; + self.evals = result_var + } +} + +#[cfg(test)] +mod tests { + use crate::alloc::AllocVar; + use crate::fields::fp::FpVar; + use crate::poly::domain::Radix2Domain; + use crate::poly::evaluations::univariate::EvaluationsVar; + use crate::R1CSVar; + use ark_ff::{FftField, Field, One, UniformRand}; + use ark_poly::polynomial::univariate::DensePolynomial; + use ark_poly::{Polynomial, UVPolynomial}; + use ark_relations::r1cs::ConstraintSystem; + use ark_std::test_rng; + use ark_test_curves::bls12_381::Fr; + + #[test] + fn test_interpolate() { + let mut rng = test_rng(); + let poly = DensePolynomial::rand(15, &mut rng); + let gen = Fr::get_root_of_unity(1 << 4).unwrap(); + assert_eq!(gen.pow(&[1 << 4]), Fr::one()); + let domain = Radix2Domain { + gen, + offset: Fr::multiplicative_generator(), + dim: 4, // 2^4 = 16 + }; + let mut coset_point = domain.offset; + let mut oracle_evals = Vec::new(); + for _ in 0..(1 << 4) { + oracle_evals.push(poly.evaluate(&coset_point)); + coset_point *= gen; + } + let cs = ConstraintSystem::new_ref(); + let evaluations_fp: Vec<_> = oracle_evals + .iter() + .map(|x| FpVar::new_input(ns!(cs, "evaluations"), || Ok(x)).unwrap()) + .collect(); + let evaluations_var = EvaluationsVar::from_vec_and_domain(evaluations_fp, domain, true); + + let interpolate_point = Fr::rand(&mut rng); + let interpolate_point_fp = + FpVar::new_input(ns!(cs, "interpolate point"), || Ok(interpolate_point)).unwrap(); + + let expected = poly.evaluate(&interpolate_point); + + let actual = evaluations_var + .interpolate_and_evaluate(&interpolate_point_fp) + .unwrap() + .value() + .unwrap(); + + assert_eq!(actual, expected); + assert!(cs.is_satisfied().unwrap()); + } + + #[test] + fn test_division() { + let mut rng = test_rng(); + let gen = Fr::get_root_of_unity(1 << 4).unwrap(); + assert_eq!(gen.pow(&[1 << 4]), Fr::one()); + let domain = Radix2Domain { + gen, + offset: Fr::multiplicative_generator(), + dim: 4, // 2^4 = 16 + }; + + let cs = ConstraintSystem::new_ref(); + + let ev_a = EvaluationsVar::from_vec_and_domain( + (0..16) + .map(|_| FpVar::new_input(ns!(cs, "poly_a"), || Ok(Fr::rand(&mut rng))).unwrap()) + .collect(), + domain.clone(), + false, + ); + let ev_b = EvaluationsVar::from_vec_and_domain( + (0..16) + .map(|_| FpVar::new_input(ns!(cs, "poly_a"), || Ok(Fr::rand(&mut rng))).unwrap()) + .collect(), + domain.clone(), + false, + ); + + let a_div_b = (&ev_a) / (&ev_b); + assert!(cs.is_satisfied().unwrap()); + let b_div_a = (&ev_b) / (&ev_a); + + let one = &a_div_b * &b_div_a; + for ev in one.evals.iter() { + assert!(Fr::is_one(&ev.value().unwrap())) + } + + assert!(cs.is_satisfied().unwrap()); + } +} diff --git a/src/poly/mod.rs b/src/poly/mod.rs index fde6e3f..f17cb27 100644 --- a/src/poly/mod.rs +++ b/src/poly/mod.rs @@ -1,2 +1,4 @@ +pub mod domain; +pub mod evaluations; /// Modules for working with polynomials in coefficient forms. pub mod polynomial;