1use crate::boolean::Boolean;
  2use crate::eq::EqGadget;
  3use crate::fields::fp::FpVar;
  4use crate::fields::FieldVar;
  5use ark_ff::PrimeField;
  6use ark_relations::r1cs::SynthesisError;
  7use ark_std::vec::Vec;
  8
  9pub mod vanishing_poly;
 10
 11#[derive(Clone, Debug)]
 12/// Defines an evaluation domain over a prime field. The domain is a coset of size `1<<dim`.
 13///
 14/// Native code corresponds to `ark-poly::univariate::domain::radix2`, but `ark-poly` only supports
 15/// subgroup for now.
 16///
 17// TODO: support cosets in `ark-poly`.
 18pub struct Radix2DomainVar<F: PrimeField> {
 19    /// generator of subgroup g
 20    pub gen: F,
 21    /// index of the quotient group (i.e. the `offset`)
 22    offset: FpVar<F>,
 23    /// dimension of evaluation domain, which is log2(size of coset)
 24    pub dim: u64,
 25}
 26impl<F: PrimeField> Radix2DomainVar<F> {
 27    /// Construct an evaluation domain with the given offset.
 28    pub fn new(gen: F, dimension: u64, offset: FpVar<F>) -> Result<Self, SynthesisError> {
 29        offset.enforce_not_equal(&FpVar::zero())?;
 30        Ok(Self {
 31            gen,
 32            offset,
 33            dim: dimension,
 34        })
 35    }
 36
 37    /// What is the offset of `self`?
 38    pub fn offset(&self) -> &FpVar<F> {
 39        &self.offset
 40    }
 41}
 42
 43impl<F: PrimeField> EqGadget<F> for Radix2DomainVar<F> {
 44    fn is_eq(&self, other: &Self) -> Result<Boolean<F>, SynthesisError> {
 45        if self.gen != other.gen || self.dim != other.dim {
 46            Ok(Boolean::constant(false))
 47        } else {
 48            self.offset.is_eq(&other.offset)
 49        }
 50    }
 51}
 52
 53impl<F: PrimeField> Radix2DomainVar<F> {
 54    /// order of the domain
 55    pub fn order(&self) -> usize {
 56        1 << self.dim
 57    }
 58
 59    /// Returns offset, offset*g, offset*g^2, ..., offset*g^{coset_size}
 60    pub fn elements(&self) -> Vec<FpVar<F>> {
 61        let mut result = Vec::new();
 62        result.push(self.offset.clone());
 63        for _ in 1..(1 << self.dim) {
 64            let new_element = result.last().unwrap() * self.gen;
 65            result.push(new_element);
 66        }
 67        result
 68    }
 69
 70    /// Size of the domain
 71    pub fn size(&self) -> u64 {
 72        1 << self.dim
 73    }
 74
 75    /// For domain `h<g>` with dimension `n`, `position` represented by `query_pos` in big endian form,
 76    /// returns all points of `h*g^{position}<g^{2^{n-coset_dim}}>`. The result is the query coset at index `query_pos`
 77    /// for the FRI protocol.
 78    ///
 79    /// # Panics
 80    /// This function panics when `query_pos.len() != coset_dim` or `query_pos.len() != self.dim`.  
 81    pub fn query_position_to_coset_elements(
 82        &self,
 83        query_pos: &[Boolean<F>],
 84        coset_dim: u64,
 85    ) -> Result<Vec<FpVar<F>>, SynthesisError> {
 86        Ok(self
 87            .query_position_to_coset(query_pos, coset_dim)?
 88            .elements())
 89    }
 90
 91    /// For domain `h<g>` with dimension `n`, `position` represented by `query_pos` in big endian form,
 92    /// returns all points of `h*g^{position}<g^{n-query_pos.len()}>`
 93    ///
 94    /// # Panics
 95    /// This function panics when `query_pos.len() != coset_dim` or `query_pos.len() != self.dim`.  
 96    pub fn query_position_to_coset(
 97        &self,
 98        query_pos: &[Boolean<F>],
 99        coset_dim: u64,
100    ) -> Result<Self, SynthesisError> {
101        let coset_index = truncate_to_coset_index(query_pos, self.dim, coset_dim);
102        let offset_var = &self.offset * FpVar::Constant(self.gen).pow_le(&coset_index)?;
103        Ok(Self {
104            gen: self.gen.pow(&[1 << (self.dim - coset_dim)]), // distance between coset
105            offset: offset_var,
106            dim: coset_dim,
107        })
108    }
109}
110
111fn truncate_to_coset_index<F: PrimeField>(
112    query_pos: &[Boolean<F>],
113    codeword_dim: u64,
114    coset_dim: u64,
115) -> Vec<Boolean<F>> {
116    assert!(query_pos.len() == codeword_dim as usize || query_pos.len() == coset_dim as usize);
117    if query_pos.len() == codeword_dim as usize {
118        query_pos[0..(query_pos.len() - coset_dim as usize)].to_vec()
119    } else {
120        query_pos.to_vec()
121    }
122}
123
124#[cfg(test)]
125mod tests {
126    use crate::prelude::*;
127    use ark_ff::PrimeField;
128    use ark_relations::r1cs::ConstraintSystem;
129    use ark_std::{rand::Rng, test_rng};
130
131    use crate::{alloc::AllocVar, fields::fp::FpVar, poly::domain::Radix2DomainVar, R1CSVar};
132
133    fn test_query_coset_template<F: PrimeField>() {
134        const COSET_DIM: u64 = 7;
135        const COSET_SIZE: usize = 1 << COSET_DIM;
136        const LOCALIZATION: u64 = 3;
137        let cs = ConstraintSystem::new_ref();
138        let mut rng = test_rng();
139        let gen = F::get_root_of_unity(COSET_SIZE).unwrap();
140        let offset = F::rand(&mut rng);
141        let offset_var = FpVar::new_witness(cs.clone(), || Ok(offset)).unwrap();
142        let domain = Radix2DomainVar::new(gen, COSET_DIM, offset_var).unwrap();
143
144        let num_cosets = 1 << (COSET_DIM - LOCALIZATION);
145
146        let coset_index = rng.gen_range(0..num_cosets);
147        let coset_index_var = UInt32::new_witness(cs.clone(), || Ok(coset_index))
148            .unwrap()
149            .to_bits_le()
150            .into_iter()
151            .take(COSET_DIM as usize)
152            .collect::<Vec<_>>();
153
154        let elements_actual = domain
155            .query_position_to_coset(&coset_index_var, LOCALIZATION)
156            .unwrap()
157            .elements();
158
159        let elements_expected = domain
160            .elements()
161            .into_iter()
162            .skip(coset_index as usize)
163            .step_by(1 << (COSET_DIM - LOCALIZATION))
164            .collect::<Vec<_>>();
165
166        assert_eq!(elements_expected.len(), elements_actual.len());
167
168        elements_expected
169            .into_iter()
170            .zip(elements_actual.into_iter())
171            .for_each(|(left, right)| {
172                assert_eq!(left.value().unwrap(), right.value().unwrap());
173            });
174    }
175
176    #[test]
177    fn test_on_bls12_381() {
178        test_query_coset_template::<ark_bls12_381::Fr>();
179    }
180
181    #[test]
182    fn test_on_bls12_377() {
183        test_query_coset_template::<ark_bls12_377::Fr>();
184    }
185}