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