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    pub offset: FpVar<F>,
23    /// dimension of evaluation domain
24    pub dim: u64,
25}
26
27impl<F: PrimeField> EqGadget<F> for Radix2DomainVar<F> {
28    fn is_eq(&self, other: &Self) -> Result<Boolean<F>, SynthesisError> {
29        if self.gen != other.gen || self.dim != other.dim {
30            Ok(Boolean::constant(false))
31        } else {
32            self.offset.is_eq(&other.offset)
33        }
34    }
35}
36
37impl<F: PrimeField> Radix2DomainVar<F> {
38    /// order of the domain
39    pub fn order(&self) -> usize {
40        1 << self.dim
41    }
42
43    /// Returns g, g^2, ..., g^{dim}
44    fn powers_of_gen(&self, dim: usize) -> Vec<F> {
45        let mut result = Vec::new();
46        let mut cur = self.gen;
47        for _ in 0..dim {
48            result.push(cur);
49            cur = cur * cur;
50        }
51        result
52    }
53
54    /// Size of the domain
55    pub fn size(&self) -> u64 {
56        1 << self.dim
57    }
58
59    /// For domain `h<g>` with dimension `n`, `position` represented by `query_pos` in big endian form,
60    /// returns `h*g^{position}<g^{n-query_pos.len()}>`
61    pub fn query_position_to_coset(
62        &self,
63        query_pos: &[Boolean<F>],
64        coset_dim: u64,
65    ) -> Result<Vec<FpVar<F>>, SynthesisError> {
66        let mut coset_index = query_pos;
67        assert!(
68            query_pos.len() == self.dim as usize
69                || query_pos.len() == (self.dim - coset_dim) as usize
70        );
71        if query_pos.len() == self.dim as usize {
72            coset_index = &coset_index[0..(coset_index.len() - coset_dim as usize)];
73        }
74        let mut coset = Vec::new();
75        let powers_of_g = &self.powers_of_gen(self.dim as usize)[(coset_dim as usize)..];
76
77        let mut first_point_in_coset: FpVar<F> = FpVar::zero();
78        for i in 0..coset_index.len() {
79            let term = coset_index[i].select(&FpVar::constant(powers_of_g[i]), &FpVar::zero())?;
80            first_point_in_coset += &term;
81        }
82
83        first_point_in_coset *= &self.offset;
84
85        coset.push(first_point_in_coset);
86        for i in 1..(1 << (coset_dim as usize)) {
87            let new_elem = &coset[i - 1] * &FpVar::Constant(self.gen);
88            coset.push(new_elem);
89        }
90
91        Ok(coset)
92    }
93}