p3_mersenne_31/
mersenne_31.rs

1use alloc::vec;
2use alloc::vec::Vec;
3use core::fmt::{Debug, Display, Formatter};
4use core::hash::{Hash, Hasher};
5use core::iter::{Product, Sum};
6use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
7use core::{array, fmt, iter};
8
9use num_bigint::BigUint;
10use p3_field::exponentiation::exp_1717986917;
11use p3_field::integers::QuotientMap;
12use p3_field::op_assign_macros::{
13    impl_add_assign, impl_div_methods, impl_mul_methods, impl_sub_assign,
14};
15use p3_field::{
16    Field, InjectiveMonomial, Packable, PermutationMonomial, PrimeCharacteristicRing, PrimeField,
17    PrimeField32, PrimeField64, RawDataSerializable, halve_u32, impl_raw_serializable_primefield32,
18    quotient_map_large_iint, quotient_map_large_uint, quotient_map_small_int,
19};
20use p3_util::{flatten_to_base, gcd_inversion_prime_field_32};
21use rand::Rng;
22use rand::distr::{Distribution, StandardUniform};
23use serde::de::Error;
24use serde::{Deserialize, Deserializer, Serialize};
25
26/// The Mersenne31 prime
27const P: u32 = (1 << 31) - 1;
28
29/// The prime field `F_p` where `p = 2^31 - 1`.
30#[derive(Copy, Clone, Default)]
31#[repr(transparent)] // Important for reasoning about memory layout.
32#[must_use]
33pub struct Mersenne31 {
34    /// Not necessarily canonical, but must fit in 31 bits.
35    pub(crate) value: u32,
36}
37
38impl Mersenne31 {
39    /// Convert a u32 element into a Mersenne31 element.
40    ///
41    /// # Safety
42    /// The element must lie in the range: `[0, 2^31 - 1]`.
43    #[inline]
44    pub(crate) const fn new(value: u32) -> Self {
45        debug_assert!((value >> 31) == 0);
46        Self { value }
47    }
48
49    /// Convert a u32 element into a Mersenne31 element.
50    ///
51    /// # Panics
52    /// This will panic if the element does not lie in the range: `[0, 2^31 - 1]`.
53    #[inline]
54    pub const fn new_checked(value: u32) -> Option<Self> {
55        if (value >> 31) == 0 {
56            Some(Self { value })
57        } else {
58            None
59        }
60    }
61
62    /// Convert a constant `u32` array into a constant array of field elements.
63    /// This allows inputs to be `> 2^31`, and just reduces them `mod P`.
64    ///
65    /// This means that this will be slower than `array.map(Mersenne31::new_checked)` but
66    /// has the advantage of being able to be used in `const` environments.
67    #[inline]
68    pub const fn new_array<const N: usize>(input: [u32; N]) -> [Self; N] {
69        let mut output = [Self::ZERO; N];
70        let mut i = 0;
71        while i < N {
72            output[i].value = input[i] % P;
73            i += 1;
74        }
75        output
76    }
77
78    /// Precomputed table of generators for two-adic subgroups of the degree two extension field over Mersenne31.
79    /// The `i`'th element is a generator of the subgroup of order `2^i`.
80    pub const EXT_TWO_ADIC_GENERATORS: [[Self; 2]; 33] = [
81        [Self::ONE, Self::ZERO],
82        [Self::new(2_147_483_646), Self::new(0)],
83        [Self::new(0), Self::new(2_147_483_646)],
84        [Self::new(32_768), Self::new(2_147_450_879)],
85        [Self::new(590_768_354), Self::new(978_592_373)],
86        [Self::new(1_179_735_656), Self::new(1_241_207_368)],
87        [Self::new(1_567_857_810), Self::new(456_695_729)],
88        [Self::new(1_774_253_895), Self::new(1_309_288_441)],
89        [Self::new(736_262_640), Self::new(1_553_669_210)],
90        [Self::new(1_819_216_575), Self::new(1_662_816_114)],
91        [Self::new(1_323_191_254), Self::new(1_936_974_060)],
92        [Self::new(605_622_498), Self::new(1_964_232_216)],
93        [Self::new(343_674_985), Self::new(501_786_993)],
94        [Self::new(1_995_316_534), Self::new(149_306_621)],
95        [Self::new(2_107_600_913), Self::new(1_378_821_388)],
96        [Self::new(541_476_169), Self::new(2_101_081_972)],
97        [Self::new(2_135_874_973), Self::new(483_411_332)],
98        [Self::new(2_097_144_245), Self::new(1_684_033_590)],
99        [Self::new(1_662_322_247), Self::new(670_236_780)],
100        [Self::new(1_172_215_635), Self::new(595_888_646)],
101        [Self::new(241_940_101), Self::new(323_856_519)],
102        [Self::new(1_957_194_259), Self::new(2_139_647_100)],
103        [Self::new(1_957_419_629), Self::new(1_541_039_442)],
104        [Self::new(1_062_045_235), Self::new(1_824_580_421)],
105        [Self::new(1_929_382_196), Self::new(1_664_698_822)],
106        [Self::new(1_889_294_251), Self::new(331_248_939)],
107        [Self::new(1_214_231_414), Self::new(1_646_302_518)],
108        [Self::new(1_765_392_370), Self::new(461_136_547)],
109        [Self::new(1_629_751_483), Self::new(66_485_474)],
110        [Self::new(1_501_355_827), Self::new(1_439_063_420)],
111        [Self::new(509_778_402), Self::new(800_467_507)],
112        [Self::new(311_014_874), Self::new(1_584_694_829)],
113        [Self::new(1_166_849_849), Self::new(1_117_296_306)],
114    ];
115}
116
117impl PartialEq for Mersenne31 {
118    #[inline]
119    fn eq(&self, other: &Self) -> bool {
120        self.as_canonical_u32() == other.as_canonical_u32()
121    }
122}
123
124impl Eq for Mersenne31 {}
125
126impl Packable for Mersenne31 {}
127
128impl Hash for Mersenne31 {
129    fn hash<H: Hasher>(&self, state: &mut H) {
130        state.write_u32(self.to_unique_u32());
131    }
132}
133
134impl Ord for Mersenne31 {
135    #[inline]
136    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
137        self.as_canonical_u32().cmp(&other.as_canonical_u32())
138    }
139}
140
141impl PartialOrd for Mersenne31 {
142    #[inline]
143    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
144        Some(self.cmp(other))
145    }
146}
147
148impl Display for Mersenne31 {
149    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
150        Display::fmt(&self.value, f)
151    }
152}
153
154impl Debug for Mersenne31 {
155    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
156        Debug::fmt(&self.value, f)
157    }
158}
159
160impl Distribution<Mersenne31> for StandardUniform {
161    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Mersenne31 {
162        loop {
163            let next_u31 = rng.next_u32() >> 1;
164            let is_canonical = next_u31 != Mersenne31::ORDER_U32;
165            if is_canonical {
166                return Mersenne31::new(next_u31);
167            }
168        }
169    }
170}
171
172impl Serialize for Mersenne31 {
173    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
174        // No need to convert to canonical.
175        serializer.serialize_u32(self.value)
176    }
177}
178
179impl<'a> Deserialize<'a> for Mersenne31 {
180    fn deserialize<D: Deserializer<'a>>(d: D) -> Result<Self, D::Error> {
181        let val = u32::deserialize(d)?;
182        // Ensure that `val` satisfies our invariant. i.e. Not necessarily canonical, but must fit in 31 bits.
183        if val <= P {
184            Ok(Self::new(val))
185        } else {
186            Err(D::Error::custom("Value is out of range"))
187        }
188    }
189}
190
191impl RawDataSerializable for Mersenne31 {
192    impl_raw_serializable_primefield32!();
193}
194
195impl PrimeCharacteristicRing for Mersenne31 {
196    type PrimeSubfield = Self;
197
198    const ZERO: Self = Self { value: 0 };
199    const ONE: Self = Self { value: 1 };
200    const TWO: Self = Self { value: 2 };
201    const NEG_ONE: Self = Self {
202        value: Self::ORDER_U32 - 1,
203    };
204
205    #[inline]
206    fn from_prime_subfield(f: Self::PrimeSubfield) -> Self {
207        f
208    }
209
210    #[inline]
211    fn from_bool(b: bool) -> Self {
212        Self::new(b as u32)
213    }
214
215    #[inline]
216    fn halve(&self) -> Self {
217        Self::new(halve_u32::<P>(self.value))
218    }
219
220    #[inline]
221    fn mul_2exp_u64(&self, exp: u64) -> Self {
222        // In a Mersenne field, multiplication by 2^k is just a left rotation by k bits.
223        let exp = exp % 31;
224        let left = (self.value << exp) & ((1 << 31) - 1);
225        let right = self.value >> (31 - exp);
226        let rotated = left | right;
227        Self::new(rotated)
228    }
229
230    #[inline]
231    fn div_2exp_u64(&self, exp: u64) -> Self {
232        // In a Mersenne field, division by 2^k is just a right rotation by k bits.
233        let exp = (exp % 31) as u8;
234        let left = self.value >> exp;
235        let right = (self.value << (31 - exp)) & ((1 << 31) - 1);
236        let rotated = left | right;
237        Self::new(rotated)
238    }
239
240    #[inline]
241    fn sum_array<const N: usize>(input: &[Self]) -> Self {
242        assert_eq!(N, input.len());
243        // Benchmarking shows that for N <= 5 it's faster to sum the elements directly
244        // but for N > 5 it's faster to use the .sum() methods which passes through u64's
245        // allowing for delayed reductions.
246        match N {
247            0 => Self::ZERO,
248            1 => input[0],
249            2 => input[0] + input[1],
250            3 => input[0] + input[1] + input[2],
251            4 => (input[0] + input[1]) + (input[2] + input[3]),
252            5 => {
253                let lhs = input[0] + input[1];
254                let rhs = input[2] + input[3];
255                lhs + rhs + input[4]
256            }
257            _ => input.iter().copied().sum(),
258        }
259    }
260
261    #[inline]
262    fn zero_vec(len: usize) -> Vec<Self> {
263        // SAFETY:
264        // Due to `#[repr(transparent)]`, Mersenne31 and u32 have the same size, alignment
265        // and memory layout making `flatten_to_base` safe. This this will create
266        // a vector Mersenne31 elements with value set to 0.
267        unsafe { flatten_to_base(vec![0u32; len]) }
268    }
269}
270
271// Degree of the smallest permutation polynomial for Mersenne31.
272//
273// As p - 1 = 2×3^2×7×11×... the smallest choice for a degree D satisfying gcd(p - 1, D) = 1 is 5.
274impl InjectiveMonomial<5> for Mersenne31 {}
275
276impl PermutationMonomial<5> for Mersenne31 {
277    /// In the field `Mersenne31`, `a^{1/5}` is equal to a^{1717986917}.
278    ///
279    /// This follows from the calculation `5 * 1717986917 = 4*(2^31 - 2) + 1 = 1 mod p - 1`.
280    fn injective_exp_root_n(&self) -> Self {
281        exp_1717986917(*self)
282    }
283}
284
285impl Field for Mersenne31 {
286    #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
287    type Packing = crate::PackedMersenne31Neon;
288    #[cfg(all(
289        target_arch = "x86_64",
290        target_feature = "avx2",
291        not(target_feature = "avx512f")
292    ))]
293    type Packing = crate::PackedMersenne31AVX2;
294    #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
295    type Packing = crate::PackedMersenne31AVX512;
296    #[cfg(not(any(
297        all(target_arch = "aarch64", target_feature = "neon"),
298        all(
299            target_arch = "x86_64",
300            target_feature = "avx2",
301            not(target_feature = "avx512f")
302        ),
303        all(target_arch = "x86_64", target_feature = "avx512f"),
304    )))]
305    type Packing = Self;
306
307    // Sage: GF(2^31 - 1).multiplicative_generator()
308    const GENERATOR: Self = Self::new(7);
309
310    #[inline]
311    fn is_zero(&self) -> bool {
312        self.value == 0 || self.value == Self::ORDER_U32
313    }
314
315    fn try_inverse(&self) -> Option<Self> {
316        if self.is_zero() {
317            return None;
318        }
319
320        // Number of bits in the Mersenne31 prime.
321        const NUM_PRIME_BITS: u32 = 31;
322
323        // gcd_inversion returns the inverse multiplied by 2^60 so we need to correct for that.
324        let inverse_i64 = gcd_inversion_prime_field_32::<NUM_PRIME_BITS>(self.value, P);
325        Some(Self::from_int(inverse_i64).div_2exp_u64(60))
326    }
327
328    #[inline]
329    fn order() -> BigUint {
330        P.into()
331    }
332}
333
334// We can use some macros to implement QuotientMap<Int> for all integer types except for u32 and i32's.
335quotient_map_small_int!(Mersenne31, u32, [u8, u16]);
336quotient_map_small_int!(Mersenne31, i32, [i8, i16]);
337quotient_map_large_uint!(
338    Mersenne31,
339    u32,
340    Mersenne31::ORDER_U32,
341    "`[0, 2^31 - 2]`",
342    "`[0, 2^31 - 1]`",
343    [u64, u128]
344);
345quotient_map_large_iint!(
346    Mersenne31,
347    i32,
348    "`[-2^30, 2^30]`",
349    "`[1 - 2^31, 2^31 - 1]`",
350    [(i64, u64), (i128, u128)]
351);
352
353// We simple need to prove custom Mersenne31 impls for QuotientMap<u32> and QuotientMap<i32>
354impl QuotientMap<u32> for Mersenne31 {
355    /// Convert a given `u32` integer into an element of the `Mersenne31` field.
356    #[inline]
357    fn from_int(int: u32) -> Self {
358        // To reduce `n` to 31 bits, we clear its MSB, then add it back in its reduced form.
359        let msb = int & (1 << 31);
360        let msb_reduced = msb >> 31;
361        Self::new(int ^ msb) + Self::new(msb_reduced)
362    }
363
364    /// Convert a given `u32` integer into an element of the `Mersenne31` field.
365    ///
366    /// Returns none if the input does not lie in the range `[0, 2^31 - 1]`.
367    #[inline]
368    fn from_canonical_checked(int: u32) -> Option<Self> {
369        (int < Self::ORDER_U32).then(|| Self::new(int))
370    }
371
372    /// Convert a given `u32` integer into an element of the `Mersenne31` field.
373    ///
374    /// # Safety
375    /// The input must lie in the range: `[0, 2^31 - 1]`.
376    #[inline(always)]
377    unsafe fn from_canonical_unchecked(int: u32) -> Self {
378        debug_assert!(int < Self::ORDER_U32);
379        Self::new(int)
380    }
381}
382
383impl QuotientMap<i32> for Mersenne31 {
384    /// Convert a given `i32` integer into an element of the `Mersenne31` field.
385    #[inline]
386    fn from_int(int: i32) -> Self {
387        if int >= 0 {
388            Self::new(int as u32)
389        } else if int > (-1 << 31) {
390            Self::new(Self::ORDER_U32.wrapping_add_signed(int))
391        } else {
392            // The only other option is int = -(2^31) = -1 mod p.
393            Self::NEG_ONE
394        }
395    }
396
397    /// Convert a given `i32` integer into an element of the `Mersenne31` field.
398    ///
399    /// Returns none if the input does not lie in the range `(-2^30, 2^30)`.
400    #[inline]
401    fn from_canonical_checked(int: i32) -> Option<Self> {
402        const TWO_EXP_30: i32 = 1 << 30;
403        const NEG_TWO_EXP_30_PLUS_1: i32 = (-1 << 30) + 1;
404        match int {
405            0..TWO_EXP_30 => Some(Self::new(int as u32)),
406            NEG_TWO_EXP_30_PLUS_1..0 => Some(Self::new(Self::ORDER_U32.wrapping_add_signed(int))),
407            _ => None,
408        }
409    }
410
411    /// Convert a given `i32` integer into an element of the `Mersenne31` field.
412    ///
413    /// # Safety
414    /// The input must lie in the range: `[1 - 2^31, 2^31 - 1]`.
415    #[inline(always)]
416    unsafe fn from_canonical_unchecked(int: i32) -> Self {
417        if int >= 0 {
418            Self::new(int as u32)
419        } else {
420            Self::new(Self::ORDER_U32.wrapping_add_signed(int))
421        }
422    }
423}
424
425impl PrimeField for Mersenne31 {
426    fn as_canonical_biguint(&self) -> BigUint {
427        <Self as PrimeField32>::as_canonical_u32(self).into()
428    }
429}
430
431impl PrimeField32 for Mersenne31 {
432    const ORDER_U32: u32 = P;
433
434    #[inline]
435    fn as_canonical_u32(&self) -> u32 {
436        // Since our invariant guarantees that `value` fits in 31 bits, there is only one possible
437        // `value` that is not canonical, namely 2^31 - 1 = p = 0.
438        if self.value == Self::ORDER_U32 {
439            0
440        } else {
441            self.value
442        }
443    }
444}
445
446impl PrimeField64 for Mersenne31 {
447    const ORDER_U64: u64 = <Self as PrimeField32>::ORDER_U32 as u64;
448
449    #[inline]
450    fn as_canonical_u64(&self) -> u64 {
451        self.as_canonical_u32().into()
452    }
453}
454
455impl Add for Mersenne31 {
456    type Output = Self;
457
458    #[inline]
459    fn add(self, rhs: Self) -> Self {
460        // See the following for a way to compute the sum that avoids
461        // the conditional which may be preferable on some
462        // architectures.
463        // https://github.com/Plonky3/Plonky3/blob/6049a30c3b1f5351c3eb0f7c994dc97e8f68d10d/mersenne-31/src/lib.rs#L249
464
465        // Working with i32 means we get a flag which informs us if overflow happened.
466        let (sum_i32, over) = (self.value as i32).overflowing_add(rhs.value as i32);
467        let sum_u32 = sum_i32 as u32;
468        let sum_corr = sum_u32.wrapping_sub(Self::ORDER_U32);
469
470        // If self + rhs did not overflow, return it.
471        // If self + rhs overflowed, sum_corr = self + rhs - (2**31 - 1).
472        Self::new(if over { sum_corr } else { sum_u32 })
473    }
474}
475
476impl Sub for Mersenne31 {
477    type Output = Self;
478
479    #[inline]
480    fn sub(self, rhs: Self) -> Self {
481        let (mut sub, over) = self.value.overflowing_sub(rhs.value);
482
483        // If we didn't overflow we have the correct value.
484        // Otherwise we have added 2**32 = 2**31 + 1 mod 2**31 - 1.
485        // Hence we need to remove the most significant bit and subtract 1.
486        sub -= over as u32;
487        Self::new(sub & Self::ORDER_U32)
488    }
489}
490
491impl Neg for Mersenne31 {
492    type Output = Self;
493
494    #[inline]
495    fn neg(self) -> Self::Output {
496        // Can't underflow, since self.value is 31-bits and thus can't exceed ORDER.
497        Self::new(Self::ORDER_U32 - self.value)
498    }
499}
500
501impl Mul for Mersenne31 {
502    type Output = Self;
503
504    #[inline]
505    #[allow(clippy::cast_possible_truncation)]
506    fn mul(self, rhs: Self) -> Self {
507        let prod = u64::from(self.value) * u64::from(rhs.value);
508        from_u62(prod)
509    }
510}
511
512impl_add_assign!(Mersenne31);
513impl_sub_assign!(Mersenne31);
514impl_mul_methods!(Mersenne31);
515impl_div_methods!(Mersenne31, Mersenne31);
516
517impl Sum for Mersenne31 {
518    #[inline]
519    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
520        // This is faster than iter.reduce(|x, y| x + y).unwrap_or(Self::ZERO) for iterators of length >= 6.
521        // It assumes that iter.len() < 2^31.
522
523        // This sum will not overflow so long as iter.len() < 2^33.
524        let sum = iter.map(|x| x.value as u64).sum::<u64>();
525
526        // sum is < 2^62 provided iter.len() < 2^31.
527        from_u62(sum)
528    }
529}
530
531#[inline(always)]
532pub(crate) fn from_u62(input: u64) -> Mersenne31 {
533    debug_assert!(input < (1 << 62));
534    let input_lo = (input & ((1 << 31) - 1)) as u32;
535    let input_high = (input >> 31) as u32;
536    Mersenne31::new(input_lo) + Mersenne31::new(input_high)
537}
538
539#[cfg(test)]
540mod tests {
541    use num_bigint::BigUint;
542    use p3_field::{InjectiveMonomial, PermutationMonomial, PrimeCharacteristicRing};
543    use p3_field_testing::{
544        test_field, test_prime_field, test_prime_field_32, test_prime_field_64,
545    };
546
547    use crate::Mersenne31;
548
549    type F = Mersenne31;
550
551    #[test]
552    fn exp_root() {
553        // Confirm that (x^{1/5})^5 = x
554
555        let m1 = F::from_u32(0x34167c58);
556        let m2 = F::from_u32(0x61f3207b);
557
558        assert_eq!(m1.injective_exp_n().injective_exp_root_n(), m1);
559        assert_eq!(m2.injective_exp_n().injective_exp_root_n(), m2);
560        assert_eq!(F::TWO.injective_exp_n().injective_exp_root_n(), F::TWO);
561    }
562
563    // Mersenne31 has a redundant representation of Zero but no redundant representation of One.
564    const ZEROS: [Mersenne31; 2] = [Mersenne31::ZERO, Mersenne31::new((1_u32 << 31) - 1)];
565    const ONES: [Mersenne31; 1] = [Mersenne31::ONE];
566
567    // Get the prime factorization of the order of the multiplicative group.
568    // i.e. the prime factorization of P - 1.
569    fn multiplicative_group_prime_factorization() -> [(BigUint, u32); 7] {
570        [
571            (BigUint::from(2u8), 1),
572            (BigUint::from(3u8), 2),
573            (BigUint::from(7u8), 1),
574            (BigUint::from(11u8), 1),
575            (BigUint::from(31u8), 1),
576            (BigUint::from(151u8), 1),
577            (BigUint::from(331u16), 1),
578        ]
579    }
580
581    test_field!(
582        crate::Mersenne31,
583        &super::ZEROS,
584        &super::ONES,
585        &super::multiplicative_group_prime_factorization()
586    );
587    test_prime_field!(crate::Mersenne31);
588    test_prime_field_64!(crate::Mersenne31, &super::ZEROS, &super::ONES);
589    test_prime_field_32!(crate::Mersenne31, &super::ZEROS, &super::ONES);
590}
p3_mersenne_31/mersenne_31.rs

p3_mersenne_31/
mersenne_31.rs
