p3_monty_31/
extension.rs

1use p3_field::extension::{
2    BinomiallyExtendable, BinomiallyExtendableAlgebra, HasTwoAdicBinomialExtension,
3};
4use p3_field::{
5    PrimeCharacteristicRing, TwoAdicField, field_to_array, packed_mod_add, packed_mod_sub,
6};
7
8use crate::utils::{add, sub};
9use crate::{
10    BinomialExtensionData, FieldParameters, MontyField31, TwoAdicData, base_mul_packed,
11    octic_mul_packed, quartic_mul_packed, quintic_mul_packed,
12};
13
14// If a field implements BinomialExtensionData<WIDTH> then there is a natural
15// field extension of degree WIDTH we can define.
16// We perform no checks to make sure the data given in BinomialExtensionData<WIDTH> is valid and
17// corresponds to an actual field extension. Ensuring that is left to the implementer.
18
19impl<const WIDTH: usize, FP> BinomiallyExtendableAlgebra<Self, WIDTH> for MontyField31<FP>
20where
21    FP: BinomialExtensionData<WIDTH> + FieldParameters,
22{
23    #[inline(always)]
24    fn binomial_mul(a: &[Self; WIDTH], b: &[Self; WIDTH], res: &mut [Self; WIDTH], _w: Self) {
25        match WIDTH {
26            4 => quartic_mul_packed(a, b, res),
27            5 => quintic_mul_packed(a, b, res),
28            8 => octic_mul_packed(a, b, res),
29            _ => panic!("Unsupported binomial extension degree: {}", WIDTH),
30        }
31    }
32
33    #[inline(always)]
34    fn binomial_add(a: &[Self; WIDTH], b: &[Self; WIDTH]) -> [Self; WIDTH] {
35        let mut res = [Self::ZERO; WIDTH];
36        unsafe {
37            // Safe as Self is repr(transparent) and stores a single u32.
38            let a: &[u32; WIDTH] = &*(a.as_ptr() as *const [u32; WIDTH]);
39            let b: &[u32; WIDTH] = &*(b.as_ptr() as *const [u32; WIDTH]);
40            let res: &mut [u32; WIDTH] = &mut *(res.as_mut_ptr() as *mut [u32; WIDTH]);
41
42            packed_mod_add(a, b, res, FP::PRIME, add::<FP>);
43        }
44        res
45    }
46
47    #[inline(always)]
48    fn binomial_sub(a: &[Self; WIDTH], b: &[Self; WIDTH]) -> [Self; WIDTH] {
49        let mut res = [Self::ZERO; WIDTH];
50        unsafe {
51            // Safe as Self is repr(transparent) and stores a single u32.
52            let a: &[u32; WIDTH] = &*(a.as_ptr() as *const [u32; WIDTH]);
53            let b: &[u32; WIDTH] = &*(b.as_ptr() as *const [u32; WIDTH]);
54            let res: &mut [u32; WIDTH] = &mut *(res.as_mut_ptr() as *mut [u32; WIDTH]);
55
56            packed_mod_sub(a, b, res, FP::PRIME, sub::<FP>);
57        }
58        res
59    }
60
61    #[inline(always)]
62    fn binomial_base_mul(lhs: [Self; WIDTH], rhs: Self) -> [Self; WIDTH] {
63        let mut res = [Self::ZERO; WIDTH];
64        base_mul_packed(lhs, rhs, &mut res);
65        res
66    }
67}
68
69impl<const WIDTH: usize, FP> BinomiallyExtendable<WIDTH> for MontyField31<FP>
70where
71    FP: BinomialExtensionData<WIDTH> + FieldParameters,
72{
73    const W: Self = FP::W;
74
75    const DTH_ROOT: Self = FP::DTH_ROOT;
76
77    const EXT_GENERATOR: [Self; WIDTH] = FP::EXT_GENERATOR;
78}
79
80impl<const WIDTH: usize, FP> HasTwoAdicBinomialExtension<WIDTH> for MontyField31<FP>
81where
82    FP: BinomialExtensionData<WIDTH> + TwoAdicData + FieldParameters,
83{
84    const EXT_TWO_ADICITY: usize = FP::EXT_TWO_ADICITY;
85
86    fn ext_two_adic_generator(bits: usize) -> [Self; WIDTH] {
87        assert!(bits <= Self::EXT_TWO_ADICITY);
88        if bits <= FP::TWO_ADICITY {
89            field_to_array(Self::two_adic_generator(bits))
90        } else {
91            FP::TWO_ADIC_EXTENSION_GENERATORS.as_ref()[bits - FP::TWO_ADICITY - 1]
92        }
93    }
94}