Square trait (#409)

* `Squarable` trait

* Minor style

* Further minor style (Squarable -> Square to match Rust convention)

field/src/arch/x86_64/avx2_goldilocks_field.rs

@@ -7,6 +7,7 @@ use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssi
 
 use crate::field_types::{Field, PrimeField};
 use crate::goldilocks_field::GoldilocksField;
+use crate::ops::Square;
 use crate::packed_field::PackedField;
 
 // Ideally `Avx2GoldilocksField` would wrap `__m256i`. Unfortunately, `__m256i` has an alignment of
@@ -194,7 +195,9 @@ unsafe impl PackedField for Avx2GoldilocksField {
         };
         (Self::new(res0), Self::new(res1))
     }
+}
 
+impl Square for Avx2GoldilocksField {
     #[inline]
     fn square(&self) -> Self {
         Self::new(unsafe { square(self.get()) })
@@ -509,6 +512,7 @@ mod tests {
     use crate::arch::x86_64::avx2_goldilocks_field::Avx2GoldilocksField;
     use crate::field_types::PrimeField;
     use crate::goldilocks_field::GoldilocksField;
+    use crate::ops::Square;
     use crate::packed_field::PackedField;
 
     fn test_vals_a() -> [GoldilocksField; 4] {

field/src/extension_field/quadratic.rs

@@ -9,6 +9,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::extension_field::{Extendable, FieldExtension, Frobenius, OEF};
 use crate::field_types::Field;
+use crate::ops::Square;
 
 #[derive(Copy, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
 #[serde(bound = "")]
@@ -73,19 +74,6 @@ impl<F: Extendable<2>> Field for QuadraticExtension<F> {
         F::characteristic()
     }
 
-    #[inline(always)]
-    fn square(&self) -> Self {
-        // Specialising mul reduces the computation of c1 from 2 muls
-        // and one add to one mul and a shift
-
-        let Self([a0, a1]) = *self;
-
-        let c0 = a0.square() + <Self as OEF<2>>::W * a1.square();
-        let c1 = a0 * a1.double();
-
-        Self([c0, c1])
-    }
-
     // Algorithm 11.3.4 in Handbook of Elliptic and Hyperelliptic Curve Cryptography.
     fn try_inverse(&self) -> Option<Self> {
         if self.is_zero() {
@@ -204,6 +192,21 @@ impl<F: Extendable<2>> MulAssign for QuadraticExtension<F> {
     }
 }
 
+impl<F: Extendable<2>> Square for QuadraticExtension<F> {
+    #[inline(always)]
+    fn square(&self) -> Self {
+        // Specialising mul reduces the computation of c1 from 2 muls
+        // and one add to one mul and a shift
+
+        let Self([a0, a1]) = *self;
+
+        let c0 = a0.square() + <Self as OEF<2>>::W * a1.square();
+        let c1 = a0 * a1.double();
+
+        Self([c0, c1])
+    }
+}
+
 impl<F: Extendable<2>> Product for QuadraticExtension<F> {
     fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
         iter.fold(Self::ONE, |acc, x| acc * x)

field/src/extension_field/quartic.rs

@@ -10,6 +10,7 @@ use serde::{Deserialize, Serialize};
 
 use crate::extension_field::{Extendable, FieldExtension, Frobenius, OEF};
 use crate::field_types::Field;
+use crate::ops::Square;
 
 #[derive(Copy, Clone, Eq, PartialEq, Hash, Serialize, Deserialize)]
 #[serde(bound = "")]
@@ -75,19 +76,6 @@ impl<F: Extendable<4>> Field for QuarticExtension<F> {
         F::characteristic()
     }
 
-    #[inline(always)]
-    fn square(&self) -> Self {
-        let Self([a0, a1, a2, a3]) = *self;
-        let w = <Self as OEF<4>>::W;
-
-        let c0 = a0.square() + w * (a1 * a3.double() + a2.square());
-        let c1 = (a0 * a1 + w * a2 * a3).double();
-        let c2 = a0 * a2.double() + a1.square() + w * a3.square();
-        let c3 = (a0 * a3 + a1 * a2).double();
-
-        Self([c0, c1, c2, c3])
-    }
-
     // Algorithm 11.3.4 in Handbook of Elliptic and Hyperelliptic Curve Cryptography.
     fn try_inverse(&self) -> Option<Self> {
         if self.is_zero() {
@@ -241,6 +229,21 @@ impl<F: Extendable<4>> MulAssign for QuarticExtension<F> {
     }
 }
 
+impl<F: Extendable<4>> Square for QuarticExtension<F> {
+    #[inline(always)]
+    fn square(&self) -> Self {
+        let Self([a0, a1, a2, a3]) = *self;
+        let w = <Self as OEF<4>>::W;
+
+        let c0 = a0.square() + w * (a1 * a3.double() + a2.square());
+        let c1 = (a0 * a1 + w * a2 * a3).double();
+        let c2 = a0 * a2.double() + a1.square() + w * a3.square();
+        let c3 = (a0 * a3 + a1 * a2).double();
+
+        Self([c0, c1, c2, c3])
+    }
+}
+
 impl<F: Extendable<4>> Product for QuarticExtension<F> {
     fn product<I: Iterator<Item = Self>>(iter: I) -> Self {
         iter.fold(Self::ONE, |acc, x| acc * x)

field/src/field_testing.rs

@@ -1,6 +1,7 @@
 use crate::extension_field::Extendable;
 use crate::extension_field::Frobenius;
 use crate::field_types::Field;
+use crate::ops::Square;
 
 #[macro_export]
 macro_rules! test_field_arithmetic {

field/src/field_types.rs

@@ -11,6 +11,7 @@ use serde::de::DeserializeOwned;
 use serde::Serialize;
 
 use crate::extension_field::Frobenius;
+use crate::ops::Square;
 
 /// A finite field.
 pub trait Field:
@@ -26,6 +27,7 @@ pub trait Field:
     + SubAssign<Self>
     + Mul<Self, Output = Self>
     + MulAssign<Self>
+    + Square
     + Product
     + Div<Self, Output = Self>
     + DivAssign<Self>
@@ -80,11 +82,6 @@ pub trait Field:
         *self + *self
     }
 
-    #[inline]
-    fn square(&self) -> Self {
-        *self * *self
-    }
-
     #[inline]
     fn cube(&self) -> Self {
         self.square() * *self

field/src/lib.rs

@@ -17,6 +17,7 @@ pub mod field_types;
 pub mod goldilocks_field;
 pub mod interpolation;
 mod inversion;
+pub mod ops;
 pub mod packable;
 pub mod packed_field;
 pub mod polynomial;

field/src/ops.rs

@@ -0,0 +1,11 @@
+use std::ops::Mul;
+
+pub trait Square {
+    fn square(&self) -> Self;
+}
+
+impl<F: Mul<F, Output = Self> + Copy> Square for F {
+    default fn square(&self) -> Self {
+        *self * *self
+    }
+}

field/src/packed_field.rs

@@ -4,6 +4,7 @@ use std::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign};
 use std::slice;
 
 use crate::field_types::Field;
+use crate::ops::Square;
 
 /// # Safety
 /// - WIDTH is assumed to be a power of 2.
@@ -24,6 +25,7 @@ pub unsafe trait PackedField:
     + Mul<Self::Scalar, Output = Self>
     + MulAssign<Self>
     + MulAssign<Self::Scalar>
+    + Square
     + Neg<Output = Self>
     + Product
     + Send
@@ -44,10 +46,6 @@ where
     const ZEROS: Self;
     const ONES: Self;
 
-    fn square(&self) -> Self {
-        *self * *self
-    }
-
     fn from_arr(arr: [Self::Scalar; Self::WIDTH]) -> Self;
     fn as_arr(&self) -> [Self::Scalar; Self::WIDTH];
 
@@ -106,10 +104,6 @@ unsafe impl<F: Field> PackedField for F {
     const ZEROS: Self = F::ZERO;
     const ONES: Self = F::ONE;
 
-    fn square(&self) -> Self {
-        <Self as Field>::square(self)
-    }
-
     fn from_arr(arr: [Self::Scalar; Self::WIDTH]) -> Self {
         arr[0]
     }

field/src/prime_field_testing.rs

@@ -71,6 +71,7 @@ macro_rules! test_prime_field_arithmetic {
             use std::ops::{Add, Mul, Neg, Sub};
 
             use crate::field_types::{Field, PrimeField};
+            use crate::ops::Square;
 
             #[test]
             fn arithmetic_addition() {

plonky2/src/curve/curve_adds.rs

@@ -1,6 +1,7 @@
 use std::ops::Add;
 
 use plonky2_field::field_types::Field;
+use plonky2_field::ops::Square;
 
 use crate::curve::curve_types::{AffinePoint, Curve, ProjectivePoint};
 

plonky2/src/curve/curve_summation.rs

@@ -1,6 +1,7 @@
 use std::iter::Sum;
 
 use plonky2_field::field_types::Field;
+use plonky2_field::ops::Square;
 
 use crate::curve::curve_types::{AffinePoint, Curve, ProjectivePoint};
 

plonky2/src/curve/curve_types.rs

@@ -2,6 +2,7 @@ use std::fmt::Debug;
 use std::ops::Neg;
 
 use plonky2_field::field_types::Field;
+use plonky2_field::ops::Square;
 
 // To avoid implementation conflicts from associated types,
 // see https://github.com/rust-lang/rust/issues/20400

plonky2/src/gates/exponentiation.rs

@@ -2,6 +2,7 @@ use std::marker::PhantomData;
 
 use plonky2_field::extension_field::Extendable;
 use plonky2_field::field_types::Field;
+use plonky2_field::ops::Square;
 use plonky2_field::packed_field::PackedField;
 
 use crate::gates::gate::Gate;
@@ -90,7 +91,7 @@ impl<F: RichField + Extendable<D>, const D: usize> Gate<F, D> for Exponentiation
             let prev_intermediate_value = if i == 0 {
                 F::Extension::ONE
             } else {
-                <F::Extension as Field>::square(&intermediate_values[i - 1])
+                intermediate_values[i - 1].square()
             };
 
             // power_bits is in LE order, but we accumulate in BE order.