p3_monty_31/
extension.rs

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

p3_monty_31/
extension.rs
