stwo_gpu/core/fields/
cm31.rs

1use core::fmt::{Debug, Display};
2use core::ops::{
3    Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Rem, RemAssign, Sub, SubAssign,
4};
5
6use serde::{Deserialize, Serialize};
7
8use super::{ComplexConjugate, FieldExpOps};
9use crate::core::fields::m31::M31;
10use crate::{impl_extension_field, impl_field};
11pub const P2: u64 = 4611686014132420609; // (2 ** 31 - 1) ** 2
12
13/// Complex extension field of M31.
14/// Equivalent to M31\[x\] over (x^2 + 1) as the irreducible polynomial.
15/// Represented as (a, b) of a + bi.
16#[derive(Copy, Clone, Default, PartialEq, Eq, PartialOrd, Ord, Hash, Deserialize, Serialize)]
17pub struct CM31(pub M31, pub M31);
18
19impl_field!(CM31, P2);
20impl_extension_field!(CM31, M31);
21
22impl CM31 {
23    pub const fn from_u32_unchecked(a: u32, b: u32) -> CM31 {
24        Self(M31::from_u32_unchecked(a), M31::from_u32_unchecked(b))
25    }
26
27    pub const fn from_m31(a: M31, b: M31) -> CM31 {
28        Self(a, b)
29    }
30}
31
32impl Display for CM31 {
33    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
34        write!(f, "{} + {}i", self.0, self.1)
35    }
36}
37
38impl Debug for CM31 {
39    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
40        write!(f, "{} + {}i", self.0, self.1)
41    }
42}
43
44impl Mul for CM31 {
45    type Output = Self;
46
47    fn mul(self, rhs: Self) -> Self::Output {
48        // (a + bi) * (c + di) = (ac - bd) + (ad + bc)i.
49        Self(
50            self.0 * rhs.0 - self.1 * rhs.1,
51            self.0 * rhs.1 + self.1 * rhs.0,
52        )
53    }
54}
55
56impl TryInto<M31> for CM31 {
57    type Error = ();
58
59    fn try_into(self) -> Result<M31, Self::Error> {
60        if self.1 != M31::zero() {
61            return Err(());
62        }
63        Ok(self.0)
64    }
65}
66
67impl FieldExpOps for CM31 {
68    fn inverse(&self) -> Self {
69        assert!(!self.is_zero(), "0 has no inverse");
70        // 1 / (a + bi) = (a - bi) / (a^2 + b^2).
71        Self(self.0, -self.1) * (self.0.square() + self.1.square()).inverse()
72    }
73}
74
75#[cfg(test)]
76#[macro_export]
77macro_rules! cm31 {
78    ($m0:expr, $m1:expr) => {
79        CM31::from_u32_unchecked($m0, $m1)
80    };
81}
82
83#[cfg(test)]
84mod tests {
85    use super::CM31;
86    use crate::core::fields::m31::P;
87    use crate::core::fields::FieldExpOps;
88    use crate::m31;
89
90    #[test]
91    fn test_inverse() {
92        let cm = cm31!(1, 2);
93        let cm_inv = cm.inverse();
94        assert_eq!(cm * cm_inv, cm31!(1, 0));
95    }
96
97    #[test]
98    fn test_ops() {
99        let cm0 = cm31!(1, 2);
100        let cm1 = cm31!(4, 5);
101        let m = m31!(8);
102        let cm = CM31::from(m);
103        let cm0_x_cm1 = cm31!(P - 6, 13);
104
105        assert_eq!(cm0 + cm1, cm31!(5, 7));
106        assert_eq!(cm1 + m, cm1 + cm);
107        assert_eq!(cm0 * cm1, cm0_x_cm1);
108        assert_eq!(cm1 * m, cm1 * cm);
109        assert_eq!(-cm0, cm31!(P - 1, P - 2));
110        assert_eq!(cm0 - cm1, cm31!(P - 3, P - 3));
111        assert_eq!(cm1 - m, cm1 - cm);
112        assert_eq!(cm0_x_cm1 / cm1, cm31!(1, 2));
113        assert_eq!(cm1 / m, cm1 / cm);
114    }
115}
stwo_gpu/core/fields/cm31.rs

stwo_gpu/core/fields/
cm31.rs
