Bug Fix in domain and API changes (#86)

* Bug Fix in `domain` and API changes

Currently, `domain.query_position_to_coset` is not consistent native code in `ark-ldt`. This commit does the following changes:
* change `query_position_to_coset` to `query_position_to_coset_elements`, and make the return result consistent with native code
* add function `query_position_to_coset` to just return generator and offset instead of coset elements. This is useful when evaluating vanishing poly.

* minor bug fix

* update CHANGELOG

* `query_position_to_coset_elements` use `query_position_to_coset`

* Update CHANGELOG.md

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

CHANGELOG.md

@@ -4,6 +4,8 @@
 
 ### Breaking changes
 
+- [\#86](https://github.com/arkworks-rs/r1cs-std/pull/86) Change the API for domains for coset.
+
 ### Features
 
 - [\#79](https://github.com/arkworks-rs/r1cs-std/pull/79) Move `NonNativeFieldVar` from `ark-nonnative` to `ark-r1cs-std`.
@@ -14,6 +16,7 @@
 
 ### Bug Fixes
 
+- [\#86](https://github.com/arkworks-rs/r1cs-std/pull/86) Make result of `query_position_to_coset` consistent with `ark-ldt`.
 - [\#77](https://github.com/arkworks-rs/r1cs-std/pull/77) Fix BLS12 `G2PreparedGadget`'s `AllocVar` when G2 uses a divisive twist.
 
 ## v0.3.1

src/poly/domain/mod.rs

@@ -20,7 +20,7 @@ pub struct Radix2DomainVar<F: PrimeField> {
     pub gen: F,
     /// index of the quotient group (i.e. the `offset`)
     offset: FpVar<F>,
-    /// dimension of evaluation domain
+    /// dimension of evaluation domain, which is log2(size of coset)
     pub dim: u64,
 }
 impl<F: PrimeField> Radix2DomainVar<F> {
@@ -56,13 +56,13 @@ impl<F: PrimeField> Radix2DomainVar<F> {
         1 << self.dim
     }
 
-    /// Returns g, g^2, ..., g^{dim}
-    fn powers_of_gen(&self, dim: usize) -> Vec<F> {
+    /// Returns offset, offset*g, offset*g^2, ..., offset*g^{coset_size}
+    pub fn elements(&self) -> Vec<FpVar<F>> {
         let mut result = Vec::new();
-        let mut cur = self.gen;
-        for _ in 0..dim {
-            result.push(cur);
-            cur = cur * cur;
+        result.push(self.offset.clone());
+        for _ in 1..(1 << self.dim) {
+            let new_element = result.last().unwrap() * self.gen;
+            result.push(new_element);
         }
         result
     }
@@ -73,37 +73,113 @@ impl<F: PrimeField> Radix2DomainVar<F> {
     }
 
     /// 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(
+    /// returns all points of `h*g^{position}<g^{2^{n-coset_dim}}>`. The result is the query coset at index `query_pos`
+    /// for the FRI protocol.
+    ///
+    /// # Panics
+    /// This function panics when `query_pos.len() != coset_dim` or `query_pos.len() != self.dim`.  
+    pub fn query_position_to_coset_elements(
         &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)..];
+        Ok(self
+            .query_position_to_coset(query_pos, coset_dim)?
+            .elements())
+    }
 
-        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 *= &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)
+    /// For domain `h<g>` with dimension `n`, `position` represented by `query_pos` in big endian form,
+    /// returns all points of `h*g^{position}<g^{n-query_pos.len()}>`
+    ///
+    /// # Panics
+    /// This function panics when `query_pos.len() != coset_dim` or `query_pos.len() != self.dim`.  
+    pub fn query_position_to_coset(
+        &self,
+        query_pos: &[Boolean<F>],
+        coset_dim: u64,
+    ) -> Result<Self, SynthesisError> {
+        let coset_index = truncate_to_coset_index(query_pos, self.dim, coset_dim);
+        let offset_var = &self.offset * FpVar::Constant(self.gen).pow_le(&coset_index)?;
+        Ok(Self {
+            gen: self.gen.pow(&[1 << (self.dim - coset_dim)]), // distance between coset
+            offset: offset_var,
+            dim: coset_dim,
+        })
+    }
+}
+
+fn truncate_to_coset_index<F: PrimeField>(
+    query_pos: &[Boolean<F>],
+    codeword_dim: u64,
+    coset_dim: u64,
+) -> Vec<Boolean<F>> {
+    assert!(query_pos.len() == codeword_dim as usize || query_pos.len() == coset_dim as usize);
+    if query_pos.len() == codeword_dim as usize {
+        query_pos[0..(query_pos.len() - coset_dim as usize)].to_vec()
+    } else {
+        query_pos.to_vec()
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::prelude::*;
+    use ark_ff::PrimeField;
+    use ark_relations::r1cs::ConstraintSystem;
+    use ark_std::{rand::Rng, test_rng};
+
+    use crate::{alloc::AllocVar, fields::fp::FpVar, poly::domain::Radix2DomainVar, R1CSVar};
+
+    fn test_query_coset_template<F: PrimeField>() {
+        const COSET_DIM: u64 = 7;
+        const COSET_SIZE: usize = 1 << COSET_DIM;
+        const LOCALIZATION: u64 = 3;
+        let cs = ConstraintSystem::new_ref();
+        let mut rng = test_rng();
+        let gen = F::get_root_of_unity(COSET_SIZE).unwrap();
+        let offset = F::rand(&mut rng);
+        let offset_var = FpVar::new_witness(cs.clone(), || Ok(offset)).unwrap();
+        let domain = Radix2DomainVar::new(gen, COSET_DIM, offset_var).unwrap();
+
+        let num_cosets = 1 << (COSET_DIM - LOCALIZATION);
+
+        let coset_index = rng.gen_range(0..num_cosets);
+        let coset_index_var = UInt32::new_witness(cs.clone(), || Ok(coset_index))
+            .unwrap()
+            .to_bits_le()
+            .into_iter()
+            .take(COSET_DIM as usize)
+            .collect::<Vec<_>>();
+
+        let elements_actual = domain
+            .query_position_to_coset(&coset_index_var, LOCALIZATION)
+            .unwrap()
+            .elements();
+
+        let elements_expected = domain
+            .elements()
+            .into_iter()
+            .skip(coset_index as usize)
+            .step_by(1 << (COSET_DIM - LOCALIZATION))
+            .collect::<Vec<_>>();
+
+        assert_eq!(elements_expected.len(), elements_actual.len());
+
+        elements_expected
+            .into_iter()
+            .zip(elements_actual.into_iter())
+            .for_each(|(left, right)| {
+                assert_eq!(left.value().unwrap(), right.value().unwrap());
+            });
+    }
+
+    #[test]
+    fn test_on_bls12_381() {
+        test_query_coset_template::<ark_bls12_381::Fr>();
+    }
+
+    #[test]
+    fn test_on_bls12_377() {
+        test_query_coset_template::<ark_bls12_377::Fr>();
     }
 }