Perf: Optimize Sqrt in Fp2 for all fields (#757)

Author: yelhousni

ecc/bls12-377/fp/element.go

@@ -16,3 +16,140 @@ func (z *Element) MulByNonResidueInv(x *Element) *Element {
 	z.Mul(x, &qnrInv)
 	return z
 }
+
+// SqrtAndInverse computes both the square root and the inverse of x in E2 by
+// doing a single exponentiation to the power (p-2^e-1)/2^(e+1) where e is the
+// field 2-adicity.
+func (z *Element) SqrtAndInverse(x, inv *Element) (*Element, *Element) {
+	// q ≡ 1 (mod 4)
+	// see modSqrtTonelliShanks in math/big/int.go
+
+	var y, b, t, w Element
+	r := uint64(46)
+
+	// w = x^((s-1)/2))
+	w.ExpBySqrtExp(*x)
+
+	inv.expByC1(x)
+
+	y.Square(&w).
+		Square(&y).
+		Mul(&y, x)
+
+	for i := uint64(0); i < r-1; i++ {
+		y.Square(&y)
+	}
+	inv.Mul(inv, &y)
+
+	// y = x^((s+1)/2)) = w * x
+	y.Mul(x, &w)
+
+	// b = xˢ = w * w * x = y * x
+	b.Mul(&w, &y)
+
+	// g = nonResidue ^ s
+	var g = Element{
+		7563926049028936178,
+		2688164645460651601,
+		12112688591437172399,
+		3177973240564633687,
+		14764383749841851163,
+		52487407124055189,
+	}
+
+	// compute legendre symbol
+	// t = x^((q-1)/2) = r-1 squaring of xˢ
+	t = b
+	for i := uint64(0); i < r-1; i++ {
+		t.Square(&t)
+	}
+	if t.IsZero() {
+		return z.SetZero(), inv
+	}
+	if !t.IsOne() {
+		// t != 1, we don't have a square root
+		return nil, inv
+	}
+	for {
+		var m uint64
+		t = b
+
+		// for t != 1
+		for !t.IsOne() {
+			t.Square(&t)
+			m++
+		}
+
+		if m == 0 {
+			return z.Set(&y), inv
+		}
+		// t = g^(2^(r-m-1)) (mod q)
+		ge := int(r - m - 1)
+		t = g
+		for ge > 0 {
+			t.Square(&t)
+			ge--
+		}
+
+		g.Square(&t)
+		y.Mul(&y, &t)
+		b.Mul(&b, &g)
+		r = m
+	}
+}
+
+// expByC1 set z to x^c1 and return z
+// where c1 = 2^(e-1)-1 and e = 46 is the 2-adicity of Fp.
+func (z *Element) expByC1(x *Element) *Element {
+	// addition chain:
+	//
+	//	_10    = 2*1
+	//	_11    = 1 + _10
+	//	_1100  = _11 << 2
+	//	_1111  = _11 + _1100
+	//	_11110 = 2*_1111
+	//	_11111 = 1 + _11110
+	//	x10    = _11111 << 5 + _11111
+	//	x20    = x10 << 10 + x10
+	//	x40    = x20 << 20 + x20
+	//	return   x40 << 5 + _11111
+	//
+	// Operations: 44 squares 7 multiplies
+	//
+	// Generated by github.com/mmcloughlin/addchain v0.4.0.
+
+	var (
+		t0 = new(Element)
+		t1 = new(Element)
+	)
+	z.Square(x)
+	z.Mul(x, z)
+	t0.Square(z)
+	for s := 1; s < 2; s++ {
+		t0.Square(t0)
+	}
+	z.Mul(z, t0)
+	z.Square(z)
+	z.Mul(x, z)
+	t0.Square(z)
+	for s := 1; s < 5; s++ {
+		t0.Square(t0)
+	}
+	t0.Mul(z, t0)
+	t1.Square(t0)
+	for s := 1; s < 10; s++ {
+		t1.Square(t1)
+	}
+	t0.Mul(t0, t1)
+	t1.Square(t0)
+	for s := 1; s < 20; s++ {
+		t1.Square(t1)
+	}
+	t0.Mul(t0, t1)
+	for s := 0; s < 5; s++ {
+		t0.Square(t0)
+	}
+	z.Mul(z, t0)
+
+	return z
+}