Take advantage of special prime modulus to optimize sqrt

2 files changed, 107 insertions, 4 deletions

diff --git a/crypto/src/math/ec/custom/sec/SecP256K1FieldElement.cs b/crypto/src/math/ec/custom/sec/SecP256K1FieldElement.cs
index 11cf6ff6a..a5481c925 100644
--- a/crypto/src/math/ec/custom/sec/SecP256K1FieldElement.cs
+++ b/crypto/src/math/ec/custom/sec/SecP256K1FieldElement.cs
@@ -126,8 +126,77 @@ namespace Org.BouncyCastle.Math.EC.Custom.Sec
          */
         public override ECFieldElement Sqrt()
         {
-            ECFieldElement root = new FpFieldElement(Q, ToBigInteger()).Sqrt();
-            return root == null ? null : new SecP256K1FieldElement(root.ToBigInteger());
+            /*
+             * Raise this element to the exponent 2^254 - 2^30 - 2^7 - 2^6 - 2^5 - 2^4 - 2^2
+             * 
+             * Breaking up the exponent's binary representation into "repunits", we get:
+             * { 223 1s } { 1 0s } { 22 1s } { 4 0s } { 2 1s } { 2 0s}
+             * 
+             * Therefore we need an addition chain containing 2, 22, 223 (the lengths of the repunits)
+             * We use: 1, [2], 3, 6, 9, 11, [22], 44, 88, 176, 220, [223]
+             */
+
+            uint[] x1 = this.x;
+            if (Nat256.IsZero(x1) || Nat256.IsOne(x1))
+            {
+                return this;
+            }
+
+            uint[] x2 = Nat256.Create();
+            SecP256K1Field.Square(x1, x2);
+            SecP256K1Field.Multiply(x2, x1, x2);
+
+            uint[] x3 = Nat256.Create();
+            SecP256K1Field.Square(x2, x3);
+            SecP256K1Field.Multiply(x3, x1, x3);
+
+            uint[] x6 = Nat256.Create();
+            SecP256K1Field.SquareN(x3, 3, x6);
+            SecP256K1Field.Multiply(x6, x3, x6);
+
+            uint[] x9 = x6;
+            SecP256K1Field.SquareN(x6, 3, x9);
+            SecP256K1Field.Multiply(x9, x3, x9);
+
+            uint[] x11 = x9;
+            SecP256K1Field.SquareN(x9, 2, x11);
+            SecP256K1Field.Multiply(x11, x2, x11);
+
+            uint[] x22 = Nat256.Create();
+            SecP256K1Field.SquareN(x11, 11, x22);
+            SecP256K1Field.Multiply(x22, x11, x22);
+
+            uint[] x44 = x11;
+            SecP256K1Field.SquareN(x22, 22, x44);
+            SecP256K1Field.Multiply(x44, x22, x44);
+
+            uint[] x88 = Nat256.Create();
+            SecP256K1Field.SquareN(x44, 44, x88);
+            SecP256K1Field.Multiply(x88, x44, x88);
+
+            uint[] x176 = Nat256.Create();
+            SecP256K1Field.SquareN(x88, 88, x176);
+            SecP256K1Field.Multiply(x176, x88, x176);
+
+            uint[] x220 = x88;
+            SecP256K1Field.SquareN(x176, 44, x220);
+            SecP256K1Field.Multiply(x220, x44, x220);
+
+            uint[] x223 = x44;
+            SecP256K1Field.SquareN(x220, 3, x223);
+            SecP256K1Field.Multiply(x223, x3, x223);
+
+            uint[] t1 = x223;
+            SecP256K1Field.SquareN(t1, 23, t1);
+            SecP256K1Field.Multiply(t1, x22, t1);
+            SecP256K1Field.SquareN(t1, 6, t1);
+            SecP256K1Field.Multiply(t1, x2, t1);
+            SecP256K1Field.SquareN(t1, 2, t1);
+
+            uint[] t2 = x2;
+            SecP256K1Field.Square(t1, t2);
+
+            return Arrays.AreEqual(x1, t2) ? new SecP256K1FieldElement(t1) : null;
         }
 
         public override bool Equals(object obj)
diff --git a/crypto/src/math/ec/custom/sec/SecP256R1FieldElement.cs b/crypto/src/math/ec/custom/sec/SecP256R1FieldElement.cs
index a4a7004c0..650f12aaf 100644
--- a/crypto/src/math/ec/custom/sec/SecP256R1FieldElement.cs
+++ b/crypto/src/math/ec/custom/sec/SecP256R1FieldElement.cs
@@ -126,8 +126,42 @@ namespace Org.BouncyCastle.Math.EC.Custom.Sec
          */
         public override ECFieldElement Sqrt()
         {
-            ECFieldElement root = new FpFieldElement(Q, ToBigInteger()).Sqrt();
-            return root == null ? null : new SecP256R1FieldElement(root.ToBigInteger());
+            // Raise this element to the exponent 2^254 - 2^222 + 2^190 + 2^94
+
+            uint[] x1 = this.x;
+            if (Nat256.IsZero(x1) || Nat256.IsOne(x1))
+            {
+                return this;
+            }
+
+            uint[] t1 = Nat256.Create();
+            uint[] t2 = Nat256.Create();
+
+            SecP256R1Field.Square(x1, t1);
+            SecP256R1Field.Multiply(t1, x1, t1);
+
+            SecP256R1Field.SquareN(t1, 2, t2);
+            SecP256R1Field.Multiply(t2, t1, t2);
+
+            SecP256R1Field.SquareN(t2, 4, t1);
+            SecP256R1Field.Multiply(t1, t2, t1);
+
+            SecP256R1Field.SquareN(t1, 8, t2);
+            SecP256R1Field.Multiply(t2, t1, t2);
+
+            SecP256R1Field.SquareN(t2, 16, t1);
+            SecP256R1Field.Multiply(t1, t2, t1);
+
+            SecP256R1Field.SquareN(t1, 32, t1);
+            SecP256R1Field.Multiply(t1, x1, t1);
+
+            SecP256R1Field.SquareN(t1, 96, t1);
+            SecP256R1Field.Multiply(t1, x1, t1);
+
+            SecP256R1Field.SquareN(t1, 94, t1);
+            SecP256R1Field.Multiply(t1, t1, t2);
+
+            return Arrays.AreEqual(x1, t2) ? new SecP256R1FieldElement(t1) : null;
         }
 
         public override bool Equals(object obj)