Let `Radix2Domain::offset` to be `FpVar` instead of `F` (#65)

* restructure code

* done

* add changelog

* add the changelog to mark this as a breaking change

* add the CHANGELOG

* tweak

* add `EqGadget`

* rename generate_interpolate_cache to generate_interpolation_cache

* address the comment

Co-authored-by: weikeng <w.k@berkeley.edu>

CHANGELOG.md

@@ -2,29 +2,35 @@
 
 ### Breaking changes
 
-- #60 Rename `AllocatedBit` to `AllocatedBool` for consistency with the `Boolean` variable.
-You can update downstream usage with `grep -rl 'AllocatedBit' . | xargs env LANG=C env LC_CTYPE=C sed -i '' 's/AllocatedBit/AllocatedBool/g'`.
+- [\#60](https://github.com/arkworks-rs/r1cs-std/pull/60) Rename `AllocatedBit` to `AllocatedBool` for consistency with the `Boolean` variable.
+  You can update downstream usage with `grep -rl 'AllocatedBit' . | xargs env LANG=C env LC_CTYPE=C sed -i '' 's/AllocatedBit/AllocatedBool/g'`.
+- [\#65](https://github.com/arkworks-rs/r1cs-std/pull/65) Rename `Radix2Domain` in `r1cs-std` to `Radix2DomainVar`.
 
 ### Features
 
-- [\#53](https://github.com/arkworks-rs/r1cs-std/pull/53) Add univariate evaluation domain and lagrange interpolation. 
+- [\#53](https://github.com/arkworks-rs/r1cs-std/pull/53) Add univariate evaluation domain and Lagrange interpolation. 
 
 ### Improvements
 
+- [\#65](https://github.com/arkworks-rs/r1cs-std/pull/65) Add support for non-constant coset offset in `Radix2DomainVar`.
+
 ### Bug Fixes
 
 
 ## v0.2.0
 
 ### Breaking changes
+
 - [\#12](https://github.com/arkworks-rs/r1cs-std/pull/12) Make the output of the `ToBitsGadget` impl for `FpVar` fixed-size
 - [\#48](https://github.com/arkworks-rs/r1cs-std/pull/48) Add `Clone` trait bound to `CondSelectGadget`.
 
 ### Features
+
 - [\#21](https://github.com/arkworks-rs/r1cs-std/pull/21) Add `UInt128`
 - [\#50](https://github.com/arkworks-rs/r1cs-std/pull/50) Add `DensePolynomialVar`
 
 ### Improvements
+
 - [\#5](https://github.com/arkworks-rs/r1cs-std/pull/5) Speedup BLS-12 pairing
 - [\#13](https://github.com/arkworks-rs/r1cs-std/pull/13) Add `ToConstraintFieldGadget` to `ProjectiveVar`
 - [\#15](https://github.com/arkworks-rs/r1cs-std/pull/15), #16 Allow `cs` to be `None` when converting a Montgomery point into a Twisted Edwards point
@@ -38,6 +44,7 @@ You can update downstream usage with `grep -rl 'AllocatedBit' . | xargs env LANG
 - [\#46](https://github.com/arkworks-rs/r1cs-std/pull/46) Add mux gadget as an auto-impl in `CondSelectGadget` to support random access of an array
 
 ### Bug fixes
+
 - [\#8](https://github.com/arkworks-rs/r1cs-std/pull/8) Fix bug in `three_bit_cond_neg_lookup` when using a constant lookup bit
 - [\#9](https://github.com/arkworks-rs/r1cs-std/pull/9) Fix bug in `short_weierstrass::ProjectiveVar::to_affine`
 - [\#29](https://github.com/arkworks-rs/r1cs-std/pull/29) Fix `to_non_unique_bytes` for `BLS12::G1Prepared`

src/poly/domain/mod.rs

@@ -1,4 +1,5 @@
 use crate::boolean::Boolean;
+use crate::eq::EqGadget;
 use crate::fields::fp::FpVar;
 use crate::fields::FieldVar;
 use ark_ff::PrimeField;
@@ -7,23 +8,33 @@ use ark_std::vec::Vec;
 
 pub mod vanishing_poly;
 
-#[derive(Copy, Clone, Hash, Eq, PartialEq, Debug)]
+#[derive(Clone, 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> {
+pub struct Radix2DomainVar<F: PrimeField> {
     /// generator of subgroup g
     pub gen: F,
     /// index of the quotient group (i.e. the `offset`)
-    pub offset: F,
+    pub offset: FpVar<F>,
     /// dimension of evaluation domain
     pub dim: u64,
 }
 
-impl<F: PrimeField> Radix2Domain<F> {
+impl<F: PrimeField> EqGadget<F> for Radix2DomainVar<F> {
+    fn is_eq(&self, other: &Self) -> Result<Boolean<F>, SynthesisError> {
+        if self.gen != other.gen || self.dim != other.dim {
+            Ok(Boolean::constant(false))
+        } else {
+            self.offset.is_eq(&other.offset)
+        }
+    }
+}
+
+impl<F: PrimeField> Radix2DomainVar<F> {
     /// order of the domain
     pub fn order(&self) -> usize {
         1 << self.dim
@@ -40,6 +51,11 @@ impl Radix2Domain {
         result
     }
 
+    /// Size of the domain
+    pub fn size(&self) -> u64 {
+        1 << self.dim
+    }
+
     /// 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(
@@ -64,7 +80,7 @@ impl Radix2Domain {
             first_point_in_coset += &term;
         }
 
-        first_point_in_coset *= &FpVar::Constant(self.offset);
+        first_point_in_coset *= &self.offset;
 
         coset.push(first_point_in_coset);
         for i in 1..(1 << (coset_dim as usize)) {

src/poly/evaluations/univariate/lagrange_interpolator.rs

@@ -98,8 +98,11 @@ impl LagrangeInterpolator {
 
 #[cfg(test)]
 mod tests {
-    use crate::poly::domain::Radix2Domain;
+    use crate::fields::fp::FpVar;
+    use crate::fields::FieldVar;
+    use crate::poly::domain::Radix2DomainVar;
     use crate::poly::evaluations::univariate::lagrange_interpolator::LagrangeInterpolator;
+    use crate::R1CSVar;
     use ark_ff::{FftField, Field, One};
     use ark_poly::univariate::DensePolynomial;
     use ark_poly::{Polynomial, UVPolynomial};
@@ -112,21 +115,25 @@ mod tests {
         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 {
+        let domain = Radix2DomainVar {
             gen,
-            offset: Fr::multiplicative_generator(),
+            offset: FpVar::constant(Fr::multiplicative_generator()),
             dim: 4, // 2^4 = 16
         };
         // generate evaluations of `poly` on this domain
-        let mut coset_point = domain.offset;
+        let mut coset_point = domain.offset.value().unwrap();
         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);
+        let interpolator = LagrangeInterpolator::new(
+            domain.offset.value().unwrap(),
+            domain.gen,
+            domain.dim,
+            oracle_evals,
+        );
 
         // the point to evaluate at
         let interpolate_point = Fr::rand(&mut rng);

src/poly/evaluations/univariate/mod.rs

@@ -3,7 +3,7 @@ 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::domain::Radix2DomainVar;
 use crate::poly::evaluations::univariate::lagrange_interpolator::LagrangeInterpolator;
 use crate::R1CSVar;
 use ark_ff::{batch_inversion, PrimeField};
@@ -18,7 +18,12 @@ pub struct EvaluationsVar {
     pub evals: Vec<FpVar<F>>,
     /// Optional Lagrange Interpolator. Useful for lagrange interpolation.
     pub lagrange_interpolator: Option<LagrangeInterpolator<F>>,
-    domain: Radix2Domain<F>,
+    domain: Radix2DomainVar<F>,
+    /// Contains all domain elements of `domain.base_domain`.
+    ///
+    /// This is a cache for lagrange interpolation when offset is non-constant. Will be `None` if offset is constant
+    /// or `interpolate` is set to `false`.
+    subgroup_points: Option<Vec<F>>,
 }
 
 impl<F: PrimeField> EvaluationsVar<F> {
@@ -27,7 +32,7 @@ impl EvaluationsVar {
     /// using lagrange interpolation.
     pub fn from_vec_and_domain(
         evaluations: Vec<FpVar<F>>,
-        domain: Radix2Domain<F>,
+        domain: Radix2DomainVar<F>,
         interpolate: bool,
     ) -> Self {
         assert_eq!(
@@ -40,22 +45,39 @@ impl EvaluationsVar {
             evals: evaluations,
             lagrange_interpolator: None,
             domain,
+            subgroup_points: None,
         };
         if interpolate {
-            ev.generate_lagrange_interpolator();
+            ev.generate_interpolation_cache();
         }
         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)
+    /// Precompute necessary calculation for lagrange interpolation and mark it ready to interpolate
+    pub fn generate_interpolation_cache(&mut self) {
+        if self.domain.offset.is_constant() {
+            let poly_evaluations_val: Vec<_> =
+                self.evals.iter().map(|v| v.value().unwrap()).collect();
+            let domain = &self.domain;
+            let lagrange_interpolator = if let FpVar::Constant(x) = domain.offset {
+                LagrangeInterpolator::new(x, domain.gen, domain.dim, poly_evaluations_val)
+            } else {
+                panic!("Domain offset needs to be constant.")
+            };
+            self.lagrange_interpolator = Some(lagrange_interpolator)
+        } else {
+            // calculate all elements of base subgroup so that in later part we don't need to calculate the exponents again
+            let mut subgroup_points = Vec::with_capacity(self.domain.size() as usize);
+            subgroup_points.push(F::one());
+            for i in 1..self.domain.size() as usize {
+                subgroup_points.push(subgroup_points[i - 1] * self.domain.gen)
+            }
+            self.subgroup_points = Some(subgroup_points)
+        }
     }
 
+    /// Compute lagrange coefficients for each evaluation, given `interpolation_point`.
+    /// Only valid if the domain offset is constant.
     fn compute_lagrange_coefficients(
         &self,
         interpolation_point: &FpVar<F>,
@@ -67,7 +89,7 @@ impl EvaluationsVar {
             .lagrange_interpolator
             .as_ref()
             .expect("lagrange interpolator has not been initialized. \
-            Call `self.generate_lagrange_interpolator` first or set `interpolate` to true in constructor. ");
+            Call `self.generate_interpolation_cache` 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();
@@ -105,6 +127,19 @@ impl EvaluationsVar {
     pub fn interpolate_and_evaluate(
         &self,
         interpolation_point: &FpVar<F>,
+    ) -> Result<FpVar<F>, SynthesisError> {
+        // specialize: if domain offset is constant, we can optimize to have fewer constraints
+        if self.domain.offset.is_constant() {
+            self.lagrange_interpolate_with_constant_offset(interpolation_point)
+        } else {
+            // if domain offset is not constant, then we use standard lagrange interpolation code
+            self.lagrange_interpolate_with_non_constant_offset(interpolation_point)
+        }
+    }
+
+    fn lagrange_interpolate_with_constant_offset(
+        &self,
+        interpolation_point: &FpVar<F>,
     ) -> Result<FpVar<F>, SynthesisError> {
         let lagrange_interpolator = self
             .lagrange_interpolator
@@ -119,6 +154,50 @@ impl EvaluationsVar {
 
         Ok(interpolation)
     }
+
+    /// Generate interpolation constraints. We assume at compile time we know the base coset (i.e. `gen`) but not know `offset`.
+    fn lagrange_interpolate_with_non_constant_offset(
+        &self,
+        interpolation_point: &FpVar<F>,
+    ) -> Result<FpVar<F>, SynthesisError> {
+        // first, make sure `subgroup_points` is made
+        let subgroup_points = self.subgroup_points.as_ref()
+            .expect("lagrange interpolator has not been initialized. \
+            Call `self.generate_interpolation_cache` first or set `interpolate` to true in constructor. ");
+        // Let denote interpolation_point as alpha.
+        // Lagrange polynomial for coset element `a` is
+        // \frac{1}{size * offset ^ size} * \frac{alpha^size - offset^size}{alpha * a^{-1} - 1}
+        // Notice that a = (offset * a') where a' is the corresponding element of base coset
+
+        // let `lhs` become \frac{alpha^size - offset^size}{size * offset ^ size}. This part is shared by all lagrange polynomials
+        let coset_offset_to_size = self.domain.offset.pow_by_constant(&[self.domain.size()])?; // offset^size
+        let alpha_to_s = interpolation_point.pow_by_constant(&[self.domain.size()])?;
+        let lhs_numerator = &alpha_to_s - &coset_offset_to_size;
+        let lhs_denominator = &coset_offset_to_size * FpVar::constant(F::from(self.domain.size()));
+
+        let lhs = lhs_numerator.mul_by_inverse(&lhs_denominator)?;
+
+        // `rhs` for coset element `a` is \frac{1}{alpha * a^{-1} - 1} = \frac{1}{alpha * offset^{-1} * a'^{-1} - 1}
+        let alpha_coset_offset_inv = interpolation_point.mul_by_inverse(&self.domain.offset)?;
+
+        // `res` stores the sum of all lagrange polynomials evaluated at alpha
+        let mut res = FpVar::<F>::zero();
+
+        let domain_size = self.domain.size() as usize;
+        for i in 0..domain_size {
+            // a'^{-1} where a is the base coset element
+            let subgroup_point_inv = subgroup_points[(domain_size - i) % domain_size];
+            debug_assert_eq!(subgroup_points[i] * subgroup_point_inv, F::one());
+            // alpha * offset^{-1} * a'^{-1} - 1
+            let lag_donom = &alpha_coset_offset_inv * subgroup_point_inv - F::one();
+            let lag_coeff = lhs.mul_by_inverse(&lag_donom)?;
+
+            let lag_interpoland = &self.evals[i] * lag_coeff;
+            res += lag_interpoland
+        }
+
+        Ok(res)
+    }
 }
 
 impl<'a, 'b, F: PrimeField> Add<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
@@ -132,8 +211,13 @@ impl<'a, 'b, F: PrimeField> Add<&'a EvaluationsVar> for &'b EvaluationsVar
 }
 
 impl<'a, F: PrimeField> AddAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
+    /// Performs the `+=` operations, assuming `domain.offset` is equal.
     fn add_assign(&mut self, other: &'a EvaluationsVar<F>) {
-        assert_eq!(self.domain, other.domain, "domains are unequal");
+        // offset might be unknown at compile time, so we assume offset is equal
+        assert!(
+            self.domain.gen == other.domain.gen && self.domain.dim == other.domain.dim,
+            "domains are unequal"
+        );
 
         self.lagrange_interpolator = None;
         self.evals
@@ -154,8 +238,13 @@ impl<'a, 'b, F: PrimeField> Sub<&'a EvaluationsVar> for &'b EvaluationsVar
 }
 
 impl<'a, F: PrimeField> SubAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
+    /// Performs the `-=` operations, assuming `domain.offset` is equal.
     fn sub_assign(&mut self, other: &'a EvaluationsVar<F>) {
-        assert_eq!(self.domain, other.domain, "domains are unequal");
+        // offset might be unknown at compile time, so we assume offset is equal
+        assert!(
+            self.domain.gen == other.domain.gen && self.domain.dim == other.domain.dim,
+            "domains are unequal"
+        );
 
         self.lagrange_interpolator = None;
         self.evals
@@ -168,6 +257,7 @@ impl<'a, F: PrimeField> SubAssign<&'a EvaluationsVar> for EvaluationsVar {
 impl<'a, 'b, F: PrimeField> Mul<&'a EvaluationsVar<F>> for &'b EvaluationsVar<F> {
     type Output = EvaluationsVar<F>;
 
+    /// Performs the `*` operations, assuming `domain.offset` is equal.
     fn mul(self, rhs: &'a EvaluationsVar<F>) -> Self::Output {
         let mut result = self.clone();
         result *= rhs;
@@ -176,8 +266,13 @@ impl<'a, 'b, F: PrimeField> Mul<&'a EvaluationsVar> for &'b EvaluationsVar
 }
 
 impl<'a, F: PrimeField> MulAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
+    /// Performs the `*=` operations, assuming `domain.offset` is equal.
     fn mul_assign(&mut self, other: &'a EvaluationsVar<F>) {
-        assert_eq!(self.domain, other.domain, "domains are unequal");
+        // offset might be unknown at compile time, so we assume offset is equal
+        assert!(
+            self.domain.gen == other.domain.gen && self.domain.dim == other.domain.dim,
+            "domains are unequal"
+        );
 
         self.lagrange_interpolator = None;
         self.evals
@@ -198,8 +293,13 @@ impl<'a, 'b, F: PrimeField> Div<&'a EvaluationsVar> for &'b EvaluationsVar
 }
 
 impl<'a, F: PrimeField> DivAssign<&'a EvaluationsVar<F>> for EvaluationsVar<F> {
+    /// Performs the `/=` operations, assuming `domain.offset` is equal.
     fn div_assign(&mut self, other: &'a EvaluationsVar<F>) {
-        assert_eq!(self.domain, other.domain, "domains are unequal");
+        // offset might be unknown at compile time, so we assume offset is equal
+        assert!(
+            self.domain.gen == other.domain.gen && self.domain.dim == other.domain.dim,
+            "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();
@@ -228,7 +328,8 @@ impl<'a, F: PrimeField> DivAssign<&'a EvaluationsVar> for EvaluationsVar {
 mod tests {
     use crate::alloc::AllocVar;
     use crate::fields::fp::FpVar;
-    use crate::poly::domain::Radix2Domain;
+    use crate::fields::FieldVar;
+    use crate::poly::domain::Radix2DomainVar;
     use crate::poly::evaluations::univariate::EvaluationsVar;
     use crate::R1CSVar;
     use ark_ff::{FftField, Field, One, UniformRand};
@@ -239,17 +340,17 @@ mod tests {
     use ark_test_curves::bls12_381::Fr;
 
     #[test]
-    fn test_interpolate() {
+    fn test_interpolate_constant_offset() {
         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 {
+        let domain = Radix2DomainVar {
             gen,
-            offset: Fr::multiplicative_generator(),
+            offset: FpVar::constant(Fr::rand(&mut rng)),
             dim: 4, // 2^4 = 16
         };
-        let mut coset_point = domain.offset;
+        let mut coset_point = domain.offset.value().unwrap();
         let mut oracle_evals = Vec::new();
         for _ in 0..(1 << 4) {
             oracle_evals.push(poly.evaluate(&coset_point));
@@ -276,6 +377,49 @@ mod tests {
 
         assert_eq!(actual, expected);
         assert!(cs.is_satisfied().unwrap());
+        println!("number of constraints: {}", cs.num_constraints())
+    }
+
+    #[test]
+    fn test_interpolate_non_constant_offset() {
+        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 cs = ConstraintSystem::new_ref();
+        let domain = Radix2DomainVar {
+            gen,
+            offset: FpVar::new_witness(ns!(cs, "offset"), || Ok(Fr::rand(&mut rng))).unwrap(),
+            dim: 4, // 2^4 = 16
+        };
+        let mut coset_point = domain.offset.value().unwrap();
+        let mut oracle_evals = Vec::new();
+        for _ in 0..(1 << 4) {
+            oracle_evals.push(poly.evaluate(&coset_point));
+            coset_point *= gen;
+        }
+
+        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());
+        println!("number of constraints: {}", cs.num_constraints())
     }
 
     #[test]
@@ -283,9 +427,9 @@ mod tests {
         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 {
+        let domain = Radix2DomainVar {
             gen,
-            offset: Fr::multiplicative_generator(),
+            offset: FpVar::constant(Fr::multiplicative_generator()),
             dim: 4, // 2^4 = 16
         };