Skip to main content

p3_monty_31/
monty_31.rs

1//! An abstraction of 31-bit fields which use a MONTY approach for faster multiplication.
2
3use alloc::vec;
4use alloc::vec::Vec;
5use core::fmt::{self, Debug, Display, Formatter};
6use core::hash::Hash;
7use core::iter::{Product, Sum};
8use core::marker::PhantomData;
9use core::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign};
10use core::{array, iter};
11
12use num_bigint::BigUint;
13use p3_field::integers::QuotientMap;
14use p3_field::op_assign_macros::{
15    impl_add_assign, impl_div_methods, impl_mul_methods, impl_sub_assign,
16};
17use p3_field::{
18    Field, InjectiveMonomial, Packable, PermutationMonomial, PrimeCharacteristicRing, PrimeField,
19    PrimeField32, PrimeField64, RawDataSerializable, TwoAdicField, UniformSamplingField,
20    impl_raw_serializable_primefield32, quotient_map_small_int, tonelli_shanks_two_adic,
21};
22use p3_util::{flatten_to_base, gcd_inversion_prime_field_32};
23use rand::Rng;
24use rand::distr::{Distribution, StandardUniform};
25use serde::de::Error;
26use serde::{Deserialize, Deserializer, Serialize};
27
28use crate::utils::{
29    add, from_monty, halve_u32, large_monty_reduce, monty_reduce, monty_reduce_u128, sub, to_monty,
30    to_monty_64, to_monty_64_signed, to_monty_signed,
31};
32use crate::{FieldParameters, MontyParameters, RelativelyPrimePower, TwoAdicData};
33
34#[derive(Clone, Copy, Default, Eq, Hash, PartialEq)]
35#[repr(transparent)] // Important for reasoning about memory layout.
36#[must_use]
37pub struct MontyField31<MP: MontyParameters> {
38    /// The MONTY form of the field element, saved as a positive integer less than `P`.
39    ///
40    /// This is `pub(crate)` for tests and delayed reduction strategies. If you're accessing `value` outside of those, you're
41    /// likely doing something fishy.
42    pub(crate) value: u32,
43    _phantom: PhantomData<MP>,
44}
45
46impl<MP: MontyParameters> MontyField31<MP> {
47    /// Create a new field element from any `u32`.
48    ///
49    /// Any `u32` value is accepted and automatically converted to Montgomery form.
50    #[inline(always)]
51    pub const fn new(value: u32) -> Self {
52        const {
53            assert!(MP::PRIME % 2 == 1, "PRIME must be odd");
54            assert!(MP::PRIME < (1 << 31), "PRIME must be a 31-bit prime");
55            assert!(MP::MONTY_BITS == 32);
56            assert!(
57                MP::PRIME.wrapping_mul(MP::MONTY_MU) == 1,
58                "MONTY_MU must satisfy PRIME * MONTY_MU ≡ 1 (mod 2^32)"
59            );
60        }
61        Self {
62            value: to_monty::<MP>(value),
63            _phantom: PhantomData,
64        }
65    }
66
67    /// Create a new field element from something already in MONTY form.
68    /// This is `pub(crate)` for tests and delayed reduction strategies. If you're using it outside of those, you're
69    /// likely doing something fishy.
70    #[inline(always)]
71    pub(crate) const fn new_monty(value: u32) -> Self {
72        Self {
73            value,
74            _phantom: PhantomData,
75        }
76    }
77
78    /// Produce a u32 in range [0, P) from a field element corresponding to the true value.
79    #[inline(always)]
80    pub(crate) const fn to_u32(elem: &Self) -> u32 {
81        from_monty::<MP>(elem.value)
82    }
83
84    /// Convert a `[u32; N]` array to an array of field elements.
85    ///
86    /// Const version of `input.map(MontyField31::new)`.
87    #[inline]
88    pub const fn new_array<const N: usize>(input: [u32; N]) -> [Self; N] {
89        let mut output = [Self::new_monty(0); N];
90        let mut i = 0;
91        while i < N {
92            output[i] = Self::new(input[i]);
93            i += 1;
94        }
95        output
96    }
97
98    /// Convert a constant 2d u32 array into a constant 2d array of field elements.
99    /// Constant version of array.map(MontyField31::new_array).
100    #[inline]
101    pub const fn new_2d_array<const N: usize, const M: usize>(
102        input: [[u32; N]; M],
103    ) -> [[Self; N]; M] {
104        let mut output = [[Self::new_monty(0); N]; M];
105        let mut i = 0;
106        while i < M {
107            output[i] = Self::new_array(input[i]);
108            i += 1;
109        }
110        output
111    }
112}
113
114impl<FP: FieldParameters> MontyField31<FP> {
115    const MONTY_POWERS_OF_TWO: [Self; 64] = {
116        let mut powers_of_two = [FP::MONTY_ONE; 64];
117        let mut i = 1;
118        while i < 64 {
119            powers_of_two[i] = Self::new_monty(to_monty_64::<FP>(1 << i));
120            i += 1;
121        }
122        powers_of_two
123    };
124
125    const HALF: Self = Self::new(FP::HALF_P_PLUS_1);
126}
127
128impl<FP: MontyParameters> Ord for MontyField31<FP> {
129    #[inline]
130    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
131        Self::to_u32(self).cmp(&Self::to_u32(other))
132    }
133}
134
135impl<FP: MontyParameters> PartialOrd for MontyField31<FP> {
136    #[inline]
137    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
138        Some(self.cmp(other))
139    }
140}
141
142impl<FP: MontyParameters> Display for MontyField31<FP> {
143    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
144        Display::fmt(&Self::to_u32(self), f)
145    }
146}
147
148impl<FP: MontyParameters> Debug for MontyField31<FP> {
149    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
150        Debug::fmt(&Self::to_u32(self), f)
151    }
152}
153
154impl<FP: MontyParameters> Distribution<MontyField31<FP>> for StandardUniform {
155    #[inline]
156    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> MontyField31<FP> {
157        loop {
158            let next_u31 = rng.next_u32() >> 1;
159            let is_canonical = next_u31 < FP::PRIME;
160            if is_canonical {
161                return MontyField31::new_monty(next_u31);
162            }
163        }
164    }
165}
166
167impl<FP: FieldParameters> Serialize for MontyField31<FP> {
168    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
169        // It's faster to Serialize and Deserialize in monty form.
170        // Binary (non human-readable) formats get a fixed 4-byte encoding instead of
171        // the serializer's default varint, since every value here is a near-uniform
172        // 31-bit integer and varint saves nothing on average.
173        if serializer.is_human_readable() {
174            serializer.serialize_u32(self.value)
175        } else {
176            self.value.to_le_bytes().serialize(serializer)
177        }
178    }
179}
180
181impl<'de, FP: FieldParameters> Deserialize<'de> for MontyField31<FP> {
182    fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
183        // It's faster to Serialize and Deserialize in monty form.
184        let human_readable = d.is_human_readable();
185        let val = if human_readable {
186            u32::deserialize(d)?
187        } else {
188            u32::from_le_bytes(<[u8; 4]>::deserialize(d)?)
189        };
190        if val < FP::PRIME {
191            Ok(Self::new_monty(val))
192        } else {
193            Err(D::Error::custom("Value is out of range"))
194        }
195    }
196}
197
198impl<MP: MontyParameters> Packable for MontyField31<MP> {}
199
200// Provide a blanket implementation for Monty31 fields here, which forwards the
201// implementation of the variables to the generic argument `<Field>Parameter`,
202// for which we implement the trait (KoalaBear, BabyBear).
203impl<MP> UniformSamplingField for MontyField31<MP>
204where
205    MP: UniformSamplingField + MontyParameters,
206{
207    const MAX_SINGLE_SAMPLE_BITS: usize = MP::MAX_SINGLE_SAMPLE_BITS;
208    const SAMPLING_BITS_M: [u64; 64] = MP::SAMPLING_BITS_M;
209}
210
211impl<FP: FieldParameters> PrimeCharacteristicRing for MontyField31<FP> {
212    type PrimeSubfield = Self;
213
214    const ZERO: Self = FP::MONTY_ZERO;
215    const ONE: Self = FP::MONTY_ONE;
216    const TWO: Self = FP::MONTY_TWO;
217    const NEG_ONE: Self = FP::MONTY_NEG_ONE;
218
219    #[inline(always)]
220    fn from_prime_subfield(f: Self) -> Self {
221        f
222    }
223
224    #[inline]
225    fn halve(&self) -> Self {
226        Self::new_monty(halve_u32::<FP>(self.value))
227    }
228
229    #[inline]
230    fn mul_2exp_u64(&self, exp: u64) -> Self {
231        // The array FP::MONTY_POWERS_OF_TWO contains the powers of 2
232        // from 2^0 to 2^63 in monty form. We can use this to quickly
233        // compute 2^exp.
234        match exp {
235            0 => *self,
236            1 => *self + *self,
237            _ => {
238                if exp < 64 {
239                    *self * Self::MONTY_POWERS_OF_TWO[exp as usize]
240                } else {
241                    // For larger values we use the default method.
242                    *self * Self::TWO.exp_u64(exp)
243                }
244            }
245        }
246    }
247
248    #[inline]
249    fn div_2exp_u64(&self, exp: u64) -> Self {
250        if exp <= 32 {
251            // As the monty form of 2^{-exp} is 2^{32 - exp} mod P, for
252            // 0 <= exp <= 32, we can multiply by 2^{-exp} by doing a shift
253            // followed by a monty reduction.
254            let long_prod = (self.value as u64) << (32 - exp);
255            Self::new_monty(monty_reduce::<FP>(long_prod))
256        } else {
257            // For larger values we use a slower method though this is
258            // still much faster than the default method as it avoids the inverse().
259            *self * Self::HALF.exp_u64(exp)
260        }
261    }
262
263    #[inline]
264    fn zero_vec(len: usize) -> Vec<Self> {
265        // SAFETY:
266        // Due to `#[repr(transparent)]`, MontyField31 and u32 have the same size, alignment
267        // and memory layout making `flatten_to_base` safe. This this will create
268        // a vector MontyField31 elements with value set to 0 which is the
269        // MONTY form of 0.
270        unsafe { flatten_to_base(vec![0u32; len]) }
271    }
272
273    #[inline]
274    fn sum_array<const N: usize>(input: &[Self]) -> Self {
275        assert_eq!(N, input.len());
276        // Benchmarking shows that for N <= 7 it's faster to sum the elements directly
277        // but for N > 7 it's faster to use the .sum() methods which passes through u64's
278        // allowing for delayed reductions.
279        match N {
280            0 => Self::ZERO,
281            1 => input[0],
282            2 => input[0] + input[1],
283            3 => input[0] + input[1] + input[2],
284            4 => (input[0] + input[1]) + (input[2] + input[3]),
285            5 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<1>(&input[4..]),
286            6 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<2>(&input[4..]),
287            7 => Self::sum_array::<4>(&input[..4]) + Self::sum_array::<3>(&input[4..]),
288            _ => input.iter().copied().sum(),
289        }
290    }
291
292    #[inline]
293    fn dot_product<const N: usize>(lhs: &[Self; N], rhs: &[Self; N]) -> Self {
294        const {
295            assert!(N as u64 <= (1 << 34));
296
297            // This code relies on assumptions about the relative size of the
298            // prime and the monty parameter. If these are changes this needs to be checked.
299            debug_assert!(FP::MONTY_BITS == 32);
300            debug_assert!((FP::PRIME as u64) < (1 << 31));
301        }
302        match N {
303            0 => Self::ZERO,
304            1 => lhs[0] * rhs[0],
305            2 => {
306                // As all values are < P < 2^31, the products are < P^2 < 2^31P.
307                // Hence, summing two together we stay below MONTY*P which means
308                // monty_reduce will produce a valid result.
309                let u64_prod_sum = (lhs[0].value as u64) * (rhs[0].value as u64)
310                    + (lhs[1].value as u64) * (rhs[1].value as u64);
311                Self::new_monty(monty_reduce::<FP>(u64_prod_sum))
312            }
313            3 => {
314                // As all values are < P < 2^31, the products are < P^2 < 2^31P.
315                // Hence, summing three together will be less than 2 * MONTY * P
316                let u64_prod_sum = (lhs[0].value as u64) * (rhs[0].value as u64)
317                    + (lhs[1].value as u64) * (rhs[1].value as u64)
318                    + (lhs[2].value as u64) * (rhs[2].value as u64);
319                Self::new_monty(large_monty_reduce::<FP>(u64_prod_sum))
320            }
321            4 => {
322                // As all values are < P < 2^31, the products are < P^2 < 2^31P.
323                // Hence, summing four together will be less than 2 * MONTY * P.
324                let u64_prod_sum = (lhs[0].value as u64) * (rhs[0].value as u64)
325                    + (lhs[1].value as u64) * (rhs[1].value as u64)
326                    + (lhs[2].value as u64) * (rhs[2].value as u64)
327                    + (lhs[3].value as u64) * (rhs[3].value as u64);
328                Self::new_monty(large_monty_reduce::<FP>(u64_prod_sum))
329            }
330            5 => {
331                let head_sum = (lhs[0].value as u64) * (rhs[0].value as u64)
332                    + (lhs[1].value as u64) * (rhs[1].value as u64)
333                    + (lhs[2].value as u64) * (rhs[2].value as u64)
334                    + (lhs[3].value as u64) * (rhs[3].value as u64);
335                let tail_sum = (lhs[4].value as u64) * (rhs[4].value as u64);
336                // head_sum < 4*P^2, tail_sum < P^2.
337                let head_sum_corr = head_sum.wrapping_sub((FP::PRIME as u64) << FP::MONTY_BITS);
338                // head_sum.min(head_sum_corr) reduces a value < 4*P^2 modulo MONTY*P,
339                // so it is < MONTY * P. Hence sum < 2 * MONTY * P.
340                let sum = head_sum.min(head_sum_corr) + tail_sum;
341                Self::new_monty(large_monty_reduce::<FP>(sum))
342            }
343            6 => {
344                let head_sum = (lhs[0].value as u64) * (rhs[0].value as u64)
345                    + (lhs[1].value as u64) * (rhs[1].value as u64)
346                    + (lhs[2].value as u64) * (rhs[2].value as u64)
347                    + (lhs[3].value as u64) * (rhs[3].value as u64);
348                let tail_sum = (lhs[4].value as u64) * (rhs[4].value as u64)
349                    + (lhs[5].value as u64) * (rhs[5].value as u64);
350                // head_sum < 4*P^2, tail_sum < 2*P^2.
351                let head_sum_corr = head_sum.wrapping_sub((FP::PRIME as u64) << FP::MONTY_BITS);
352                // head_sum.min(head_sum_corr) reduces a value < 4*P^2 modulo MONTY*P,
353                // so it is < MONTY * P. Hence sum < 2 * MONTY * P.
354                let sum = head_sum.min(head_sum_corr) + tail_sum;
355                Self::new_monty(large_monty_reduce::<FP>(sum))
356            }
357            7 => {
358                let head_sum = (lhs[0].value as u64) * (rhs[0].value as u64)
359                    + (lhs[1].value as u64) * (rhs[1].value as u64)
360                    + (lhs[2].value as u64) * (rhs[2].value as u64)
361                    + (lhs[3].value as u64) * (rhs[3].value as u64);
362                let tail_sum = (lhs[4].value as u64) * (rhs[4].value as u64)
363                    + lhs[5].value as u64 * (rhs[5].value as u64)
364                    + lhs[6].value as u64 * (rhs[6].value as u64);
365                // head_sum, tail_sum are guaranteed to be < 4*P^2.
366                let head_sum_corr = head_sum.wrapping_sub((FP::PRIME as u64) << FP::MONTY_BITS);
367                let tail_sum_corr = tail_sum.wrapping_sub((FP::PRIME as u64) << FP::MONTY_BITS);
368                // head_sum.min(head_sum_corr), tail_sum.min(tail_sum_corr) each reduce a value
369                // < 4*P^2 modulo MONTY*P, so each is < MONTY * P. Hence sum < 2 * MONTY * P.
370                let sum = head_sum.min(head_sum_corr) + tail_sum.min(tail_sum_corr);
371                Self::new_monty(large_monty_reduce::<FP>(sum))
372            }
373            8 => {
374                let head_sum = (lhs[0].value as u64) * (rhs[0].value as u64)
375                    + (lhs[1].value as u64) * (rhs[1].value as u64)
376                    + (lhs[2].value as u64) * (rhs[2].value as u64)
377                    + (lhs[3].value as u64) * (rhs[3].value as u64);
378                let tail_sum = (lhs[4].value as u64) * (rhs[4].value as u64)
379                    + lhs[5].value as u64 * (rhs[5].value as u64)
380                    + lhs[6].value as u64 * (rhs[6].value as u64)
381                    + lhs[7].value as u64 * (rhs[7].value as u64);
382                // head_sum, tail_sum are guaranteed to be < 4*P^2.
383                let head_sum_corr = head_sum.wrapping_sub((FP::PRIME as u64) << FP::MONTY_BITS);
384                let tail_sum_corr = tail_sum.wrapping_sub((FP::PRIME as u64) << FP::MONTY_BITS);
385                // head_sum.min(head_sum_corr), tail_sum.min(tail_sum_corr) each reduce a value
386                // < 4*P^2 modulo MONTY*P, so each is < MONTY * P. Hence sum < 2 * MONTY * P.
387                let sum = head_sum.min(head_sum_corr) + tail_sum.min(tail_sum_corr);
388                Self::new_monty(large_monty_reduce::<FP>(sum))
389            }
390            _ => {
391                // For large enough N, we accumulate into a u128. This helps the compiler as it lets
392                // it do a lot of computation in parallel as it knows that summing u128's is associative.
393                let acc_u128 = lhs
394                    .chunks(4)
395                    .zip(rhs.chunks(4))
396                    .map(|(l, r)| {
397                        // As all values are < P < 2^31, the products are < P^2 < 2^31P.
398                        // Hence, summing four together will not overflow a u64 but will be
399                        // larger than 2^32P.
400                        let u64_prod_sum = l
401                            .iter()
402                            .zip(r)
403                            .map(|(l, r)| (l.value as u64) * (r.value as u64))
404                            .sum::<u64>();
405                        u64_prod_sum as u128
406                    })
407                    .sum();
408                // As N <= 2^34 by the earlier assertion, acc_u128 <= 2^34 * P^2 < 2^34 * 2^62 < 2^96.
409                Self::new_monty(monty_reduce_u128::<FP>(acc_u128))
410            }
411        }
412    }
413}
414
415impl<FP: FieldParameters + RelativelyPrimePower<D>, const D: u64> InjectiveMonomial<D>
416    for MontyField31<FP>
417{
418}
419
420impl<FP: FieldParameters + RelativelyPrimePower<D>, const D: u64> PermutationMonomial<D>
421    for MontyField31<FP>
422{
423    fn injective_exp_root_n(&self) -> Self {
424        FP::exp_root_d(*self)
425    }
426}
427
428impl<FP: FieldParameters> RawDataSerializable for MontyField31<FP> {
429    impl_raw_serializable_primefield32!();
430}
431
432impl<FP: FieldParameters> Field for MontyField31<FP> {
433    #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
434    type Packing = crate::PackedMontyField31Neon<FP>;
435    #[cfg(all(
436        target_arch = "x86_64",
437        target_feature = "avx2",
438        not(target_feature = "avx512f")
439    ))]
440    type Packing = crate::PackedMontyField31AVX2<FP>;
441    #[cfg(all(target_arch = "x86_64", target_feature = "avx512f"))]
442    type Packing = crate::PackedMontyField31AVX512<FP>;
443    #[cfg(not(any(
444        all(target_arch = "aarch64", target_feature = "neon"),
445        all(
446            target_arch = "x86_64",
447            target_feature = "avx2",
448            not(target_feature = "avx512f")
449        ),
450        all(target_arch = "x86_64", target_feature = "avx512f"),
451    )))]
452    type Packing = Self;
453
454    const GENERATOR: Self = FP::MONTY_GEN;
455
456    const BENEFITS_FROM_LOCKSTEP_EVALUATION: bool = FP::BENEFITS_FROM_LOCKSTEP_EVALUATION;
457
458    fn try_inverse(&self) -> Option<Self> {
459        if self.is_zero() {
460            return None;
461        }
462
463        // The number of bits of FP::PRIME. By the very name of MontyField31 this should always be 31.
464        const NUM_PRIME_BITS: u32 = 31;
465
466        // Get the inverse using a gcd algorithm.
467        // We use `val` to denote the input to `gcd_inversion_prime_field_32` and `R = 2^{MONTY_BITS}`
468        // our monty constant.
469        // The function gcd_inversion_31_bit_field maps `val mod P -> 2^{60}val mod P`
470        let gcd_inverse = gcd_inversion_prime_field_32::<NUM_PRIME_BITS>(self.value, FP::PRIME);
471
472        // Currently |gcd_inverse| <= 2^{NUM_PRIME_BITS - 2} <= 2^{60}
473        // As P > 2^{30}, 0 < 2^{30}P + gcd_inverse < 2^61
474        let pos_inverse = (((FP::PRIME as i64) << 30) + gcd_inverse) as u64;
475
476        // We could do a % operation here, but monty reduction is faster.
477        // This does remove a factor of `R` from the result so we will need to
478        // correct for that.
479        let uncorrected_value = Self::new_monty(monty_reduce::<FP>(pos_inverse));
480
481        // Currently, uncorrected_value = R^{-1} * 2^{60} * val^{-1} mod P = 2^{28} * val^{-1} mod P`.
482        // But `val` is really the monty form of some value `x` satisfying `val = xR mod P`. We want
483        // `x^{-1}R mod P = R^2 x^{-1}R^{-1} mod P = 2^{64} val^{-1} mod P`.
484        // Hence we need to multiply by 2^{64 - 28} = 2^{36}.
485
486        // Unrolling the definitions a little, this 36 comes from: 3 * FP::MONTY_BITS - (2 * NUM_PRIME_BITS - 2)
487
488        Some(uncorrected_value.mul_2exp_u64((3 * FP::MONTY_BITS - (2 * NUM_PRIME_BITS - 2)) as u64))
489    }
490
491    #[inline]
492    fn order() -> BigUint {
493        FP::PRIME.into()
494    }
495}
496
497impl<FP: FieldParameters + TwoAdicData> MontyField31<FP> {
498    /// A square root of this field element, if one exists.
499    ///
500    /// This specializes the generic [`Field::try_sqrt`] for two-adic Monty-31
501    /// fields: it seeds Tonelli–Shanks from the precomputed
502    /// [`TwoAdicField::two_adic_generator`] instead of recomputing `GENERATOR^q`.
503    /// As an inherent method it shadows the trait method for concrete field types
504    /// such as `BabyBear` and `KoalaBear`; generic `Field` callers still use the
505    /// trait default.
506    ///
507    /// See [`Field::try_sqrt`] for the returned-root semantics.
508    #[inline]
509    #[must_use]
510    pub fn try_sqrt(&self) -> Option<Self> {
511        tonelli_shanks_two_adic(*self)
512    }
513}
514
515quotient_map_small_int!(MontyField31, u32, FieldParameters, [u8, u16]);
516quotient_map_small_int!(MontyField31, i32, FieldParameters, [i8, i16]);
517
518impl<FP: FieldParameters> QuotientMap<u32> for MontyField31<FP> {
519    /// Convert a given `u32` integer into an element of the `MontyField31` field.
520    #[inline]
521    fn from_int(int: u32) -> Self {
522        Self::new(int)
523    }
524
525    /// Convert a given `u32` integer into an element of the `MontyField31` field.
526    ///
527    /// Returns `None` if the given integer is greater than the Prime.
528    #[inline]
529    fn from_canonical_checked(int: u32) -> Option<Self> {
530        (int < FP::PRIME).then(|| Self::new(int))
531    }
532
533    /// Convert a given `u32` integer into an element of the `MontyField31` field.
534    ///
535    /// # Safety
536    /// This is always safe as the conversion to monty form can accept any `u32`.
537    #[inline(always)]
538    unsafe fn from_canonical_unchecked(int: u32) -> Self {
539        Self::new(int)
540    }
541}
542
543impl<FP: FieldParameters> QuotientMap<i32> for MontyField31<FP> {
544    /// Convert a given `i32` integer into an element of the `MontyField31` field.
545    #[inline]
546    fn from_int(int: i32) -> Self {
547        Self::new_monty(to_monty_signed::<FP>(int))
548    }
549
550    /// Convert a given `i32` integer into an element of the `MontyField31` field.
551    ///
552    /// Returns `None` if the given integer does not lie in the range `[(1 - P)/2, (P - 1)/2]`.
553    #[inline]
554    fn from_canonical_checked(int: i32) -> Option<Self> {
555        let bound = (FP::PRIME >> 1) as i32;
556        if int <= bound {
557            (int >= (-bound)).then(|| Self::new_monty(to_monty_signed::<FP>(int)))
558        } else {
559            None
560        }
561    }
562
563    /// Convert a given `i32` integer into an element of the `MontyField31` field.
564    ///
565    /// # Safety
566    /// This is always safe as the conversion to monty form can accept any `i32`.
567    #[inline(always)]
568    unsafe fn from_canonical_unchecked(int: i32) -> Self {
569        Self::new_monty(to_monty_signed::<FP>(int))
570    }
571}
572
573impl<FP: FieldParameters> QuotientMap<u64> for MontyField31<FP> {
574    /// Convert a given `u64` integer into an element of the `MontyField31` field.
575    fn from_int(int: u64) -> Self {
576        Self::new_monty(to_monty_64::<FP>(int))
577    }
578
579    /// Convert a given `u64` integer into an element of the `MontyField31` field.
580    ///
581    /// Returns `None` if the given integer is greater than the Prime.
582    fn from_canonical_checked(int: u64) -> Option<Self> {
583        (int < FP::PRIME as u64).then(|| Self::new(int as u32))
584    }
585
586    /// Convert a given `u64` integer into an element of the `MontyField31` field.
587    ///
588    /// # Safety
589    /// This is always safe as the conversion to monty form can accept any `u64`.
590    unsafe fn from_canonical_unchecked(int: u64) -> Self {
591        Self::new_monty(to_monty_64::<FP>(int))
592    }
593}
594
595impl<FP: FieldParameters> QuotientMap<i64> for MontyField31<FP> {
596    /// Convert a given `i64` integer into an element of the `MontyField31` field.
597    fn from_int(int: i64) -> Self {
598        Self::new_monty(to_monty_64_signed::<FP>(int))
599    }
600
601    /// Convert a given `i64` integer into an element of the `MontyField31` field.
602    ///
603    /// Returns `None` if the given integer does not lie in the range `[(1 - P)/2, (P - 1)/2]`.
604    fn from_canonical_checked(int: i64) -> Option<Self> {
605        let bound = (FP::PRIME >> 1) as i64;
606        if int <= bound {
607            (int >= (-bound)).then(|| Self::new_monty(to_monty_signed::<FP>(int as i32)))
608        } else {
609            None
610        }
611    }
612
613    /// Convert a given `i64` integer into an element of the `MontyField31` field.
614    ///
615    /// # Safety
616    /// This is always safe as the conversion to monty form can accept any `i64`.
617    unsafe fn from_canonical_unchecked(int: i64) -> Self {
618        Self::new_monty(to_monty_64_signed::<FP>(int))
619    }
620}
621
622impl<FP: FieldParameters> QuotientMap<u128> for MontyField31<FP> {
623    /// Convert a given `u128` integer into an element of the `MontyField31` field.
624    fn from_int(int: u128) -> Self {
625        Self::new_monty(to_monty::<FP>((int % (FP::PRIME as u128)) as u32))
626    }
627
628    /// Convert a given `u128` integer into an element of the `MontyField31` field.
629    ///
630    /// Returns `None` if the given integer is greater than the Prime.
631    fn from_canonical_checked(int: u128) -> Option<Self> {
632        (int < FP::PRIME as u128).then(|| Self::new(int as u32))
633    }
634
635    /// Convert a given `u128` integer into an element of the `MontyField31` field.
636    ///
637    /// # Safety
638    /// The input must be a valid `u64` element.
639    unsafe fn from_canonical_unchecked(int: u128) -> Self {
640        Self::new_monty(to_monty_64::<FP>(int as u64))
641    }
642}
643
644impl<FP: FieldParameters> QuotientMap<i128> for MontyField31<FP> {
645    /// Convert a given `i128` integer into an element of the `MontyField31` field.
646    fn from_int(int: i128) -> Self {
647        Self::new_monty(to_monty_signed::<FP>((int % (FP::PRIME as i128)) as i32))
648    }
649
650    /// Convert a given `i128` integer into an element of the `MontyField31` field.
651    ///
652    /// Returns `None` if the given integer does not lie in the range `[(1 - P)/2, (P - 1)/2]`.
653    fn from_canonical_checked(int: i128) -> Option<Self> {
654        let bound = (FP::PRIME >> 1) as i128;
655        if int <= bound {
656            (int >= (-bound)).then(|| Self::new_monty(to_monty_signed::<FP>(int as i32)))
657        } else {
658            None
659        }
660    }
661
662    /// Convert a given `i128` integer into an element of the `MontyField31` field.
663    ///
664    /// # Safety
665    /// The input must be a valid `i64` element.
666    unsafe fn from_canonical_unchecked(int: i128) -> Self {
667        Self::new_monty(to_monty_64_signed::<FP>(int as i64))
668    }
669}
670
671impl<FP: FieldParameters> PrimeField for MontyField31<FP> {
672    fn as_canonical_biguint(&self) -> BigUint {
673        self.as_canonical_u32().into()
674    }
675}
676
677impl<FP: FieldParameters> PrimeField64 for MontyField31<FP> {
678    const ORDER_U64: u64 = FP::PRIME as u64;
679
680    #[inline]
681    fn as_canonical_u64(&self) -> u64 {
682        self.as_canonical_u32().into()
683    }
684
685    #[inline]
686    fn to_unique_u64(&self) -> u64 {
687        // The internal representation is already a unique u32 for each field element.
688        // It's fine to hash things in monty form.
689        self.value as u64
690    }
691}
692
693impl<FP: FieldParameters> PrimeField32 for MontyField31<FP> {
694    const ORDER_U32: u32 = FP::PRIME;
695
696    #[inline]
697    fn as_canonical_u32(&self) -> u32 {
698        Self::to_u32(self)
699    }
700
701    #[inline]
702    fn to_unique_u32(&self) -> u32 {
703        // The internal representation is already a unique u32 for each field element.
704        // It's fine to hash things in monty form.
705        self.value
706    }
707}
708
709impl<FP: FieldParameters + TwoAdicData> TwoAdicField for MontyField31<FP> {
710    const TWO_ADICITY: usize = FP::TWO_ADICITY;
711    fn two_adic_generator(bits: usize) -> Self {
712        const {
713            // Verify that 2^TWO_ADICITY divides PRIME - 1.
714            assert!(
715                (FP::PRIME as u64 - 1).is_multiple_of(1u64 << FP::TWO_ADICITY),
716                "2^TWO_ADICITY must divide PRIME - 1"
717            );
718            // Verify maximality: 2^(TWO_ADICITY+1) must NOT divide PRIME - 1.
719            assert!(
720                ((FP::PRIME as u64 - 1) >> FP::TWO_ADICITY) % 2 == 1,
721                "TWO_ADICITY must be maximal"
722            );
723        }
724        assert!(bits <= Self::TWO_ADICITY);
725        FP::TWO_ADIC_GENERATORS.as_ref()[bits]
726    }
727}
728
729impl<FP: MontyParameters> Add for MontyField31<FP> {
730    type Output = Self;
731
732    #[inline]
733    fn add(self, rhs: Self) -> Self {
734        Self::new_monty(add::<FP>(self.value, rhs.value))
735    }
736}
737
738impl<FP: MontyParameters> Sub for MontyField31<FP> {
739    type Output = Self;
740
741    #[inline]
742    fn sub(self, rhs: Self) -> Self {
743        Self::new_monty(sub::<FP>(self.value, rhs.value))
744    }
745}
746
747impl<FP: FieldParameters> Neg for MontyField31<FP> {
748    type Output = Self;
749
750    #[inline]
751    fn neg(self) -> Self::Output {
752        Self::ZERO - self
753    }
754}
755
756impl<FP: MontyParameters> Mul for MontyField31<FP> {
757    type Output = Self;
758
759    #[inline]
760    fn mul(self, rhs: Self) -> Self {
761        let long_prod = self.value as u64 * rhs.value as u64;
762        Self::new_monty(monty_reduce::<FP>(long_prod))
763    }
764}
765
766impl_add_assign!(MontyField31, (MontyParameters, MP));
767impl_sub_assign!(MontyField31, (MontyParameters, MP));
768impl_mul_methods!(MontyField31, (FieldParameters, FP));
769impl_div_methods!(MontyField31, MontyField31, (FieldParameters, FP));
770
771impl<FP: MontyParameters> Sum for MontyField31<FP> {
772    #[inline]
773    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
774        // This is faster than iter.reduce(|x, y| x + y).unwrap_or(Self::ZERO) for iterators of length > 2.
775        // There might be a faster reduction method possible for lengths <= 16 which avoids %.
776
777        // This sum will not overflow so long as iter.len() < 2^33.
778        let sum = iter.map(|x| x.value as u64).sum::<u64>();
779        Self::new_monty((sum % FP::PRIME as u64) as u32)
780    }
781}