From c406f464b8133f5e214d07f827aaed647d50d79b Mon Sep 17 00:00:00 2001
From: Hamish Ivey-Law <426294+unzvfu@users.noreply.github.com>
Date: Fri, 22 Oct 2021 16:05:00 +1100
Subject: [PATCH] Faster squaring in extension fields (#318)

* Faster squaring for extension fields; use in cube().

* Add squaring to arithmetic benchmark.

* Rename variable to remove warning; adjust comment.

* Inline always (improves latency measurements a lot).
---
 benches/field_arithmetic.rs            | 26 ++++++++++++++++++++++++++
 src/field/extension_field/quadratic.rs | 13 +++++++++++++
 src/field/extension_field/quartic.rs   | 13 +++++++++++++
 src/field/field_types.rs               |  6 +++++-
 4 files changed, 57 insertions(+), 1 deletion(-)

diff --git a/benches/field_arithmetic.rs b/benches/field_arithmetic.rs
index 7990973c..e25a8382 100644
--- a/benches/field_arithmetic.rs
+++ b/benches/field_arithmetic.rs
@@ -34,6 +34,32 @@ pub(crate) fn bench_field<F: Field>(c: &mut Criterion) {
         )
     });
 
+    c.bench_function(&format!("sqr-throughput<{}>", type_name::<F>()), |b| {
+        b.iter_batched(
+            || (F::rand(), F::rand(), F::rand(), F::rand()),
+            |(mut x, mut y, mut z, mut w)| {
+                for _ in 0..25 {
+                    (x, y, z, w) = (x.square(), y.square(), z.square(), w.square());
+                }
+                (x, y, z, w)
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    c.bench_function(&format!("sqr-latency<{}>", type_name::<F>()), |b| {
+        b.iter_batched(
+            || F::rand(),
+            |mut x| {
+                for _ in 0..100 {
+                    x = x.square();
+                }
+                x
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
     c.bench_function(&format!("add-throughput<{}>", type_name::<F>()), |b| {
         b.iter_batched(
             || {
diff --git a/src/field/extension_field/quadratic.rs b/src/field/extension_field/quadratic.rs
index a60a1a7f..24cb98cb 100644
--- a/src/field/extension_field/quadratic.rs
+++ b/src/field/extension_field/quadratic.rs
@@ -70,6 +70,19 @@ impl<F: Extendable<2>> Field for QuadraticExtension<F> {
         F::order() * F::order()
     }
 
+    #[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() {
diff --git a/src/field/extension_field/quartic.rs b/src/field/extension_field/quartic.rs
index 0a16ff02..fcd398cd 100644
--- a/src/field/extension_field/quartic.rs
+++ b/src/field/extension_field/quartic.rs
@@ -72,6 +72,19 @@ impl<F: Extendable<4>> Field for QuarticExtension<F> {
         F::order().pow(4u32)
     }
 
+    #[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() {
diff --git a/src/field/field_types.rs b/src/field/field_types.rs
index f23a95f3..8fda4fc6 100644
--- a/src/field/field_types.rs
+++ b/src/field/field_types.rs
@@ -74,12 +74,16 @@ pub trait Field:
         *self == Self::ONE
     }
 
+    fn double(&self) -> Self {
+        *self + *self
+    }
+
     fn square(&self) -> Self {
         *self * *self
     }
 
     fn cube(&self) -> Self {
-        *self * *self * *self
+        self.square() * *self
     }
 
     /// Compute the multiplicative inverse of this field element.