math/big: use optimized formula in ModSqrt for 3 mod 4 primes

For primes which are 3 mod 4, using Tonelli-Shanks is slower
and more complicated than using the identity

     a**((p+1)/4) mod p == sqrt(a)

For 2^450-2^225-1 and 2^10860-2^5430-1, which are 3 mod 4:

BenchmarkModSqrt225_TonelliTri      1000     1135375 ns/op
BenchmarkModSqrt225_3Mod4          10000      156009 ns/op
BenchmarkModSqrt5430_Tonelli           1  3448851386 ns/op
BenchmarkModSqrt5430_3Mod4             2   914616710 ns/op

~2.6x to 7x faster.

Fixes #11437 (which is a prime choice of issues to fix)

Change-Id: I813fb29454160483ec29825469e0370d517850c2
Reviewed-on: https://go-review.googlesource.com/11522
Reviewed-by: Adam Langley <agl@golang.org>

diff --git a/src/math/big/int.go b/src/math/big/int.go
index 65334e0ef5501a9caeceffb2067644293abca1d8..84485a27502bd5c1cda16edab3daa561ebf1cf14 100644 (file)
--- a/src/math/big/int.go
+++ b/src/math/big/int.go
@@ -640,23 +640,23 @@ func Jacobi(x, y *Int) int {
        }
 }
 
-// ModSqrt sets z to a square root of x mod p if such a square root exists, and
-// returns z. The modulus p must be an odd prime. If x is not a square mod p,
-// ModSqrt leaves z unchanged and returns nil. This function panics if p is
-// not an odd integer.
-func (z *Int) ModSqrt(x, p *Int) *Int {
-       switch Jacobi(x, p) {
-       case -1:
-               return nil // x is not a square mod p
-       case 0:
-               return z.SetInt64(0) // sqrt(0) mod p = 0
-       case 1:
-               break
-       }
-       if x.neg || x.Cmp(p) >= 0 { // ensure 0 <= x < p
-               x = new(Int).Mod(x, p)
-       }
+// modSqrt3Mod4 uses the identity
+//      (a^((p+1)/4))^2  mod p
+//   == u^(p+1)          mod p
+//   == u^2              mod p
+// to calculate the square root of any quadratic residue mod p quickly for 3
+// mod 4 primes.
+func (z *Int) modSqrt3Mod4Prime(x, p *Int) *Int {
+       z.Set(p)         // z = p
+       z.Add(z, intOne) // z = p + 1
+       z.Rsh(z, 2)      // z = (p + 1) / 4
+       z.Exp(x, z, p)   // z = x^z mod p
+       return z
+}
 
+// modSqrtTonelliShanks uses the Tonelli-Shanks algorithm to find the square
+// root of a quadratic residue modulo any prime.
+func (z *Int) modSqrtTonelliShanks(x, p *Int) *Int {
        // Break p-1 into s*2^e such that s is odd.
        var s Int
        s.Sub(p, intOne)
@@ -703,6 +703,31 @@ func (z *Int) ModSqrt(x, p *Int) *Int {
        }
 }
 
+// ModSqrt sets z to a square root of x mod p if such a square root exists, and
+// returns z. The modulus p must be an odd prime. If x is not a square mod p,
+// ModSqrt leaves z unchanged and returns nil. This function panics if p is
+// not an odd integer.
+func (z *Int) ModSqrt(x, p *Int) *Int {
+       switch Jacobi(x, p) {
+       case -1:
+               return nil // x is not a square mod p
+       case 0:
+               return z.SetInt64(0) // sqrt(0) mod p = 0
+       case 1:
+               break
+       }
+       if x.neg || x.Cmp(p) >= 0 { // ensure 0 <= x < p
+               x = new(Int).Mod(x, p)
+       }
+
+       // Check whether p is 3 mod 4, and if so, use the faster algorithm.
+       if len(p.abs) > 0 && p.abs[0]%4 == 3 {
+               return z.modSqrt3Mod4Prime(x, p)
+       }
+       // Otherwise, use Tonelli-Shanks.
+       return z.modSqrtTonelliShanks(x, p)
+}
+
 // Lsh sets z = x << n and returns z.
 func (z *Int) Lsh(x *Int, n uint) *Int {
        z.abs = z.abs.shl(x.abs, n)

diff --git a/src/math/big/int_test.go b/src/math/big/int_test.go
index 97874626f336fddee773f0369673faeb2aaba89c..5b80509a318f59f75e5d7bac01984625c24fd707 100644 (file)
--- a/src/math/big/int_test.go
+++ b/src/math/big/int_test.go
@@ -1185,6 +1185,53 @@ func BenchmarkBitsetNegOrig(b *testing.B) {
        }
 }
 
+// tri generates the trinomial 2**(n*2) - 2**n - 1, which is always 3 mod 4 and
+// 7 mod 8, so that 2 is always a quadratic residue.
+func tri(n uint) *Int {
+       x := NewInt(1)
+       x.Lsh(x, n)
+       x2 := new(Int).Lsh(x, n)
+       x2.Sub(x2, x)
+       x2.Sub(x2, intOne)
+       return x2
+}
+
+func BenchmarkModSqrt225_Tonelli(b *testing.B) {
+       p := tri(225)
+       x := NewInt(2)
+       for i := 0; i < b.N; i++ {
+               x.SetUint64(2)
+               x.modSqrtTonelliShanks(x, p)
+       }
+}
+
+func BenchmarkModSqrt224_3Mod4(b *testing.B) {
+       p := tri(225)
+       x := new(Int).SetUint64(2)
+       for i := 0; i < b.N; i++ {
+               x.SetUint64(2)
+               x.modSqrt3Mod4Prime(x, p)
+       }
+}
+
+func BenchmarkModSqrt5430_Tonelli(b *testing.B) {
+       p := tri(5430)
+       x := new(Int).SetUint64(2)
+       for i := 0; i < b.N; i++ {
+               x.SetUint64(2)
+               x.modSqrtTonelliShanks(x, p)
+       }
+}
+
+func BenchmarkModSqrt5430_3Mod4(b *testing.B) {
+       p := tri(5430)
+       x := new(Int).SetUint64(2)
+       for i := 0; i < b.N; i++ {
+               x.SetUint64(2)
+               x.modSqrt3Mod4Prime(x, p)
+       }
+}
+
 func TestBitwise(t *testing.T) {
        x := new(Int)
        y := new(Int)