1use 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)] #[must_use]
37pub struct MontyField31<MP: MontyParameters> {
38 pub(crate) value: u32,
43 _phantom: PhantomData<MP>,
44}
45
46impl<MP: MontyParameters> MontyField31<MP> {
47 #[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 #[inline(always)]
71 pub(crate) const fn new_monty(value: u32) -> Self {
72 Self {
73 value,
74 _phantom: PhantomData,
75 }
76 }
77
78 #[inline(always)]
80 pub(crate) const fn to_u32(elem: &Self) -> u32 {
81 from_monty::<MP>(elem.value)
82 }
83
84 #[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 #[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 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 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
200impl<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 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 *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 let long_prod = (self.value as u64) << (32 - exp);
255 Self::new_monty(monty_reduce::<FP>(long_prod))
256 } else {
257 *self * Self::HALF.exp_u64(exp)
260 }
261 }
262
263 #[inline]
264 fn zero_vec(len: usize) -> Vec<Self> {
265 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 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 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 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 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 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 let head_sum_corr = head_sum.wrapping_sub((FP::PRIME as u64) << FP::MONTY_BITS);
338 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 let head_sum_corr = head_sum.wrapping_sub((FP::PRIME as u64) << FP::MONTY_BITS);
352 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 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 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 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 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 let acc_u128 = lhs
394 .chunks(4)
395 .zip(rhs.chunks(4))
396 .map(|(l, r)| {
397 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 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 const NUM_PRIME_BITS: u32 = 31;
465
466 let gcd_inverse = gcd_inversion_prime_field_32::<NUM_PRIME_BITS>(self.value, FP::PRIME);
471
472 let pos_inverse = (((FP::PRIME as i64) << 30) + gcd_inverse) as u64;
475
476 let uncorrected_value = Self::new_monty(monty_reduce::<FP>(pos_inverse));
480
481 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 #[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 #[inline]
521 fn from_int(int: u32) -> Self {
522 Self::new(int)
523 }
524
525 #[inline]
529 fn from_canonical_checked(int: u32) -> Option<Self> {
530 (int < FP::PRIME).then(|| Self::new(int))
531 }
532
533 #[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 #[inline]
546 fn from_int(int: i32) -> Self {
547 Self::new_monty(to_monty_signed::<FP>(int))
548 }
549
550 #[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 #[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 fn from_int(int: u64) -> Self {
576 Self::new_monty(to_monty_64::<FP>(int))
577 }
578
579 fn from_canonical_checked(int: u64) -> Option<Self> {
583 (int < FP::PRIME as u64).then(|| Self::new(int as u32))
584 }
585
586 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 fn from_int(int: i64) -> Self {
598 Self::new_monty(to_monty_64_signed::<FP>(int))
599 }
600
601 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 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 fn from_int(int: u128) -> Self {
625 Self::new_monty(to_monty::<FP>((int % (FP::PRIME as u128)) as u32))
626 }
627
628 fn from_canonical_checked(int: u128) -> Option<Self> {
632 (int < FP::PRIME as u128).then(|| Self::new(int as u32))
633 }
634
635 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 fn from_int(int: i128) -> Self {
647 Self::new_monty(to_monty_signed::<FP>((int % (FP::PRIME as i128)) as i32))
648 }
649
650 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 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 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 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 assert!(
715 (FP::PRIME as u64 - 1).is_multiple_of(1u64 << FP::TWO_ADICITY),
716 "2^TWO_ADICITY must divide PRIME - 1"
717 );
718 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 let sum = iter.map(|x| x.value as u64).sum::<u64>();
779 Self::new_monty((sum % FP::PRIME as u64) as u32)
780 }
781}