p3_koala_bear/
koala_bear.rs

1use p3_field::{exp_1420470955, exp_u64_by_squaring, Field, FieldAlgebra};
2use p3_monty_31::{
3    BarrettParameters, BinomialExtensionData, FieldParameters, MontyField31, MontyParameters,
4    PackedMontyParameters, TwoAdicData,
5};
6
7/// The prime field `2^31 - 2^24 + 1`, a.k.a. the Koala Bear field.
8pub type KoalaBear = MontyField31<KoalaBearParameters>;
9
10#[derive(Copy, Clone, Default, Debug, Eq, Hash, PartialEq)]
11pub struct KoalaBearParameters;
12
13impl MontyParameters for KoalaBearParameters {
14    /// The KoalaBear prime: 2^31 - 2^24 + 1
15    /// This is a 31-bit prime with the highest possible two adicity if we additionally demand that
16    /// the cube map (x -> x^3) is an automorphism of the multiplicative group.
17    /// It's not unique, as there is one other option with equal 2 adicity: 2^30 + 2^27 + 2^24 + 1.
18    /// There is also one 29-bit prime with higher two adicity which might be appropriate for some applications: 2^29 - 2^26 + 1.
19    const PRIME: u32 = 0x7f000001;
20
21    const MONTY_BITS: u32 = 32;
22    const MONTY_MU: u32 = 0x81000001;
23}
24
25impl PackedMontyParameters for KoalaBearParameters {}
26
27impl BarrettParameters for KoalaBearParameters {}
28
29impl FieldParameters for KoalaBearParameters {
30    const MONTY_GEN: KoalaBear = KoalaBear::new(3);
31
32    fn exp_u64_generic<FA: FieldAlgebra>(val: FA, power: u64) -> FA {
33        match power {
34            1420470955 => exp_1420470955(val), // used to compute x^{1/7}
35            _ => exp_u64_by_squaring(val, power),
36        }
37    }
38
39    fn try_inverse<F: Field>(p1: F) -> Option<F> {
40        if p1.is_zero() {
41            return None;
42        }
43
44        // From Fermat's little theorem, in a prime field `F_p`, the inverse of `a` is `a^(p-2)`.
45        // Here p-2 = 2130706431 = 1111110111111111111111111111111_2
46        // Uses 29 Squares + 7 Multiplications => 36 Operations total.
47
48        let p10 = p1.square();
49        let p11 = p10 * p1;
50        let p1100 = p11.exp_power_of_2(2);
51        let p1111 = p1100 * p11;
52        let p110000 = p1100.exp_power_of_2(2);
53        let p111111 = p110000 * p1111;
54        let p1111110000 = p111111.exp_power_of_2(4);
55        let p1111111111 = p1111110000 * p1111;
56        let p11111101111 = p1111111111 * p1111110000;
57        let p111111011110000000000 = p11111101111.exp_power_of_2(10);
58        let p111111011111111111111 = p111111011110000000000 * p1111111111;
59        let p1111110111111111111110000000000 = p111111011111111111111.exp_power_of_2(10);
60        let p1111110111111111111111111111111 = p1111110111111111111110000000000 * p1111111111;
61
62        Some(p1111110111111111111111111111111)
63    }
64}
65
66impl TwoAdicData for KoalaBearParameters {
67    const TWO_ADICITY: usize = 24;
68
69    type ArrayLike = &'static [KoalaBear];
70
71    const TWO_ADIC_GENERATORS: Self::ArrayLike = &KoalaBear::new_array([
72        0x1, 0x7f000000, 0x7e010002, 0x6832fe4a, 0x8dbd69c, 0xa28f031, 0x5c4a5b99, 0x29b75a80,
73        0x17668b8a, 0x27ad539b, 0x334d48c7, 0x7744959c, 0x768fc6fa, 0x303964b2, 0x3e687d4d,
74        0x45a60e61, 0x6e2f4d7a, 0x163bd499, 0x6c4a8a45, 0x143ef899, 0x514ddcad, 0x484ef19b,
75        0x205d63c3, 0x68e7dd49, 0x6ac49f88,
76    ]);
77
78    const ROOTS_8: Self::ArrayLike =
79        &KoalaBear::new_array([0x1, 0x6832fe4a, 0x7e010002, 0x174e3650]);
80    const INV_ROOTS_8: Self::ArrayLike =
81        &KoalaBear::new_array([0x1, 0x67b1c9b1, 0xfeffff, 0x16cd01b7]);
82
83    const ROOTS_16: Self::ArrayLike = &KoalaBear::new_array([
84        0x1, 0x8dbd69c, 0x6832fe4a, 0x27ae21e2, 0x7e010002, 0x3a89a025, 0x174e3650, 0x27dfce22,
85    ]);
86    const INV_ROOTS_16: Self::ArrayLike = &KoalaBear::new_array([
87        0x1, 0x572031df, 0x67b1c9b1, 0x44765fdc, 0xfeffff, 0x5751de1f, 0x16cd01b7, 0x76242965,
88    ]);
89}
90
91impl BinomialExtensionData<4> for KoalaBearParameters {
92    const W: KoalaBear = KoalaBear::new(3);
93    const DTH_ROOT: KoalaBear = KoalaBear::new(2113994754);
94    const EXT_GENERATOR: [KoalaBear; 4] = KoalaBear::new_array([2, 1, 0, 0]);
95    const EXT_TWO_ADICITY: usize = 26;
96
97    type ArrayLike = [[KoalaBear; 4]; 2];
98
99    const TWO_ADIC_EXTENSION_GENERATORS: Self::ArrayLike =
100        KoalaBear::new_2d_array([[0, 0, 1759267465, 0], [0, 0, 0, 777715144]]);
101}
102
103#[cfg(test)]
104mod tests {
105    use p3_field::{PrimeField32, PrimeField64, TwoAdicField};
106    use p3_field_testing::{test_field, test_field_dft, test_two_adic_field};
107
108    use super::*;
109
110    type F = KoalaBear;
111
112    #[test]
113    fn test_koala_bear_two_adicity_generators() {
114        let base = KoalaBear::from_canonical_u32(0x6ac49f88);
115        for bits in 0..=KoalaBear::TWO_ADICITY {
116            assert_eq!(
117                KoalaBear::two_adic_generator(bits),
118                base.exp_power_of_2(KoalaBear::TWO_ADICITY - bits)
119            );
120        }
121    }
122
123    #[test]
124    fn test_koala_bear() {
125        let f = F::from_canonical_u32(100);
126        assert_eq!(f.as_canonical_u64(), 100);
127
128        let f = F::from_canonical_u32(0);
129        assert!(f.is_zero());
130
131        let f = F::from_wrapped_u32(F::ORDER_U32);
132        assert!(f.is_zero());
133
134        let f_1 = F::ONE;
135        let f_1_copy = F::from_canonical_u32(1);
136
137        let expected_result = F::ZERO;
138        assert_eq!(f_1 - f_1_copy, expected_result);
139
140        let expected_result = F::TWO;
141        assert_eq!(f_1 + f_1_copy, expected_result);
142
143        let f_2 = F::from_canonical_u32(2);
144        let expected_result = F::from_canonical_u32(3);
145        assert_eq!(f_1 + f_1_copy * f_2, expected_result);
146
147        let expected_result = F::from_canonical_u32(5);
148        assert_eq!(f_1 + f_2 * f_2, expected_result);
149
150        let f_p_minus_1 = F::from_canonical_u32(F::ORDER_U32 - 1);
151        let expected_result = F::ZERO;
152        assert_eq!(f_1 + f_p_minus_1, expected_result);
153
154        let f_p_minus_2 = F::from_canonical_u32(F::ORDER_U32 - 2);
155        let expected_result = F::from_canonical_u32(F::ORDER_U32 - 3);
156        assert_eq!(f_p_minus_1 + f_p_minus_2, expected_result);
157
158        let expected_result = F::from_canonical_u32(1);
159        assert_eq!(f_p_minus_1 - f_p_minus_2, expected_result);
160
161        let expected_result = f_p_minus_1;
162        assert_eq!(f_p_minus_2 - f_p_minus_1, expected_result);
163
164        let expected_result = f_p_minus_2;
165        assert_eq!(f_p_minus_1 - f_1, expected_result);
166
167        let m1 = F::from_canonical_u32(0x34167c58);
168        let m2 = F::from_canonical_u32(0x61f3207b);
169        let expected_prod = F::from_canonical_u32(0x54b46b81);
170        assert_eq!(m1 * m2, expected_prod);
171
172        assert_eq!(m1.exp_u64(1420470955).exp_const_u64::<3>(), m1);
173        assert_eq!(m2.exp_u64(1420470955).exp_const_u64::<3>(), m2);
174        assert_eq!(f_2.exp_u64(1420470955).exp_const_u64::<3>(), f_2);
175
176        let f_serialized = serde_json::to_string(&f).unwrap();
177        let f_deserialized: F = serde_json::from_str(&f_serialized).unwrap();
178        assert_eq!(f, f_deserialized);
179
180        let f_1_serialized = serde_json::to_string(&f_1).unwrap();
181        let f_1_deserialized: F = serde_json::from_str(&f_1_serialized).unwrap();
182        let f_1_serialized_again = serde_json::to_string(&f_1_deserialized).unwrap();
183        let f_1_deserialized_again: F = serde_json::from_str(&f_1_serialized_again).unwrap();
184        assert_eq!(f_1, f_1_deserialized);
185        assert_eq!(f_1, f_1_deserialized_again);
186
187        let f_2_serialized = serde_json::to_string(&f_2).unwrap();
188        let f_2_deserialized: F = serde_json::from_str(&f_2_serialized).unwrap();
189        assert_eq!(f_2, f_2_deserialized);
190
191        let f_p_minus_1_serialized = serde_json::to_string(&f_p_minus_1).unwrap();
192        let f_p_minus_1_deserialized: F = serde_json::from_str(&f_p_minus_1_serialized).unwrap();
193        assert_eq!(f_p_minus_1, f_p_minus_1_deserialized);
194
195        let f_p_minus_2_serialized = serde_json::to_string(&f_p_minus_2).unwrap();
196        let f_p_minus_2_deserialized: F = serde_json::from_str(&f_p_minus_2_serialized).unwrap();
197        assert_eq!(f_p_minus_2, f_p_minus_2_deserialized);
198
199        let m1_serialized = serde_json::to_string(&m1).unwrap();
200        let m1_deserialized: F = serde_json::from_str(&m1_serialized).unwrap();
201        assert_eq!(m1, m1_deserialized);
202
203        let m2_serialized = serde_json::to_string(&m2).unwrap();
204        let m2_deserialized: F = serde_json::from_str(&m2_serialized).unwrap();
205        assert_eq!(m2, m2_deserialized);
206    }
207
208    test_field!(crate::KoalaBear);
209    test_two_adic_field!(crate::KoalaBear);
210
211    test_field_dft!(radix2dit, crate::KoalaBear, p3_dft::Radix2Dit<_>);
212    test_field_dft!(bowers, crate::KoalaBear, p3_dft::Radix2Bowers);
213    test_field_dft!(
214        parallel,
215        crate::KoalaBear,
216        p3_dft::Radix2DitParallel::<crate::KoalaBear>
217    );
218    test_field_dft!(
219        recur_dft,
220        crate::KoalaBear,
221        p3_monty_31::dft::RecursiveDft<_>
222    );
223}