Skip to main content

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    const W: Self = FP::W;
82
83    const DTH_ROOT: Self = FP::DTH_ROOT;
84
85    const EXT_GENERATOR: [Self; WIDTH] = FP::EXT_GENERATOR;
86}
87
88impl<const WIDTH: usize, FP> HasTwoAdicBinomialExtension<WIDTH> for MontyField31<FP>
89where
90    FP: BinomialExtensionData<WIDTH> + TwoAdicData + FieldParameters,
91{
92    const EXT_TWO_ADICITY: usize = FP::EXT_TWO_ADICITY;
93
94    fn ext_two_adic_generator(bits: usize) -> [Self; WIDTH] {
95        const {
96            assert!(FP::EXT_TWO_ADICITY >= FP::TWO_ADICITY);
97        }
98        assert!(bits <= Self::EXT_TWO_ADICITY);
99        if bits <= FP::TWO_ADICITY {
100            field_to_array(Self::two_adic_generator(bits))
101        } else {
102            FP::TWO_ADIC_EXTENSION_GENERATORS.as_ref()[bits - FP::TWO_ADICITY - 1]
103        }
104    }
105}
106
107impl<FP> ExtensionAlgebra<Self, 5, QuinticTrinomial> for MontyField31<FP>
108where
109    FP: TrinomialQuinticData + FieldParameters,
110{
111    #[inline(always)]
112    fn ext_mul(a: &[Self; 5], b: &[Self; 5], res: &mut [Self; 5]) {
113        quintic_mul_packed_trinomial(a, b, res);
114    }
115
116    #[inline(always)]
117    fn ext_square(a: &[Self; 5], res: &mut [Self; 5]) {
118        quintic_square(a, res);
119    }
120
121    #[inline(always)]
122    fn ext_add(a: &[Self; 5], b: &[Self; 5]) -> [Self; 5] {
123        let mut res = [Self::ZERO; 5];
124        unsafe {
125            let a: &[u32; 5] = &*(a.as_ptr() as *const [u32; 5]);
126            let b: &[u32; 5] = &*(b.as_ptr() as *const [u32; 5]);
127            let res: &mut [u32; 5] = &mut *(res.as_mut_ptr() as *mut [u32; 5]);
128            packed_mod_add(a, b, res, FP::PRIME, add::<FP>);
129        }
130        res
131    }
132
133    #[inline(always)]
134    fn ext_sub(a: &[Self; 5], b: &[Self; 5]) -> [Self; 5] {
135        let mut res = [Self::ZERO; 5];
136        unsafe {
137            let a: &[u32; 5] = &*(a.as_ptr() as *const [u32; 5]);
138            let b: &[u32; 5] = &*(b.as_ptr() as *const [u32; 5]);
139            let res: &mut [u32; 5] = &mut *(res.as_mut_ptr() as *mut [u32; 5]);
140            packed_mod_sub(a, b, res, FP::PRIME, sub::<FP>);
141        }
142        res
143    }
144
145    #[inline(always)]
146    fn ext_base_mul(lhs: [Self; 5], rhs: Self) -> [Self; 5] {
147        let mut res = [Self::ZERO; 5];
148        base_mul_packed(lhs, rhs, &mut res);
149        res
150    }
151}
152
153impl<FP> QuinticTrinomialExtendable for MontyField31<FP>
154where
155    FP: TrinomialQuinticData + FieldParameters,
156{
157    const FROBENIUS_COEFFS: [[Self; 5]; 4] = FP::FROBENIUS_COEFFS;
158
159    const EXT_GENERATOR: [Self; 5] = FP::EXT_GENERATOR;
160}
161
162impl<FP> HasTwoAdicQuinticExtension for MontyField31<FP>
163where
164    FP: TrinomialQuinticData + TwoAdicData + FieldParameters,
165{
166    const EXT_TWO_ADICITY: usize = FP::EXT_TWO_ADICITY;
167
168    fn ext_two_adic_generator(bits: usize) -> [Self; 5] {
169        assert!(bits <= Self::EXT_TWO_ADICITY);
170        if bits <= FP::TWO_ADICITY {
171            field_to_array(Self::two_adic_generator(bits))
172        } else {
173            FP::TWO_ADIC_EXTENSION_GENERATORS.as_ref()[bits - FP::TWO_ADICITY - 1]
174        }
175    }
176}