p3_monty_31/
extension.rs

1use p3_field::extension::{
2    Binomial, BinomiallyExtendable, ExtensionAlgebra, HasTwoAdicBinomialExtension,
3    HasTwoAdicQuinticExtension, QuinticTrinomial, QuinticTrinomialExtendable, binomial_square,
4    quintic_square,
5};
6use p3_field::{
7    PrimeCharacteristicRing, TwoAdicField, field_to_array, packed_mod_add, packed_mod_sub,
8};
9
10use crate::utils::{add, sub};
11use crate::{
12    BinomialExtensionData, FieldParameters, MontyField31, TrinomialQuinticData, TwoAdicData,
13    base_mul_packed, octic_mul_packed, quartic_mul_packed, quintic_mul_packed,
14    quintic_mul_packed_trinomial,
15};
16
17// If a field implements BinomialExtensionData<WIDTH> then there is a natural
18// field extension of degree WIDTH we can define.
19// We perform no checks to make sure the data given in BinomialExtensionData<WIDTH> is valid and
20// corresponds to an actual field extension. Ensuring that is left to the implementer.
21
22impl<const WIDTH: usize, FP> ExtensionAlgebra<Self, WIDTH, Binomial<Self>> for MontyField31<FP>
23where
24    FP: BinomialExtensionData<WIDTH> + FieldParameters,
25{
26    #[inline(always)]
27    fn ext_mul(a: &[Self; WIDTH], b: &[Self; WIDTH], res: &mut [Self; WIDTH]) {
28        match WIDTH {
29            4 => quartic_mul_packed(a, b, res),
30            5 => quintic_mul_packed(a, b, res),
31            8 => octic_mul_packed(a, b, res),
32            _ => panic!("Unsupported binomial extension degree: {}", WIDTH),
33        }
34    }
35
36    #[inline(always)]
37    fn ext_add(a: &[Self; WIDTH], b: &[Self; WIDTH]) -> [Self; WIDTH] {
38        let mut res = [Self::ZERO; WIDTH];
39        unsafe {
40            // Safe as Self is repr(transparent) and stores a single u32.
41            let a: &[u32; WIDTH] = &*(a.as_ptr() as *const [u32; WIDTH]);
42            let b: &[u32; WIDTH] = &*(b.as_ptr() as *const [u32; WIDTH]);
43            let res: &mut [u32; WIDTH] = &mut *(res.as_mut_ptr() as *mut [u32; WIDTH]);
44
45            packed_mod_add(a, b, res, FP::PRIME, add::<FP>);
46        }
47        res
48    }
49
50    #[inline(always)]
51    fn ext_sub(a: &[Self; WIDTH], b: &[Self; WIDTH]) -> [Self; WIDTH] {
52        let mut res = [Self::ZERO; WIDTH];
53        unsafe {
54            // Safe as Self is repr(transparent) and stores a single u32.
55            let a: &[u32; WIDTH] = &*(a.as_ptr() as *const [u32; WIDTH]);
56            let b: &[u32; WIDTH] = &*(b.as_ptr() as *const [u32; WIDTH]);
57            let res: &mut [u32; WIDTH] = &mut *(res.as_mut_ptr() as *mut [u32; WIDTH]);
58
59            packed_mod_sub(a, b, res, FP::PRIME, sub::<FP>);
60        }
61        res
62    }
63
64    #[inline(always)]
65    fn ext_base_mul(lhs: [Self; WIDTH], rhs: Self) -> [Self; WIDTH] {
66        let mut res = [Self::ZERO; WIDTH];
67        base_mul_packed(lhs, rhs, &mut res);
68        res
69    }
70
71    #[inline(always)]
72    fn ext_square(a: &[Self; WIDTH], res: &mut [Self; WIDTH]) {
73        binomial_square::<Self, Self, WIDTH>(a, res, <Self as BinomiallyExtendable<WIDTH>>::W);
74    }
75}
76
77impl<const WIDTH: usize, FP> BinomiallyExtendable<WIDTH> for MontyField31<FP>
78where
79    FP: BinomialExtensionData<WIDTH> + FieldParameters,
80{
81    fn binomial_algebra_id() -> alloc::vec::Vec<u8> {
82        use p3_field::PrimeField32;
83        alloc::format!("p3-power-basis-v1:X^{WIDTH}-{}", FP::W.as_canonical_u32()).into_bytes()
84    }
85
86    const W: Self = FP::W;
87
88    const DTH_ROOT: Self = FP::DTH_ROOT;
89
90    const EXT_GENERATOR: [Self; WIDTH] = FP::EXT_GENERATOR;
91}
92
93impl<const WIDTH: usize, FP> HasTwoAdicBinomialExtension<WIDTH> for MontyField31<FP>
94where
95    FP: BinomialExtensionData<WIDTH> + TwoAdicData + FieldParameters,
96{
97    const EXT_TWO_ADICITY: usize = FP::EXT_TWO_ADICITY;
98
99    fn ext_two_adic_generator(bits: usize) -> [Self; WIDTH] {
100        const {
101            assert!(FP::EXT_TWO_ADICITY >= FP::TWO_ADICITY);
102        }
103        assert!(bits <= Self::EXT_TWO_ADICITY);
104        if bits <= FP::TWO_ADICITY {
105            field_to_array(Self::two_adic_generator(bits))
106        } else {
107            FP::TWO_ADIC_EXTENSION_GENERATORS.as_ref()[bits - FP::TWO_ADICITY - 1]
108        }
109    }
110}
111
112impl<FP> ExtensionAlgebra<Self, 5, QuinticTrinomial> for MontyField31<FP>
113where
114    FP: TrinomialQuinticData + FieldParameters,
115{
116    #[inline(always)]
117    fn ext_mul(a: &[Self; 5], b: &[Self; 5], res: &mut [Self; 5]) {
118        quintic_mul_packed_trinomial(a, b, res);
119    }
120
121    #[inline(always)]
122    fn ext_square(a: &[Self; 5], res: &mut [Self; 5]) {
123        quintic_square(a, res);
124    }
125
126    #[inline(always)]
127    fn ext_add(a: &[Self; 5], b: &[Self; 5]) -> [Self; 5] {
128        let mut res = [Self::ZERO; 5];
129        unsafe {
130            let a: &[u32; 5] = &*(a.as_ptr() as *const [u32; 5]);
131            let b: &[u32; 5] = &*(b.as_ptr() as *const [u32; 5]);
132            let res: &mut [u32; 5] = &mut *(res.as_mut_ptr() as *mut [u32; 5]);
133            packed_mod_add(a, b, res, FP::PRIME, add::<FP>);
134        }
135        res
136    }
137
138    #[inline(always)]
139    fn ext_sub(a: &[Self; 5], b: &[Self; 5]) -> [Self; 5] {
140        let mut res = [Self::ZERO; 5];
141        unsafe {
142            let a: &[u32; 5] = &*(a.as_ptr() as *const [u32; 5]);
143            let b: &[u32; 5] = &*(b.as_ptr() as *const [u32; 5]);
144            let res: &mut [u32; 5] = &mut *(res.as_mut_ptr() as *mut [u32; 5]);
145            packed_mod_sub(a, b, res, FP::PRIME, sub::<FP>);
146        }
147        res
148    }
149
150    #[inline(always)]
151    fn ext_base_mul(lhs: [Self; 5], rhs: Self) -> [Self; 5] {
152        let mut res = [Self::ZERO; 5];
153        base_mul_packed(lhs, rhs, &mut res);
154        res
155    }
156}
157
158impl<FP> QuinticTrinomialExtendable for MontyField31<FP>
159where
160    FP: TrinomialQuinticData + FieldParameters,
161{
162    const FROBENIUS_COEFFS: [[Self; 5]; 4] = FP::FROBENIUS_COEFFS;
163
164    const EXT_GENERATOR: [Self; 5] = FP::EXT_GENERATOR;
165}
166
167impl<FP> HasTwoAdicQuinticExtension for MontyField31<FP>
168where
169    FP: TrinomialQuinticData + TwoAdicData + FieldParameters,
170{
171    const EXT_TWO_ADICITY: usize = FP::EXT_TWO_ADICITY;
172
173    fn ext_two_adic_generator(bits: usize) -> [Self; 5] {
174        assert!(bits <= Self::EXT_TWO_ADICITY);
175        if bits <= FP::TWO_ADICITY {
176            field_to_array(Self::two_adic_generator(bits))
177        } else {
178            FP::TWO_ADIC_EXTENSION_GENERATORS.as_ref()[bits - FP::TWO_ADICITY - 1]
179        }
180    }
181}
p3_monty_31/extension.rs

p3_monty_31/
extension.rs
