p3_monty_31/
monty_31.rs

Help
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::intrinsics::transmute;
8use core::iter::{Product, Sum};
9use core::marker::PhantomData;
10use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign};
11
12use num_bigint::BigUint;
13use p3_field::{
14    Field, FieldAlgebra, Packable, PrimeField, PrimeField32, PrimeField64, TwoAdicField,
15};
16use rand::distributions::{Distribution, Standard};
17use rand::Rng;
18use serde::de::Error;
19use serde::{Deserialize, Deserializer, Serialize};
20
21use crate::utils::{from_monty, halve_u32, monty_reduce, to_monty, to_monty_64};
22use crate::{FieldParameters, MontyParameters, TwoAdicData};
23
24#[derive(Clone, Copy, Default, Eq, Hash, PartialEq)]
25#[repr(transparent)] // Packed field implementations rely on this!
26pub struct MontyField31<MP: MontyParameters> {
27    /// The MONTY form of the field element, saved as a positive integer less than `P`.
28    ///
29    /// This is `pub(crate)` for tests and delayed reduction strategies. If you're accessing `value` outside of those, you're
30    /// likely doing something fishy.
31    pub(crate) value: u32,
32    _phantom: PhantomData<MP>,
33}
34
35impl<MP: MontyParameters> MontyField31<MP> {
36    // The standard way to crate a new element.
37    // Note that new converts the input into MONTY form so should be avoided in performance critical implementations.
38    #[inline(always)]
39    pub const fn new(value: u32) -> Self {
40        Self {
41            value: to_monty::<MP>(value),
42            _phantom: PhantomData,
43        }
44    }
45
46    // Create a new field element from something already in MONTY form.
47    // This is `pub(crate)` for tests and delayed reduction strategies. If you're using it outside of those, you're
48    // likely doing something fishy.
49    #[inline(always)]
50    pub(crate) const fn new_monty(value: u32) -> Self {
51        Self {
52            value,
53            _phantom: PhantomData,
54        }
55    }
56
57    /// Produce a u32 in range [0, P) from a field element corresponding to the true value.
58    #[inline(always)]
59    pub(crate) fn to_u32(elem: &Self) -> u32 {
60        from_monty::<MP>(elem.value)
61    }
62
63    /// Convert a constant u32 array into a constant array of field elements.
64    /// Constant version of array.map(MontyField31::new).
65    #[inline]
66    pub const fn new_array<const N: usize>(input: [u32; N]) -> [Self; N] {
67        let mut output = [MontyField31::new_monty(0); N];
68        let mut i = 0;
69        loop {
70            if i == N {
71                break;
72            }
73            output[i] = MontyField31::new(input[i]);
74            i += 1;
75        }
76        output
77    }
78
79    /// Convert a constant 2d u32 array into a constant 2d array of field elements.
80    /// Constant version of array.map(MontyField31::new_array).
81    #[inline]
82    pub const fn new_2d_array<const N: usize, const M: usize>(
83        input: [[u32; N]; M],
84    ) -> [[Self; N]; M] {
85        let mut output = [[MontyField31::new_monty(0); N]; M];
86        let mut i = 0;
87        loop {
88            if i == M {
89                break;
90            }
91            output[i] = MontyField31::new_array(input[i]);
92            i += 1;
93        }
94        output
95    }
96
97    /// Multiply the given MontyField31 element by `2^{-n}`.
98    ///
99    /// This makes use of the fact that, as the monty constant is `2^32`,
100    /// the monty form of `2^{-n}` is `2^{32 - n}`. Monty reduction works
101    /// provided the input is `< 2^32P` so this works for `0 <= n <= 32`.
102    #[inline]
103    #[must_use]
104    pub const fn mul_2exp_neg_n(&self, n: u32) -> Self {
105        assert!(n < 33);
106        let value_mul_2exp_neg_n = (self.value as u64) << (32 - n);
107        MontyField31::new_monty(monty_reduce::<MP>(value_mul_2exp_neg_n))
108    }
109}
110
111impl<FP: MontyParameters> Ord for MontyField31<FP> {
112    #[inline]
113    fn cmp(&self, other: &Self) -> core::cmp::Ordering {
114        MontyField31::to_u32(self).cmp(&MontyField31::to_u32(other))
115    }
116}
117
118impl<FP: MontyParameters> PartialOrd for MontyField31<FP> {
119    #[inline]
120    fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
121        Some(self.cmp(other))
122    }
123}
124
125impl<FP: MontyParameters> Display for MontyField31<FP> {
126    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
127        Display::fmt(&MontyField31::to_u32(self), f)
128    }
129}
130
131impl<FP: MontyParameters> Debug for MontyField31<FP> {
132    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
133        Debug::fmt(&MontyField31::to_u32(self), f)
134    }
135}
136
137impl<FP: MontyParameters> Distribution<MontyField31<FP>> for Standard {
138    #[inline]
139    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> MontyField31<FP> {
140        loop {
141            let next_u31 = rng.next_u32() >> 1;
142            let is_canonical = next_u31 < FP::PRIME;
143            if is_canonical {
144                return MontyField31::new_monty(next_u31);
145            }
146        }
147    }
148}
149
150impl<FP: FieldParameters> Serialize for MontyField31<FP> {
151    fn serialize<S: serde::Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
152        // It's faster to Serialize and Deserialize in monty form.
153        serializer.serialize_u32(self.value)
154    }
155}
156
157impl<'de, FP: FieldParameters> Deserialize<'de> for MontyField31<FP> {
158    fn deserialize<D: Deserializer<'de>>(d: D) -> Result<Self, D::Error> {
159        let val = u32::deserialize(d)?;
160        // Ensure that `val` satisfies our invariant, namely is `< P`.
161        if val < FP::PRIME {
162            // It's faster to Serialize and Deserialize in monty form.
163            Ok(MontyField31::new_monty(val))
164        } else {
165            Err(D::Error::custom("Value is out of range"))
166        }
167    }
168}
169
170impl<FP: FieldParameters> Packable for MontyField31<FP> {}
171
172impl<FP: FieldParameters> FieldAlgebra for MontyField31<FP> {
173    type F = Self;
174
175    const ZERO: Self = FP::MONTY_ZERO;
176    const ONE: Self = FP::MONTY_ONE;
177    const TWO: Self = FP::MONTY_TWO;
178    const NEG_ONE: Self = FP::MONTY_NEG_ONE;
179
180    #[inline(always)]
181    fn from_f(f: Self::F) -> Self {
182        f
183    }
184
185    #[inline(always)]
186    fn from_canonical_u8(n: u8) -> Self {
187        Self::from_canonical_u32(n as u32)
188    }
189
190    #[inline(always)]
191    fn from_canonical_u16(n: u16) -> Self {
192        Self::from_canonical_u32(n as u32)
193    }
194
195    #[inline(always)]
196    fn from_canonical_u32(n: u32) -> Self {
197        debug_assert!(n < FP::PRIME);
198        Self::from_wrapped_u32(n)
199    }
200
201    #[inline(always)]
202    fn from_canonical_u64(n: u64) -> Self {
203        debug_assert!(n < FP::PRIME as u64);
204        Self::from_canonical_u32(n as u32)
205    }
206
207    #[inline(always)]
208    fn from_canonical_usize(n: usize) -> Self {
209        debug_assert!(n < FP::PRIME as usize);
210        Self::from_canonical_u32(n as u32)
211    }
212
213    #[inline(always)]
214    fn from_wrapped_u32(n: u32) -> Self {
215        Self::new(n)
216    }
217
218    #[inline(always)]
219    fn from_wrapped_u64(n: u64) -> Self {
220        Self::new_monty(to_monty_64::<FP>(n))
221    }
222
223    #[inline]
224    fn mul_2exp_u64(&self, exp: u64) -> Self {
225        let product = (self.value as u64) << exp;
226        let value = (product % (FP::PRIME as u64)) as u32;
227        Self::new_monty(value)
228    }
229
230    #[inline]
231    fn zero_vec(len: usize) -> Vec<Self> {
232        // SAFETY: repr(transparent) ensures transmutation safety.
233        unsafe { transmute(vec![0u32; len]) }
234    }
235}
236
237impl<FP: FieldParameters> Field for MontyField31<FP> {
238    #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
239    type Packing = crate::PackedMontyField31Neon<FP>;
240    #[cfg(all(
241        target_arch = "x86_64",
242        target_feature = "avx2",
243        not(all(feature = "nightly-features", target_feature = "avx512f"))
244    ))]
245    type Packing = crate::PackedMontyField31AVX2<FP>;
246    #[cfg(all(
247        feature = "nightly-features",
248        target_arch = "x86_64",
249        target_feature = "avx512f"
250    ))]
251    type Packing = crate::PackedMontyField31AVX512<FP>;
252    #[cfg(not(any(
253        all(target_arch = "aarch64", target_feature = "neon"),
254        all(
255            target_arch = "x86_64",
256            target_feature = "avx2",
257            not(all(feature = "nightly-features", target_feature = "avx512f"))
258        ),
259        all(
260            feature = "nightly-features",
261            target_arch = "x86_64",
262            target_feature = "avx512f"
263        ),
264    )))]
265    type Packing = Self;
266
267    const GENERATOR: Self = FP::MONTY_GEN;
268
269    #[inline]
270    fn exp_u64_generic<FA: FieldAlgebra<F = Self>>(val: FA, power: u64) -> FA {
271        FP::exp_u64_generic(val, power)
272    }
273
274    fn try_inverse(&self) -> Option<Self> {
275        FP::try_inverse(*self)
276    }
277
278    #[inline]
279    fn halve(&self) -> Self {
280        Self::new_monty(halve_u32::<FP>(self.value))
281    }
282
283    #[inline]
284    fn order() -> BigUint {
285        FP::PRIME.into()
286    }
287}
288
289impl<FP: FieldParameters> PrimeField for MontyField31<FP> {
290    fn as_canonical_biguint(&self) -> BigUint {
291        <Self as PrimeField32>::as_canonical_u32(self).into()
292    }
293}
294
295impl<FP: FieldParameters> PrimeField64 for MontyField31<FP> {
296    const ORDER_U64: u64 = FP::PRIME as u64;
297
298    #[inline]
299    fn as_canonical_u64(&self) -> u64 {
300        self.as_canonical_u32().into()
301    }
302
303    #[inline]
304    fn to_unique_u64(&self) -> u64 {
305        // The internal representation is already a unique u32 for each field element.
306        // It's fine to hash things in monty form.
307        self.value as u64
308    }
309}
310
311impl<FP: FieldParameters> PrimeField32 for MontyField31<FP> {
312    const ORDER_U32: u32 = FP::PRIME;
313
314    #[inline]
315    fn as_canonical_u32(&self) -> u32 {
316        MontyField31::to_u32(self)
317    }
318
319    #[inline]
320    fn to_unique_u32(&self) -> u32 {
321        // The internal representation is already a unique u32 for each field element.
322        // It's fine to hash things in monty form.
323        self.value
324    }
325}
326
327impl<FP: FieldParameters + TwoAdicData> TwoAdicField for MontyField31<FP> {
328    const TWO_ADICITY: usize = FP::TWO_ADICITY;
329    fn two_adic_generator(bits: usize) -> Self {
330        assert!(bits <= Self::TWO_ADICITY);
331        FP::TWO_ADIC_GENERATORS.as_ref()[bits]
332    }
333}
334
335impl<FP: MontyParameters> Add for MontyField31<FP> {
336    type Output = Self;
337
338    #[inline]
339    fn add(self, rhs: Self) -> Self {
340        let mut sum = self.value + rhs.value;
341        let (corr_sum, over) = sum.overflowing_sub(FP::PRIME);
342        if !over {
343            sum = corr_sum;
344        }
345        Self::new_monty(sum)
346    }
347}
348
349impl<FP: MontyParameters> AddAssign for MontyField31<FP> {
350    #[inline]
351    fn add_assign(&mut self, rhs: Self) {
352        *self = *self + rhs;
353    }
354}
355
356impl<FP: MontyParameters> Sum for MontyField31<FP> {
357    #[inline]
358    fn sum<I: Iterator<Item = Self>>(iter: I) -> Self {
359        // This is faster than iter.reduce(|x, y| x + y).unwrap_or(Self::ZERO) for iterators of length > 2.
360        // There might be a faster reduction method possible for lengths <= 16 which avoids %.
361
362        // This sum will not overflow so long as iter.len() < 2^33.
363        let sum = iter.map(|x| x.value as u64).sum::<u64>();
364        Self::new_monty((sum % FP::PRIME as u64) as u32)
365    }
366}
367
368impl<FP: MontyParameters> Sub for MontyField31<FP> {
369    type Output = Self;
370
371    #[inline]
372    fn sub(self, rhs: Self) -> Self {
373        let (mut diff, over) = self.value.overflowing_sub(rhs.value);
374        let corr = if over { FP::PRIME } else { 0 };
375        diff = diff.wrapping_add(corr);
376        Self::new_monty(diff)
377    }
378}
379
380impl<FP: MontyParameters> SubAssign for MontyField31<FP> {
381    #[inline]
382    fn sub_assign(&mut self, rhs: Self) {
383        *self = *self - rhs;
384    }
385}
386
387impl<FP: FieldParameters> Neg for MontyField31<FP> {
388    type Output = Self;
389
390    #[inline]
391    fn neg(self) -> Self::Output {
392        Self::ZERO - self
393    }
394}
395
396impl<FP: MontyParameters> Mul for MontyField31<FP> {
397    type Output = Self;
398
399    #[inline]
400    fn mul(self, rhs: Self) -> Self {
401        let long_prod = self.value as u64 * rhs.value as u64;
402        Self::new_monty(monty_reduce::<FP>(long_prod))
403    }
404}
405
406impl<FP: MontyParameters> MulAssign for MontyField31<FP> {
407    #[inline]
408    fn mul_assign(&mut self, rhs: Self) {
409        *self = *self * rhs;
410    }
411}
412
413impl<FP: FieldParameters> Product for MontyField31<FP> {
414    #[inline]
415    fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
416        iter.reduce(|x, y| x * y).unwrap_or(Self::ONE)
417    }
418}
419
420impl<FP: FieldParameters> Div for MontyField31<FP> {
421    type Output = Self;
422
423    #[allow(clippy::suspicious_arithmetic_impl)]
424    #[inline]
425    fn div(self, rhs: Self) -> Self {
426        self * rhs.inverse()
427    }
428}
p3_monty_31/monty_31.rs

p3_monty_31/
monty_31.rs
