math/big: fix ModSqrt optimized path for x = z

name                   old time/op  new time/op  delta
ModSqrt224_3Mod4-4      153µs ± 2%   154µs ± 1%   ~     (p=0.548 n=5+5)
ModSqrt5430_3Mod4-4     776ms ± 2%   791ms ± 2%   ~     (p=0.222 n=5+5)

Fixes #22265

Change-Id: If233542716e04341990a45a1c2b7118da6d233f7
Reviewed-on: https://go-review.googlesource.com/70832
Run-TryBot: Filippo Valsorda <hi@filippo.io>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Robert Griesemer <gri@golang.org>

diff --git a/src/math/big/int.go b/src/math/big/int.go
index 92e2ae954a1e5d34abd1f65851553babcac7bd4e..73d48deb81a0acc214dc8340f02ee1a64c53ab42 100644 (file)
--- a/src/math/big/int.go
+++ b/src/math/big/int.go
@@ -667,10 +667,9 @@ func Jacobi(x, y *Int) int {
 // 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
+       e := new(Int).Add(p, intOne) // e = p + 1
+       e.Rsh(e, 2)                  // e = (p + 1) / 4
+       z.Exp(x, e, p)               // z = x^e mod p
        return z
 }
 

diff --git a/src/math/big/int_test.go b/src/math/big/int_test.go
index 65e24f1e4bfb99f1d8e84bea32f99317d9bb1703..bc2eef5f76002b6f2ae9c063767b96533e48c390 100644 (file)
--- a/src/math/big/int_test.go
+++ b/src/math/big/int_test.go
@@ -1384,6 +1384,12 @@ func testModSqrt(t *testing.T, elt, mod, sq, sqrt *Int) bool {
                t.Errorf("ModSqrt returned inconsistent value %s", z)
        }
 
+       // test x aliasing z
+       z = sqrtChk.ModSqrt(sqrtChk.Set(sq), mod)
+       if z != &sqrtChk || z.Cmp(sqrt) != 0 {
+               t.Errorf("ModSqrt returned inconsistent value %s", z)
+       }
+
        // make sure we actually got a square root
        if sqrt.Cmp(elt) == 0 {
                return true // we found the "desired" square root