lambdaworks_math/field/fields/fft_friendly/
babybear.rs

1use crate::{
2    field::{
3        element::FieldElement,
4        fields::montgomery_backed_prime_fields::{IsModulus, MontgomeryBackendPrimeField},
5        traits::IsFFTField,
6    },
7    unsigned_integer::element::{UnsignedInteger, U64},
8};
9
10pub type U64MontgomeryBackendPrimeField<T> = MontgomeryBackendPrimeField<T, 1>;
11
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct MontgomeryConfigBabybear31PrimeField;
14impl IsModulus<U64> for MontgomeryConfigBabybear31PrimeField {
15    //Babybear Prime p = 2^31 - 2^27 + 1 = 0x78000001
16    const MODULUS: U64 = U64::from_u64(2013265921);
17}
18
19pub type Babybear31PrimeField =
20    U64MontgomeryBackendPrimeField<MontgomeryConfigBabybear31PrimeField>;
21
22//a two-adic primitive root of unity is 21^(2^24)
23// 21^(2^24)=1 mod 2013265921
24// 2^27(2^4-1)+1 where n=27 (two-adicity) and k=2^4+1
25
26//In the future we should allow this with cuda feature, and just dispatch it to the CPU until the implementation is done
27#[cfg(not(feature = "cuda"))]
28impl IsFFTField for Babybear31PrimeField {
29    const TWO_ADICITY: u64 = 24;
30
31    const TWO_ADIC_PRIMITVE_ROOT_OF_UNITY: Self::BaseType = UnsignedInteger { limbs: [21] };
32
33    fn field_name() -> &'static str {
34        "babybear31"
35    }
36}
37
38impl FieldElement<Babybear31PrimeField> {
39    pub fn to_bytes_le(&self) -> [u8; 8] {
40        let limbs = self.representative().limbs;
41        limbs[0].to_le_bytes()
42    }
43
44    pub fn to_bytes_be(&self) -> [u8; 8] {
45        let limbs = self.representative().limbs;
46        limbs[0].to_be_bytes()
47    }
48}
49
50#[cfg(test)]
51mod tests {
52    use super::*;
53
54    mod test_babybear_31_bytes_ops {
55        use super::*;
56        use crate::{field::element::FieldElement, traits::ByteConversion};
57
58        #[test]
59        #[cfg(feature = "alloc")]
60        fn byte_serialization_for_a_number_matches_with_byte_conversion_implementation_le() {
61            let element =
62                FieldElement::<Babybear31PrimeField>::from_hex_unchecked("0123456701234567");
63            let bytes = element.to_bytes_le();
64            let expected_bytes: [u8; 8] = ByteConversion::to_bytes_le(&element).try_into().unwrap();
65            assert_eq!(bytes, expected_bytes);
66        }
67
68        #[test]
69        #[cfg(feature = "alloc")]
70        fn byte_serialization_for_a_number_matches_with_byte_conversion_implementation_be() {
71            let element =
72                FieldElement::<Babybear31PrimeField>::from_hex_unchecked("0123456701234567");
73            let bytes = element.to_bytes_be();
74            let expected_bytes: [u8; 8] = ByteConversion::to_bytes_be(&element).try_into().unwrap();
75            assert_eq!(bytes, expected_bytes);
76        }
77
78        #[test]
79        fn byte_serialization_and_deserialization_works_le() {
80            let element =
81                FieldElement::<Babybear31PrimeField>::from_hex_unchecked("7654321076543210");
82            let bytes = element.to_bytes_le();
83            let from_bytes = FieldElement::<Babybear31PrimeField>::from_bytes_le(&bytes).unwrap();
84            assert_eq!(element, from_bytes);
85        }
86
87        #[test]
88        fn byte_serialization_and_deserialization_works_be() {
89            let element =
90                FieldElement::<Babybear31PrimeField>::from_hex_unchecked("7654321076543210");
91            let bytes = element.to_bytes_be();
92            let from_bytes = FieldElement::<Babybear31PrimeField>::from_bytes_be(&bytes).unwrap();
93            assert_eq!(element, from_bytes);
94        }
95    }
96
97    #[cfg(all(feature = "std", not(feature = "instruments")))]
98    mod test_babybear_31_fft {
99        use super::*;
100        #[cfg(not(feature = "cuda"))]
101        use crate::fft::cpu::roots_of_unity::{
102            get_powers_of_primitive_root, get_powers_of_primitive_root_coset,
103        };
104        #[cfg(not(feature = "cuda"))]
105        use crate::field::element::FieldElement;
106        #[cfg(not(feature = "cuda"))]
107        use crate::field::traits::{IsFFTField, RootsConfig};
108        use crate::polynomial::Polynomial;
109        use proptest::{collection, prelude::*, std_facade::Vec};
110
111        #[cfg(not(feature = "cuda"))]
112        fn gen_fft_and_naive_evaluation<F: IsFFTField>(
113            poly: Polynomial<FieldElement<F>>,
114        ) -> (Vec<FieldElement<F>>, Vec<FieldElement<F>>) {
115            let len = poly.coeff_len().next_power_of_two();
116            let order = len.trailing_zeros();
117            let twiddles =
118                get_powers_of_primitive_root(order.into(), len, RootsConfig::Natural).unwrap();
119
120            let fft_eval = Polynomial::evaluate_fft::<F>(&poly, 1, None).unwrap();
121            let naive_eval = poly.evaluate_slice(&twiddles);
122
123            (fft_eval, naive_eval)
124        }
125
126        #[cfg(not(feature = "cuda"))]
127        fn gen_fft_coset_and_naive_evaluation<F: IsFFTField>(
128            poly: Polynomial<FieldElement<F>>,
129            offset: FieldElement<F>,
130            blowup_factor: usize,
131        ) -> (Vec<FieldElement<F>>, Vec<FieldElement<F>>) {
132            let len = poly.coeff_len().next_power_of_two();
133            let order = (len * blowup_factor).trailing_zeros();
134            let twiddles =
135                get_powers_of_primitive_root_coset(order.into(), len * blowup_factor, &offset)
136                    .unwrap();
137
138            let fft_eval =
139                Polynomial::evaluate_offset_fft::<F>(&poly, blowup_factor, None, &offset).unwrap();
140            let naive_eval = poly.evaluate_slice(&twiddles);
141
142            (fft_eval, naive_eval)
143        }
144
145        #[cfg(not(feature = "cuda"))]
146        fn gen_fft_and_naive_interpolate<F: IsFFTField>(
147            fft_evals: &[FieldElement<F>],
148        ) -> (Polynomial<FieldElement<F>>, Polynomial<FieldElement<F>>) {
149            let order = fft_evals.len().trailing_zeros() as u64;
150            let twiddles =
151                get_powers_of_primitive_root(order, 1 << order, RootsConfig::Natural).unwrap();
152
153            let naive_poly = Polynomial::interpolate(&twiddles, fft_evals).unwrap();
154            let fft_poly = Polynomial::interpolate_fft::<F>(fft_evals).unwrap();
155
156            (fft_poly, naive_poly)
157        }
158
159        #[cfg(not(feature = "cuda"))]
160        fn gen_fft_and_naive_coset_interpolate<F: IsFFTField>(
161            fft_evals: &[FieldElement<F>],
162            offset: &FieldElement<F>,
163        ) -> (Polynomial<FieldElement<F>>, Polynomial<FieldElement<F>>) {
164            let order = fft_evals.len().trailing_zeros() as u64;
165            let twiddles = get_powers_of_primitive_root_coset(order, 1 << order, offset).unwrap();
166
167            let naive_poly = Polynomial::interpolate(&twiddles, fft_evals).unwrap();
168            let fft_poly = Polynomial::interpolate_offset_fft(fft_evals, offset).unwrap();
169
170            (fft_poly, naive_poly)
171        }
172
173        #[cfg(not(feature = "cuda"))]
174        fn gen_fft_interpolate_and_evaluate<F: IsFFTField>(
175            poly: Polynomial<FieldElement<F>>,
176        ) -> (Polynomial<FieldElement<F>>, Polynomial<FieldElement<F>>) {
177            let eval = Polynomial::evaluate_fft::<F>(&poly, 1, None).unwrap();
178            let new_poly = Polynomial::interpolate_fft::<F>(&eval).unwrap();
179
180            (poly, new_poly)
181        }
182
183        prop_compose! {
184            fn powers_of_two(max_exp: u8)(exp in 1..max_exp) -> usize { 1 << exp }
185            // max_exp cannot be multiple of the bits that represent a usize, generally 64 or 32.
186            // also it can't exceed the test field's two-adicity.
187        }
188        prop_compose! {
189            fn field_element()(num in any::<u64>().prop_filter("Avoid null coefficients", |x| x != &0)) -> FieldElement<Babybear31PrimeField> {
190                FieldElement::<Babybear31PrimeField>::from(num)
191            }
192        }
193        prop_compose! {
194            fn offset()(num in any::<u64>(), factor in any::<u64>()) -> FieldElement<Babybear31PrimeField> { FieldElement::<Babybear31PrimeField>::from(num).pow(factor) }
195        }
196        prop_compose! {
197            fn field_vec(max_exp: u8)(vec in collection::vec(field_element(), 0..1 << max_exp)) -> Vec<FieldElement<Babybear31PrimeField>> {
198                vec
199            }
200        }
201        prop_compose! {
202            fn non_power_of_two_sized_field_vec(max_exp: u8)(vec in collection::vec(field_element(), 2..1<<max_exp).prop_filter("Avoid polynomials of size power of two", |vec| !vec.len().is_power_of_two())) -> Vec<FieldElement<Babybear31PrimeField>> {
203                vec
204            }
205        }
206        prop_compose! {
207            fn poly(max_exp: u8)(coeffs in field_vec(max_exp)) -> Polynomial<FieldElement<Babybear31PrimeField>> {
208                Polynomial::new(&coeffs)
209            }
210        }
211        prop_compose! {
212            fn poly_with_non_power_of_two_coeffs(max_exp: u8)(coeffs in non_power_of_two_sized_field_vec(max_exp)) -> Polynomial<FieldElement<Babybear31PrimeField>> {
213                Polynomial::new(&coeffs)
214            }
215        }
216
217        proptest! {
218            // Property-based test that ensures FFT eval. gives same result as a naive polynomial evaluation.
219            #[test]
220            #[cfg(not(feature = "cuda"))]
221            fn test_fft_matches_naive_evaluation(poly in poly(8)) {
222                let (fft_eval, naive_eval) = gen_fft_and_naive_evaluation(poly);
223                prop_assert_eq!(fft_eval, naive_eval);
224            }
225
226            // Property-based test that ensures FFT eval. with coset gives same result as a naive polynomial evaluation.
227            #[test]
228            #[cfg(not(feature = "cuda"))]
229            fn test_fft_coset_matches_naive_evaluation(poly in poly(4), offset in offset(), blowup_factor in powers_of_two(4)) {
230                let (fft_eval, naive_eval) = gen_fft_coset_and_naive_evaluation(poly, offset, blowup_factor);
231                prop_assert_eq!(fft_eval, naive_eval);
232            }
233
234            // Property-based test that ensures FFT interpolation is the same as naive..
235            #[test]
236            #[cfg(not(feature = "cuda"))]
237            fn test_fft_interpolate_matches_naive(fft_evals in field_vec(4)
238                                                           .prop_filter("Avoid polynomials of size not power of two",
239                                                                        |evals| evals.len().is_power_of_two())) {
240                let (fft_poly, naive_poly) = gen_fft_and_naive_interpolate(&fft_evals);
241                prop_assert_eq!(fft_poly, naive_poly);
242            }
243
244            // Property-based test that ensures FFT interpolation with an offset is the same as naive.
245            #[test]
246            #[cfg(not(feature = "cuda"))]
247            fn test_fft_interpolate_coset_matches_naive(offset in offset(), fft_evals in field_vec(4)
248                                                           .prop_filter("Avoid polynomials of size not power of two",
249                                                                        |evals| evals.len().is_power_of_two())) {
250                let (fft_poly, naive_poly) = gen_fft_and_naive_coset_interpolate(&fft_evals, &offset);
251                prop_assert_eq!(fft_poly, naive_poly);
252            }
253
254            // Property-based test that ensures interpolation is the inverse operation of evaluation.
255            #[test]
256            #[cfg(not(feature = "cuda"))]
257            fn test_fft_interpolate_is_inverse_of_evaluate(
258                poly in poly(4).prop_filter("Avoid non pows of two", |poly| poly.coeff_len().is_power_of_two())) {
259                let (poly, new_poly) = gen_fft_interpolate_and_evaluate(poly);
260                prop_assert_eq!(poly, new_poly);
261            }
262        }
263    }
264}
lambdaworks_math/field/fields/fft_friendly/babybear.rs

lambdaworks_math/field/fields/fft_friendly/
babybear.rs
