Skip to main content

p3_monty_31/
extension.rs

1use p3_field::extension::{
2    BinomiallyExtendable, BinomiallyExtendableAlgebra, HasTwoAdicBinomialExtension,
3    HasTwoAdicQuinticExtension, QuinticExtendableAlgebra, 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> BinomiallyExtendableAlgebra<Self, WIDTH> for MontyField31<FP>
22where
23    FP: BinomialExtensionData<WIDTH> + FieldParameters,
24{
25    #[inline(always)]
26    fn binomial_mul(a: &[Self; WIDTH], b: &[Self; WIDTH], res: &mut [Self; WIDTH], _w: Self) {
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 binomial_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 binomial_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 binomial_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        assert!(bits <= Self::EXT_TWO_ADICITY);
90        if bits <= FP::TWO_ADICITY {
91            field_to_array(Self::two_adic_generator(bits))
92        } else {
93            FP::TWO_ADIC_EXTENSION_GENERATORS.as_ref()[bits - FP::TWO_ADICITY - 1]
94        }
95    }
96}
97
98impl<FP> QuinticExtendableAlgebra<Self> for MontyField31<FP>
99where
100    FP: TrinomialQuinticData + FieldParameters,
101{
102    #[inline(always)]
103    fn quintic_mul(a: &[Self; 5], b: &[Self; 5], res: &mut [Self; 5]) {
104        quintic_mul_packed_trinomial(a, b, res);
105    }
106}
107
108impl<FP> QuinticTrinomialExtendable for MontyField31<FP>
109where
110    FP: TrinomialQuinticData + FieldParameters,
111{
112    const FROBENIUS_COEFFS: [[Self; 5]; 4] = FP::FROBENIUS_COEFFS;
113
114    const EXT_GENERATOR: [Self; 5] = FP::EXT_GENERATOR;
115}
116
117impl<FP> HasTwoAdicQuinticExtension for MontyField31<FP>
118where
119    FP: TrinomialQuinticData + TwoAdicData + FieldParameters,
120{
121    const EXT_TWO_ADICITY: usize = FP::EXT_TWO_ADICITY;
122
123    fn ext_two_adic_generator(bits: usize) -> [Self; 5] {
124        assert!(bits <= Self::EXT_TWO_ADICITY);
125        if bits <= FP::TWO_ADICITY {
126            field_to_array(Self::two_adic_generator(bits))
127        } else {
128            FP::TWO_ADIC_EXTENSION_GENERATORS.as_ref()[bits - FP::TWO_ADICITY - 1]
129        }
130    }
131}