Math: Implement IsSubFieldOf for field extensions (#712)

* add subfield trait

* implement IsSubfieldOf for Degree2ExtensionField struct in BLS12381

* fix mul

* avoid using field elements in issubfield impl. Add tests

* clippy and fmt

* simplify trait bounds

* add explicit type

* change iter to into_iter

* fix metal code

* remove explicit type

* refactor quadratic extensions and implement IsSubFieldOf

* refactor cubic extensions and implement issubfield

* add to_subfield_vec method to field element

* fix type

* clippy

* make to_subfield_vec available only under std feature

* fmt

* merge main

math/src/elliptic_curve/edwards/traits.rs

@@ -12,7 +12,7 @@ pub trait IsEdwards: IsEllipticCurve + Clone + Debug {
         y: &FieldElement<Self::BaseField>,
     ) -> FieldElement<Self::BaseField> {
         (Self::a() * x.pow(2_u16) + y.pow(2_u16))
-            - FieldElement::one()
+            - FieldElement::<Self::BaseField>::one()
             - Self::d() * x.pow(2_u16) * y.pow(2_u16)
     }
 }

math/src/elliptic_curve/point.rs

@@ -68,7 +68,8 @@ impl<E: IsEllipticCurve> Eq for ProjectivePoint<E> {}
 mod tests {
     use crate::cyclic_group::IsGroup;
     use crate::elliptic_curve::short_weierstrass::curves::test_curve_1::{
-        TestCurve1, TestCurveQuadraticNonResidue, TEST_CURVE_1_MAIN_SUBGROUP_ORDER,
+        TestCurve1, TestCurvePrimeField, TestCurveQuadraticNonResidue,
+        TEST_CURVE_1_MAIN_SUBGROUP_ORDER,
     };
     use crate::elliptic_curve::short_weierstrass::curves::test_curve_2::TestCurve2;
     use crate::field::element::FieldElement;
@@ -78,7 +79,7 @@ mod tests {
     use crate::field::extensions::quadratic::QuadraticExtensionFieldElement;
 
     #[allow(clippy::upper_case_acronyms)]
-    type FEE = QuadraticExtensionFieldElement<TestCurveQuadraticNonResidue>;
+    type FEE = QuadraticExtensionFieldElement<TestCurvePrimeField, TestCurveQuadraticNonResidue>;
 
     // This tests only apply for the specific curve found in the configuration file.
     #[test]

math/src/elliptic_curve/short_weierstrass/curves/bls12_381/field_extension.rs

@@ -142,6 +142,11 @@ impl IsSubFieldOf<Degree2ExtensionField> for BLS12381PrimeField {
     fn embed(a: Self::BaseType) -> <Degree2ExtensionField as IsField>::BaseType {
         [FieldElement::from_raw(a), FieldElement::zero()]
     }
+
+    #[cfg(feature = "std")]
+    fn to_subfield_vec(b: <Degree2ExtensionField as IsField>::BaseType) -> Vec<Self::BaseType> {
+        b.into_iter().map(|x| x.to_raw()).collect()
+    }
 }
 
 impl ByteConversion for FieldElement<Degree2ExtensionField> {
@@ -183,9 +188,7 @@ impl ByteConversion for FieldElement<Degree2ExtensionField> {
 ///////////////
 #[derive(Debug, Clone)]
 pub struct LevelTwoResidue;
-impl HasCubicNonResidue for LevelTwoResidue {
-    type BaseField = Degree2ExtensionField;
-
+impl HasCubicNonResidue<Degree2ExtensionField> for LevelTwoResidue {
     fn residue() -> FieldElement<Degree2ExtensionField> {
         FieldElement::new([
             FieldElement::new(U384::from("1")),
@@ -194,13 +197,11 @@ impl HasCubicNonResidue for LevelTwoResidue {
     }
 }
 
-pub type Degree6ExtensionField = CubicExtensionField<LevelTwoResidue>;
+pub type Degree6ExtensionField = CubicExtensionField<Degree2ExtensionField, LevelTwoResidue>;
 
 #[derive(Debug, Clone)]
 pub struct LevelThreeResidue;
-impl HasQuadraticNonResidue for LevelThreeResidue {
-    type BaseField = Degree6ExtensionField;
-
+impl HasQuadraticNonResidue<Degree6ExtensionField> for LevelThreeResidue {
     fn residue() -> FieldElement<Degree6ExtensionField> {
         FieldElement::new([
             FieldElement::zero(),
@@ -210,7 +211,7 @@ impl HasQuadraticNonResidue for LevelThreeResidue {
     }
 }
 
-pub type Degree12ExtensionField = QuadraticExtensionField<LevelThreeResidue>;
+pub type Degree12ExtensionField = QuadraticExtensionField<Degree6ExtensionField, LevelThreeResidue>;
 
 impl FieldElement<BLS12381PrimeField> {
     pub fn new_base(a_hex: &str) -> Self {

math/src/elliptic_curve/short_weierstrass/curves/bls12_381/pairing.rs

@@ -87,7 +87,7 @@ fn double_accumulate_line(
     let [a1, a3, a5] = y.value();
     let b0 = e - b;
     let b2 = FieldElement::new([x1_sq_30 * px, x1_sq_31 * px]);
-    let b3 = FieldElement::new([-h0 * py, -h1 * py]);
+    let b3 = FieldElement::<Degree2ExtensionField>::new([-h0 * py, -h1 * py]);
     *accumulator = FieldElement::new([
         FieldElement::new([
             a0 * &b0 + &residue * (a3 * &b3 + a4 * &b2), // w0
@@ -138,7 +138,7 @@ fn add_accumulate_line(
     let [a1, a3, a5] = y.value();
     let b0 = -lambda.clone() * y2 + theta.clone() * x2;
     let b2 = FieldElement::new([-theta0 * px, -theta1 * px]);
-    let b3 = FieldElement::new([lambda0 * py, lambda1 * py]);
+    let b3 = FieldElement::<Degree2ExtensionField>::new([lambda0 * py, lambda1 * py]);
     *accumulator = FieldElement::new([
         FieldElement::new([
             a0 * &b0 + &residue * (a3 * &b3 + a4 * &b2), // w0

math/src/elliptic_curve/short_weierstrass/curves/test_curve_1.rs

@@ -19,15 +19,15 @@ pub const TEST_CURVE_1_PRIME_FIELD_ORDER: u64 = 59;
 /// Order of the subgroup of the curve.
 pub const TEST_CURVE_1_MAIN_SUBGROUP_ORDER: u64 = 5;
 
+pub type TestCurvePrimeField = U64PrimeField<TEST_CURVE_1_PRIME_FIELD_ORDER>;
+
 /// In F59 the element -1 is not a square. We use this property
 /// to construct a Quadratic Field Extension out of it by adding
 /// its square root.
 #[derive(Debug, Clone)]
 pub struct TestCurveQuadraticNonResidue;
-impl HasQuadraticNonResidue for TestCurveQuadraticNonResidue {
-    type BaseField = U64PrimeField<TEST_CURVE_1_PRIME_FIELD_ORDER>;
-
-    fn residue() -> FieldElement<U64PrimeField<TEST_CURVE_1_PRIME_FIELD_ORDER>> {
+impl HasQuadraticNonResidue<TestCurvePrimeField> for TestCurveQuadraticNonResidue {
+    fn residue() -> FieldElement<TestCurvePrimeField> {
         -FieldElement::one()
     }
 }
@@ -37,7 +37,7 @@ impl HasQuadraticNonResidue for TestCurveQuadraticNonResidue {
 pub struct TestCurve1;
 
 impl IsEllipticCurve for TestCurve1 {
-    type BaseField = QuadraticExtensionField<TestCurveQuadraticNonResidue>;
+    type BaseField = QuadraticExtensionField<TestCurvePrimeField, TestCurveQuadraticNonResidue>;
     type PointRepresentation = ShortWeierstrassProjectivePoint<Self>;
 
     fn generator() -> Self::PointRepresentation {

math/src/elliptic_curve/short_weierstrass/curves/test_curve_2.rs

@@ -34,9 +34,7 @@ type TestCurve2PrimeField = MontgomeryBackendPrimeField<TestCurve2Modulus, 6>;
 /// its square root.
 #[derive(Debug, Clone)]
 pub struct TestCurve2QuadraticNonResidue;
-impl HasQuadraticNonResidue for TestCurve2QuadraticNonResidue {
-    type BaseField = TestCurve2PrimeField;
-
+impl HasQuadraticNonResidue<TestCurve2PrimeField> for TestCurve2QuadraticNonResidue {
     fn residue() -> FieldElement<TestCurve2PrimeField> {
         -FieldElement::one()
     }
@@ -47,7 +45,7 @@ impl HasQuadraticNonResidue for TestCurve2QuadraticNonResidue {
 pub struct TestCurve2;
 
 impl IsEllipticCurve for TestCurve2 {
-    type BaseField = QuadraticExtensionField<TestCurve2QuadraticNonResidue>;
+    type BaseField = QuadraticExtensionField<TestCurve2PrimeField, TestCurve2QuadraticNonResidue>;
     type PointRepresentation = ShortWeierstrassProjectivePoint<Self>;
 
     fn generator() -> Self::PointRepresentation {

math/src/fft/test_helpers.rs

@@ -12,7 +12,8 @@ pub fn naive_matrix_dft_test<F: IsFFTField>(input: &[FieldElement<F>]) -> Vec<Fi
     assert!(n.is_power_of_two());
     let order = n.trailing_zeros();
 
-    let twiddles = get_powers_of_primitive_root(order.into(), n, RootsConfig::Natural).unwrap();
+    let twiddles =
+        get_powers_of_primitive_root::<F>(order.into(), n, RootsConfig::Natural).unwrap();
 
     let mut output = Vec::with_capacity(n);
     for row in 0..n {

math/src/field/element.rs

@@ -63,6 +63,17 @@ impl<F: IsField> FieldElement<F> {
         numbers[0] = bi_inv;
         Ok(())
     }
+
+    #[inline(always)]
+    pub fn to_subfield_vec<S>(self) -> Vec<FieldElement<S>>
+    where
+        S: IsSubFieldOf<F>,
+    {
+        S::to_subfield_vec(self.value)
+            .into_iter()
+            .map(|x| FieldElement::from_raw(x))
+            .collect()
+    }
 }
 
 /// From overloading for field elements
@@ -434,6 +445,11 @@ where
         Self { value: F::zero() }
     }
 
+    /// Returns the raw base type
+    pub fn to_raw(self) -> F::BaseType {
+        self.value
+    }
+
     #[inline(always)]
     pub fn to_extension<L: IsField>(self) -> FieldElement<L>
     where