From b180b79a6c295c2f3508abdf386a35fc62df4370 Mon Sep 17 00:00:00 2001
From: mdvillagra <mdvillagra@gmail.com>
Date: Fri, 8 Sep 2023 19:08:57 -0400
Subject: [PATCH] added requested changes and isFFTField trait

---
 .../src/field/fields/fft_friendly/babybear.rs | 24 +++++++++++++++----
 math/src/field/fields/fft_friendly/mod.rs     |  2 +-
 math/src/unsigned_integer/element.rs          |  1 +
 3 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/math/src/field/fields/fft_friendly/babybear.rs b/math/src/field/fields/fft_friendly/babybear.rs
index bccade931..48786f05a 100644
--- a/math/src/field/fields/fft_friendly/babybear.rs
+++ b/math/src/field/fields/fft_friendly/babybear.rs
@@ -2,13 +2,12 @@ use crate::{
     field::{
         element::FieldElement,
         fields::montgomery_backed_prime_fields::{IsModulus, MontgomeryBackendPrimeField},
-        //traits::IsFFTField,
+        traits::IsFFTField, 
     },
-    unsigned_integer::element::UnsignedInteger,
+    unsigned_integer::element::{U64, UnsignedInteger},
 };
 
-pub type U64 = UnsignedInteger<1>;
-pub type U64PrimeField<T> = MontgomeryBackendPrimeField<T, 1>;
+pub type U64MontgomeryBackendPrimeField<T> = MontgomeryBackendPrimeField<T, 1>;
 
 #[derive(Debug, Clone, PartialEq, Eq)]
 pub struct MontgomeryConfigBabybear31PrimeField;
@@ -17,7 +16,21 @@ impl IsModulus<U64> for MontgomeryConfigBabybear31PrimeField {
     const MODULUS: U64 = U64::from_u64(2013265921);
 }
 
-pub type Babybear31PrimeField = U64PrimeField<MontgomeryConfigBabybear31PrimeField>;
+pub type Babybear31PrimeField = U64MontgomeryBackendPrimeField<MontgomeryConfigBabybear31PrimeField>;
+
+// 21^(2^24)=1 mod 2013265921
+// 2^27(2^4-1)+1 where n=27 (two-adicity) and k=2^4+1
+impl IsFFTField for Babybear31PrimeField {
+    const TWO_ADICITY: u64 = 27;
+
+    const TWO_ADIC_PRIMITVE_ROOT_OF_UNITY: Self::BaseType = UnsignedInteger{
+        limbs: [21],
+    };
+
+    fn field_name() -> &'static str {
+        "babybear31"
+    }
+}
 
 impl FieldElement<Babybear31PrimeField> {
     pub fn to_bytes_le(&self) -> [u8; 8] {
@@ -99,4 +112,5 @@ mod test_babybear_31_bytes_ops {
         let from_bytes = FieldElement::<Babybear31PrimeField>::from_bytes_be(&bytes).unwrap();
         assert_eq!(element, from_bytes);
     }
+    
 }
\ No newline at end of file
diff --git a/math/src/field/fields/fft_friendly/mod.rs b/math/src/field/fields/fft_friendly/mod.rs
index 9adad6e0a..c8853be92 100644
--- a/math/src/field/fields/fft_friendly/mod.rs
+++ b/math/src/field/fields/fft_friendly/mod.rs
@@ -1,4 +1,4 @@
 /// Implemenation of the Babybear Prime field p = 2^31 - 2^27 + 1
 pub mod babybear;
 /// Implementation of two-adic prime field over 256 bit unsigned integers.
-pub mod stark_252_prime_field;
\ No newline at end of file
+pub mod stark_252_prime_field;
diff --git a/math/src/unsigned_integer/element.rs b/math/src/unsigned_integer/element.rs
index faf16ca7c..fea84a98e 100644
--- a/math/src/unsigned_integer/element.rs
+++ b/math/src/unsigned_integer/element.rs
@@ -22,6 +22,7 @@ use core::fmt::{self, Debug, Display};
 pub type U384 = UnsignedInteger<6>;
 pub type U256 = UnsignedInteger<4>;
 pub type U128 = UnsignedInteger<2>;
+pub type U64 = UnsignedInteger<1>;
 
 /// A big unsigned integer in base 2^{64} represented
 /// as fixed-size array `limbs` of `u64` components.