stwo_gpu/core/poly/circle/
domain.rs

1use core::iter::Chain;
2
3use itertools::Itertools;
4use std_shims::Vec;
5
6use crate::core::circle::{
7    CirclePoint, CirclePointIndex, Coset, CosetIterator, M31_CIRCLE_LOG_ORDER,
8};
9use crate::core::fields::m31::BaseField;
10
11pub const MAX_CIRCLE_DOMAIN_LOG_SIZE: u32 = M31_CIRCLE_LOG_ORDER - 1;
12
13// A circle domain is defined using an half coset, so the minimum log size for a circle domain is
14// currently 1.
15pub const MIN_CIRCLE_DOMAIN_LOG_SIZE: u32 = 1;
16
17/// A valid domain for circle polynomial interpolation and evaluation.
18///
19/// Valid domains are a disjoint union of two conjugate cosets: `+-C + <G_n>`.
20/// The ordering defined on this domain is `C + iG_n`, and then `-C - iG_n`.
21#[derive(Copy, Clone, Debug, PartialEq, Eq)]
22pub struct CircleDomain {
23    pub half_coset: Coset,
24}
25
26impl CircleDomain {
27    /// Given a coset C + <G_n>, constructs the circle domain +-C + <G_n> (i.e.,
28    /// this coset and its conjugate).
29    pub const fn new(half_coset: Coset) -> Self {
30        Self { half_coset }
31    }
32
33    pub fn iter(&self) -> CircleDomainIterator {
34        self.half_coset
35            .iter()
36            .chain(self.half_coset.conjugate().iter())
37    }
38
39    /// Iterates over point indices.
40    pub fn iter_indices(&self) -> CircleDomainIndexIterator {
41        self.half_coset
42            .iter_indices()
43            .chain(self.half_coset.conjugate().iter_indices())
44    }
45
46    /// Returns the size of the domain.
47    pub const fn size(&self) -> usize {
48        1 << self.log_size()
49    }
50
51    /// Returns the log size of the domain.
52    pub const fn log_size(&self) -> u32 {
53        self.half_coset.log_size + 1
54    }
55
56    /// Returns the `i` th domain element.
57    pub fn at(&self, i: usize) -> CirclePoint<BaseField> {
58        self.index_at(i).to_point()
59    }
60
61    /// Returns the [CirclePointIndex] of the `i`th domain element.
62    pub fn index_at(&self, i: usize) -> CirclePointIndex {
63        if i < self.half_coset.size() {
64            self.half_coset.index_at(i)
65        } else {
66            -self.half_coset.index_at(i - self.half_coset.size())
67        }
68    }
69
70    /// Returns true if the domain is canonic.
71    ///
72    /// Canonic domains are domains with elements that are the entire set of points defined by
73    /// `G_2n + <G_n>` where `G_n` and `G_2n` are obtained by repeatedly doubling
74    /// [crate::core::circle::M31_CIRCLE_GEN].
75    pub fn is_canonic(&self) -> bool {
76        self.half_coset.initial_index * 4 == self.half_coset.step_size
77    }
78
79    /// Splits a circle domain into a smaller [CircleDomain]s, shifted by offsets.
80    pub fn split(&self, log_parts: u32) -> (CircleDomain, Vec<CirclePointIndex>) {
81        assert!(log_parts <= self.half_coset.log_size);
82        let subdomain = CircleDomain::new(Coset::new(
83            self.half_coset.initial_index,
84            self.half_coset.log_size - log_parts,
85        ));
86        let shifts = (0..1 << log_parts)
87            .map(|i| self.half_coset.step_size * i)
88            .collect_vec();
89        (subdomain, shifts)
90    }
91
92    pub fn shift(&self, shift: CirclePointIndex) -> CircleDomain {
93        CircleDomain::new(self.half_coset.shift(shift))
94    }
95}
96
97impl IntoIterator for CircleDomain {
98    type Item = CirclePoint<BaseField>;
99    type IntoIter = CircleDomainIterator;
100
101    /// Iterates over the points in the domain.
102    fn into_iter(self) -> CircleDomainIterator {
103        self.iter()
104    }
105}
106
107/// An iterator over points in a circle domain.
108///
109/// Let the domain be `+-c + <G>`. The first iterated points are `c + <G>`, then `-c + <-G>`.
110pub type CircleDomainIterator =
111    Chain<CosetIterator<CirclePoint<BaseField>>, CosetIterator<CirclePoint<BaseField>>>;
112
113/// Like [CircleDomainIterator] but returns corresponding [CirclePointIndex]s.
114type CircleDomainIndexIterator =
115    Chain<CosetIterator<CirclePointIndex>, CosetIterator<CirclePointIndex>>;
116
117#[cfg(test)]
118mod tests {
119    use itertools::Itertools;
120    use std_shims::Vec;
121
122    use super::CircleDomain;
123    use crate::core::circle::{CirclePointIndex, Coset};
124    use crate::core::poly::circle::CanonicCoset;
125
126    #[test]
127    fn test_circle_domain_iterator() {
128        let domain = CircleDomain::new(Coset::new(CirclePointIndex::generator(), 2));
129        for (i, point) in domain.iter().enumerate() {
130            if i < 4 {
131                assert_eq!(
132                    point,
133                    (CirclePointIndex::generator() + CirclePointIndex::subgroup_gen(2) * i)
134                        .to_point()
135                );
136            } else {
137                assert_eq!(
138                    point,
139                    (-(CirclePointIndex::generator() + CirclePointIndex::subgroup_gen(2) * i))
140                        .to_point()
141                );
142            }
143        }
144    }
145
146    #[test]
147    fn is_canonic_invalid_domain() {
148        let half_coset = Coset::new(CirclePointIndex::generator(), 4);
149        let not_canonic_domain = CircleDomain::new(half_coset);
150
151        assert!(!not_canonic_domain.is_canonic());
152    }
153
154    #[test]
155    fn test_at_circle_domain() {
156        let domain = CanonicCoset::new(7).circle_domain();
157        let half_domain_size = domain.size() / 2;
158
159        for i in 0..half_domain_size {
160            assert_eq!(domain.index_at(i), -domain.index_at(i + half_domain_size));
161            assert_eq!(domain.at(i), domain.at(i + half_domain_size).conjugate());
162        }
163    }
164
165    #[test]
166    fn test_domain_split() {
167        let domain = CanonicCoset::new(5).circle_domain();
168        let (subdomain, shifts) = domain.split(2);
169
170        let domain_points = domain.iter().collect::<Vec<_>>();
171        let points_for_each_domain = shifts
172            .iter()
173            .map(|&shift| (subdomain.shift(shift)).iter().collect_vec())
174            .collect::<Vec<_>>();
175        // Interleave the points from each subdomain.
176        let extended_points = (0..(1 << 3))
177            .flat_map(|point_ind| {
178                (0..(1 << 2))
179                    .map(|shift_ind| points_for_each_domain[shift_ind][point_ind])
180                    .collect_vec()
181            })
182            .collect_vec();
183        assert_eq!(domain_points, extended_points);
184    }
185}
stwo_gpu/core/poly/circle/domain.rs

stwo_gpu/core/poly/circle/
domain.rs
