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>

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
 }
 

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