stwo_gpu/core/
circle.rs

1use core::ops::{Add, Mul, Neg, Sub};
2
3use num_traits::{One, Zero};
4
5use super::fields::m31::{BaseField, M31};
6use super::fields::qm31::SecureField;
7use super::fields::{ComplexConjugate, Field, FieldExpOps};
8use crate::core::channel::Channel;
9use crate::core::fields::qm31::P4;
10
11/// A point on the complex circle. Treated as an additive group.
12#[derive(Copy, Clone, Debug, Default, PartialEq, Eq, Hash)]
13pub struct CirclePoint<F> {
14    pub x: F,
15    pub y: F,
16}
17
18impl<F: Zero + Add<Output = F> + FieldExpOps + Sub<Output = F> + Neg<Output = F>> CirclePoint<F> {
19    pub fn zero() -> Self {
20        Self {
21            x: F::one(),
22            y: F::zero(),
23        }
24    }
25
26    pub fn double(&self) -> Self {
27        self.clone() + self.clone()
28    }
29
30    /// Applies the circle's x-coordinate doubling map.
31    ///
32    /// # Examples
33    ///
34    /// ```
35    /// use stwo::core::circle::{CirclePoint, M31_CIRCLE_GEN};
36    /// use stwo::core::fields::m31::M31;
37    /// let p = M31_CIRCLE_GEN.mul(17);
38    /// assert_eq!(CirclePoint::double_x(p.x), (p + p).x);
39    /// ```
40    pub fn double_x(x: F) -> F {
41        let sx = x.square();
42        sx.clone() + sx - F::one()
43    }
44
45    /// Returns the log order of a point.
46    ///
47    /// All points have an order of the form `2^k`.
48    ///
49    /// # Examples
50    ///
51    /// ```
52    /// use stwo::core::circle::{CirclePoint, M31_CIRCLE_GEN, M31_CIRCLE_LOG_ORDER};
53    /// use stwo::core::fields::m31::M31;
54    /// assert_eq!(M31_CIRCLE_GEN.log_order(), M31_CIRCLE_LOG_ORDER);
55    /// ```
56    pub fn log_order(&self) -> u32
57    where
58        F: PartialEq + Eq,
59    {
60        // we only need the x-coordinate to check order since the only point
61        // with x=1 is the circle's identity
62        let mut res = 0;
63        let mut cur = self.x.clone();
64        while cur != F::one() {
65            cur = Self::double_x(cur);
66            res += 1;
67        }
68        res
69    }
70
71    pub fn mul(&self, mut scalar: u128) -> CirclePoint<F> {
72        let mut res = Self::zero();
73        let mut cur = self.clone();
74        while scalar > 0 {
75            if scalar & 1 == 1 {
76                res = res + cur.clone();
77            }
78            cur = cur.double();
79            scalar >>= 1;
80        }
81        res
82    }
83
84    pub fn repeated_double(&self, n: u32) -> Self {
85        let mut res = self.clone();
86        for _ in 0..n {
87            res = res.double();
88        }
89        res
90    }
91
92    pub fn conjugate(&self) -> CirclePoint<F> {
93        Self {
94            x: self.x.clone(),
95            y: -self.y.clone(),
96        }
97    }
98
99    pub fn antipode(&self) -> CirclePoint<F> {
100        Self {
101            x: -self.x.clone(),
102            y: -self.y.clone(),
103        }
104    }
105
106    pub fn into_ef<EF: From<F>>(self) -> CirclePoint<EF> {
107        CirclePoint {
108            x: self.x.clone().into(),
109            y: self.y.clone().into(),
110        }
111    }
112
113    pub fn mul_signed(&self, off: isize) -> CirclePoint<F> {
114        if off > 0 {
115            self.mul(off as u128)
116        } else {
117            self.conjugate().mul(-off as u128)
118        }
119    }
120}
121
122impl<F: Zero + Add<Output = F> + FieldExpOps + Sub<Output = F> + Neg<Output = F>> Add
123    for CirclePoint<F>
124{
125    type Output = Self;
126
127    fn add(self, rhs: Self) -> Self::Output {
128        let x = self.x.clone() * rhs.x.clone() - self.y.clone() * rhs.y.clone();
129        let y = self.x * rhs.y + self.y * rhs.x;
130        Self { x, y }
131    }
132}
133
134impl<F: Zero + Add<Output = F> + FieldExpOps + Sub<Output = F> + Neg<Output = F>> Neg
135    for CirclePoint<F>
136{
137    type Output = Self;
138
139    fn neg(self) -> Self::Output {
140        self.conjugate()
141    }
142}
143
144impl<F: Zero + Add<Output = F> + FieldExpOps + Sub<Output = F> + Neg<Output = F>> Sub
145    for CirclePoint<F>
146{
147    type Output = Self;
148
149    fn sub(self, rhs: Self) -> Self::Output {
150        self + (-rhs)
151    }
152}
153
154impl<F: Field> ComplexConjugate for CirclePoint<F> {
155    fn complex_conjugate(&self) -> Self {
156        Self {
157            x: self.x.complex_conjugate(),
158            y: self.y.complex_conjugate(),
159        }
160    }
161}
162
163impl CirclePoint<SecureField> {
164    pub fn get_point(index: u128) -> Self {
165        assert!(index < SECURE_FIELD_CIRCLE_ORDER);
166        SECURE_FIELD_CIRCLE_GEN.mul(index)
167    }
168
169    pub fn get_random_point<C: Channel>(channel: &mut C) -> Self {
170        let t = channel.draw_secure_felt();
171        let t_square = t.square();
172
173        let one_plus_tsquared_inv = t_square.add(SecureField::one()).inverse();
174
175        let x = SecureField::one()
176            .add(t_square.neg())
177            .mul(one_plus_tsquared_inv);
178        let y = t.double().mul(one_plus_tsquared_inv);
179
180        Self { x, y }
181    }
182}
183
184/// A generator for the circle group over [M31].
185///
186/// # Examples
187///
188/// ```
189/// use stwo::core::circle::{CirclePoint, M31_CIRCLE_GEN};
190/// use stwo::core::fields::m31::M31;
191///
192/// // Adding a generator to itself (2^30) times should NOT yield the identity.
193/// let circle_point = M31_CIRCLE_GEN.repeated_double(30);
194/// assert_ne!(circle_point, CirclePoint::zero());
195///
196/// // Shown above ord(M31_CIRCLE_GEN) > 2^30 . Group order is 2^31.
197/// // Ord(M31_CIRCLE_GEN) must be a divisor of it, Hence ord(M31_CIRCLE_GEN) = 2^31.
198/// // Adding the generator to itself (2^31) times should yield the identity.
199/// let circle_point = M31_CIRCLE_GEN.repeated_double(31);
200/// assert_eq!(circle_point, CirclePoint::zero());
201/// ```
202pub const M31_CIRCLE_GEN: CirclePoint<M31> = CirclePoint {
203    x: M31::from_u32_unchecked(2),
204    y: M31::from_u32_unchecked(1268011823),
205};
206
207/// Order of [M31_CIRCLE_GEN].
208pub const M31_CIRCLE_LOG_ORDER: u32 = 31;
209
210/// A generator for the circle group over [SecureField].
211pub const SECURE_FIELD_CIRCLE_GEN: CirclePoint<SecureField> = CirclePoint {
212    x: SecureField::from_u32_unchecked(1, 0, 478637715, 513582971),
213    y: SecureField::from_u32_unchecked(992285211, 649143431, 740191619, 1186584352),
214};
215
216/// Order of [SECURE_FIELD_CIRCLE_GEN].
217pub const SECURE_FIELD_CIRCLE_ORDER: u128 = P4 - 1;
218
219/// Integer i that represent the circle point i * CIRCLE_GEN. Treated as an
220/// additive ring modulo `1 << M31_CIRCLE_LOG_ORDER`.
221#[derive(Copy, Clone, Debug, PartialEq, Eq, Ord, PartialOrd)]
222pub struct CirclePointIndex(pub usize);
223
224impl CirclePointIndex {
225    pub const fn zero() -> Self {
226        Self(0)
227    }
228
229    pub const fn generator() -> Self {
230        Self(1)
231    }
232
233    pub const fn reduce(self) -> Self {
234        Self(self.0 & ((1 << M31_CIRCLE_LOG_ORDER) - 1))
235    }
236
237    pub fn subgroup_gen(log_size: u32) -> Self {
238        assert!(log_size <= M31_CIRCLE_LOG_ORDER);
239        Self(1 << (M31_CIRCLE_LOG_ORDER - log_size))
240    }
241
242    pub fn to_point(self) -> CirclePoint<M31> {
243        M31_CIRCLE_GEN.mul(self.0 as u128)
244    }
245
246    pub fn half(self) -> Self {
247        assert!(self.0 & 1 == 0);
248        Self(self.0 >> 1)
249    }
250}
251
252impl Add for CirclePointIndex {
253    type Output = Self;
254
255    fn add(self, rhs: Self) -> Self::Output {
256        Self(self.0 + rhs.0).reduce()
257    }
258}
259
260impl Sub for CirclePointIndex {
261    type Output = Self;
262
263    fn sub(self, rhs: Self) -> Self::Output {
264        Self(self.0 + (1 << M31_CIRCLE_LOG_ORDER) - rhs.0).reduce()
265    }
266}
267
268impl Mul<usize> for CirclePointIndex {
269    type Output = Self;
270
271    fn mul(self, rhs: usize) -> Self::Output {
272        Self(self.0.wrapping_mul(rhs)).reduce()
273    }
274}
275
276impl Neg for CirclePointIndex {
277    type Output = Self;
278
279    fn neg(self) -> Self::Output {
280        Self((1 << M31_CIRCLE_LOG_ORDER) - self.0).reduce()
281    }
282}
283
284/// Represents the coset initial + \<step\>.
285#[derive(Copy, Clone, Debug, PartialEq, Eq)]
286pub struct Coset {
287    pub initial_index: CirclePointIndex,
288    pub initial: CirclePoint<M31>,
289    pub step_size: CirclePointIndex,
290    pub step: CirclePoint<M31>,
291    pub log_size: u32,
292}
293
294impl Coset {
295    pub fn new(initial_index: CirclePointIndex, log_size: u32) -> Self {
296        assert!(log_size <= M31_CIRCLE_LOG_ORDER);
297        let step_size = CirclePointIndex::subgroup_gen(log_size);
298        Self {
299            initial_index,
300            initial: initial_index.to_point(),
301            step: step_size.to_point(),
302            step_size,
303            log_size,
304        }
305    }
306
307    /// Creates a coset of the form <G_n>.
308    /// For example, for n=8, we get the point indices \[0,1,2,3,4,5,6,7\].
309    pub fn subgroup(log_size: u32) -> Self {
310        Self::new(CirclePointIndex::zero(), log_size)
311    }
312
313    /// Creates a coset of the form `G_2n + <G_n>`.
314    ///
315    /// For example, let n = 8 and denote `G_16 = x, <G_8> = <2x>`.
316    /// The point indices are `[x, 3x, 5x, 7x, 9x, 11x, 13x, 15x]`.
317    pub fn odds(log_size: u32) -> Self {
318        Self::new(CirclePointIndex::subgroup_gen(log_size + 1), log_size)
319    }
320
321    /// Creates a coset of the form `G_4n + <G_n>`. It's conjugate is `3 * G_4n + <G_n>`.
322    ///
323    /// For example, let n = 8 and denote `G_32 = x, <G_8> = <4x>`.
324    /// The point indices are `[x, 5x, 9x, 13x, 17x, 21x, 25x, 29x]`.
325    /// Conjugate coset indices are `[3x, 7x, 11x, 15x, 19x, 23x, 27x, 31x]`.
326    ///
327    /// Note: This coset union with its conjugate coset is the `odds(log_size + 1)` coset.
328    pub fn half_odds(log_size: u32) -> Self {
329        Self::new(CirclePointIndex::subgroup_gen(log_size + 2), log_size)
330    }
331
332    /// Returns the size of the coset.
333    pub const fn size(&self) -> usize {
334        1 << self.log_size()
335    }
336
337    /// Returns the log size of the coset.
338    pub const fn log_size(&self) -> u32 {
339        self.log_size
340    }
341
342    pub const fn iter(&self) -> CosetIterator<CirclePoint<M31>> {
343        CosetIterator {
344            cur: self.initial,
345            step: self.step,
346            remaining: self.size(),
347        }
348    }
349
350    pub const fn iter_indices(&self) -> CosetIterator<CirclePointIndex> {
351        CosetIterator {
352            cur: self.initial_index,
353            step: self.step_size,
354            remaining: self.size(),
355        }
356    }
357
358    /// Returns a new coset comprising of all points in current coset doubled.
359    pub fn double(&self) -> Self {
360        assert!(self.log_size > 0);
361        Self {
362            initial_index: self.initial_index * 2,
363            initial: self.initial.double(),
364            step: self.step.double(),
365            step_size: self.step_size * 2,
366            log_size: self.log_size.saturating_sub(1),
367        }
368    }
369
370    pub fn repeated_double(&self, n_doubles: u32) -> Self {
371        (0..n_doubles).fold(*self, |coset, _| coset.double())
372    }
373
374    /// Note that this function panics when self.log_size == 0.
375    pub fn is_doubling_of(&self, other: Self) -> bool {
376        self.log_size <= other.log_size
377            && *self == other.repeated_double(other.log_size - self.log_size)
378    }
379
380    pub const fn initial(&self) -> CirclePoint<M31> {
381        self.initial
382    }
383
384    pub fn index_at(&self, index: usize) -> CirclePointIndex {
385        self.initial_index + self.step_size.mul(index)
386    }
387
388    pub fn at(&self, index: usize) -> CirclePoint<M31> {
389        self.index_at(index).to_point()
390    }
391
392    pub fn shift(&self, shift_size: CirclePointIndex) -> Self {
393        let initial_index = self.initial_index + shift_size;
394        Self {
395            initial_index,
396            initial: initial_index.to_point(),
397            ..*self
398        }
399    }
400
401    /// Creates the conjugate coset: -initial -\<step\>.
402    pub fn conjugate(&self) -> Self {
403        let initial_index = -self.initial_index;
404        let step_size = -self.step_size;
405        Self {
406            initial_index,
407            initial: initial_index.to_point(),
408            step_size,
409            step: step_size.to_point(),
410            log_size: self.log_size,
411        }
412    }
413}
414
415impl IntoIterator for Coset {
416    type Item = CirclePoint<BaseField>;
417    type IntoIter = CosetIterator<CirclePoint<BaseField>>;
418
419    /// Iterates over the points in the coset.
420    fn into_iter(self) -> Self::IntoIter {
421        self.iter()
422    }
423}
424
425#[derive(Clone)]
426pub struct CosetIterator<T: Add> {
427    pub cur: T,
428    pub step: T,
429    pub remaining: usize,
430}
431
432impl<T: Add<Output = T> + Copy> Iterator for CosetIterator<T> {
433    type Item = T;
434
435    fn next(&mut self) -> Option<Self::Item> {
436        if self.remaining == 0 {
437            return None;
438        }
439        self.remaining -= 1;
440        let res = self.cur;
441        self.cur = self.cur + self.step;
442        Some(res)
443    }
444}
445
446#[cfg(test)]
447mod tests {
448    use hashbrown::HashSet;
449    use num_traits::{One, Pow};
450    use std_shims::{vec, Vec};
451
452    use super::{CirclePointIndex, Coset};
453    use crate::core::channel::Blake2sChannel;
454    use crate::core::circle::{CirclePoint, SECURE_FIELD_CIRCLE_GEN};
455    use crate::core::fields::qm31::{SecureField, P4};
456    use crate::core::fields::FieldExpOps;
457    use crate::core::poly::circle::CanonicCoset;
458
459    #[test]
460    fn test_iterator() {
461        let coset = Coset::new(CirclePointIndex(1), 3);
462        let actual_indices: Vec<_> = coset.iter_indices().collect();
463        let expected_indices = vec![
464            CirclePointIndex(1),
465            CirclePointIndex(1) + CirclePointIndex::subgroup_gen(3) * 1,
466            CirclePointIndex(1) + CirclePointIndex::subgroup_gen(3) * 2,
467            CirclePointIndex(1) + CirclePointIndex::subgroup_gen(3) * 3,
468            CirclePointIndex(1) + CirclePointIndex::subgroup_gen(3) * 4,
469            CirclePointIndex(1) + CirclePointIndex::subgroup_gen(3) * 5,
470            CirclePointIndex(1) + CirclePointIndex::subgroup_gen(3) * 6,
471            CirclePointIndex(1) + CirclePointIndex::subgroup_gen(3) * 7,
472        ];
473        assert_eq!(actual_indices, expected_indices);
474
475        let actual_points = coset.iter().collect::<Vec<_>>();
476        let expected_points: Vec<_> = expected_indices.iter().map(|i| i.to_point()).collect();
477        assert_eq!(actual_points, expected_points);
478    }
479
480    #[test]
481    fn test_coset_is_half_coset_with_conjugate() {
482        let canonic_coset = CanonicCoset::new(8);
483        let coset_points: HashSet<_> = canonic_coset.coset().iter().collect();
484
485        let half_coset_points: HashSet<_> = canonic_coset.half_coset().iter().collect();
486        let half_coset_conjugate_points: HashSet<_> =
487            canonic_coset.half_coset().conjugate().iter().collect();
488
489        assert!((&half_coset_points & &half_coset_conjugate_points).is_empty());
490        assert_eq!(
491            coset_points,
492            &half_coset_points | &half_coset_conjugate_points
493        )
494    }
495
496    #[test]
497    pub fn test_get_random_circle_point() {
498        let mut channel = Blake2sChannel::default();
499
500        let first_random_circle_point = CirclePoint::get_random_point(&mut channel);
501
502        // Assert that the next random circle point is different.
503        assert_ne!(
504            first_random_circle_point,
505            CirclePoint::get_random_point(&mut channel)
506        );
507    }
508
509    #[test]
510    pub fn test_secure_field_circle_gen() {
511        let prime_factors = [
512            (2, 33),
513            (3, 2),
514            (5, 1),
515            (7, 1),
516            (11, 1),
517            (31, 1),
518            (151, 1),
519            (331, 1),
520            (733, 1),
521            (1709, 1),
522            (368140581013, 1),
523        ];
524
525        assert_eq!(
526            prime_factors
527                .iter()
528                .map(|(p, e)| p.pow(*e as u32))
529                .product::<u128>(),
530            P4 - 1
531        );
532        assert_eq!(
533            SECURE_FIELD_CIRCLE_GEN.x.square() + SECURE_FIELD_CIRCLE_GEN.y.square(),
534            SecureField::one()
535        );
536        assert_eq!(SECURE_FIELD_CIRCLE_GEN.mul(P4 - 1), CirclePoint::zero());
537        for (p, _) in prime_factors.iter() {
538            assert_ne!(
539                SECURE_FIELD_CIRCLE_GEN.mul((P4 - 1) / *p),
540                CirclePoint::zero()
541            );
542        }
543    }
544}
stwo_gpu/core/circle.rs

stwo_gpu/core/
circle.rs
