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

p3_mersenne_31/
mersenne_31.rs
