Add Univariate Domain, Vanishing Polynomial, Lagrange Interpolation (#53)

* add domain and vp

* add lagrange interpolator

* add query position to coset

* nostd

* add test assertion

* fmt

* fix test

* add Add and Sub arithmetic

* add Add and Sub arithmetic

* add unit test for mul/div arithmetic

* add more doc for clarification

* add test for native interpolate

* add test for vp constraints

* fix lagrange interpolate bug

* comment cleanup + fmt

* add CHANGELOG

* fix a compile error

* Update CHANGELOG.md

* Update CHANGELOG.md

* fix comment

* doc fix

* doc update 2

* doc update 3

* pub lagrange_interpolator

* doc fix

* rename `EvaluationDomain` to `Radix2Domain`

* tweak

* tweak

Co-authored-by: weikeng <[email protected]>

Author: tsunrise

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

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"]

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<<dim`.
+///
+/// Native code corresponds to `ark-poly::univariate::domain::radix2`, but `ark-poly` only supports
+/// subgroup for now.
+///
+/// TODO: support cosets in `ark-poly`.
+pub struct Radix2Domain<F: PrimeField> {
+    /// 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<F: PrimeField> Radix2Domain<F> {
+    /// 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<F> {
+        let mut result = Vec::new();
+        let mut cur = self.gen;
+        for _ in 0..dim {
+            result.push(cur);
+            cur = cur * cur;
+        }
+        result
+    }
+
+    /// For domain `h<g>` with dimension `n`, `position` represented by `query_pos` in big endian form,
+    /// returns `h*g^{position}<g^{n-query_pos.len()}>`
+    pub fn query_position_to_coset(
+        &self,
+        query_pos: &[Boolean<F>],
+        coset_dim: u64,
+    ) -> Result<Vec<FpVar<F>>, 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<F> = 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)
+    }
+}

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<F: Field> {
+    /// h^|H|
+    pub constant_term: F,
+    /// log_2(|H|)
+    pub dim_h: u64,
+    /// |H|
+    pub order_h: u64,
+}
+
+impl<F: PrimeField> VanishingPolynomial<F> {
+    /// returns a VanishingPolynomial of coset `H = h<g>`.
+    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<F>) -> Result<FpVar<F>, 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);
+    }
+}

src/poly/evaluations/mod.rs

@@ -0,0 +1 @@
+pub mod univariate;

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<F: PrimeField> {
+    pub(crate) domain_order: usize,
+    pub(crate) all_domain_elems: Vec<F>,
+    pub(crate) v_inv_elems: Vec<F>,
+    pub(crate) domain_vp: VanishingPolynomial<F>,
+    poly_evaluations: Vec<F>,
+}
+
+impl<F: PrimeField> LagrangeInterpolator<F> {
+    /// Returns a lagrange interpolator, given the domain specification.
+    pub fn new(
+        domain_offset: F,
+        domain_generator: F,
+        domain_dim: u64,
+        poly_evaluations: Vec<F>,
+    ) -> 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<F> = 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<F> = 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<F> {
+        /*
+        * 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<F> = 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::<F>(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)
+    }
+}