From 4e395bccbe53bbe708c09fa05d9cfef3ff2535a6 Mon Sep 17 00:00:00 2001
From: MITSUNARI Shigeo <herumi@nifty.com>
Date: Mon, 15 Aug 2016 13:03:36 +0900
Subject: [PATCH] add F2T::squareRoot

---
 include/mcl/fp_tower.hpp | 38 ++++++++++++++++++++++++++++++++++++--
 include/mcl/gmp_util.hpp |  5 +++++
 test/bn_test.cpp         | 14 ++++++++++++++
 test/fp_tower_test.cpp   | 20 ++++++++++++++++++++
 test/sq_test.cpp         |  2 +-
 5 files changed, 76 insertions(+), 3 deletions(-)

diff --git a/include/mcl/fp_tower.hpp b/include/mcl/fp_tower.hpp
index 0497cd1..556f218 100644
--- a/include/mcl/fp_tower.hpp
+++ b/include/mcl/fp_tower.hpp
@@ -114,9 +114,43 @@ public:
 		isOdd() is not good naming. QQQ
 	*/
 	bool isOdd() const { return a.isOdd(); }
-	static inline bool squareRoot(Fp2T& /*y*/, const Fp2T& /*x*/)
+	/*
+		(a + bi)^2 = (a^2 - b^2) + 2ab i = c + di
+		A = a^2
+		B = b^2
+		A = (c +/- sqrt(c^2 + d^2))/2
+		b = d / 2a
+	*/
+	static inline bool squareRoot(Fp2T& y, const Fp2T& x)
 	{
-		throw cybozu::Exception("Fp2T:squareRoot:not supported");
+		Fp t1, t2;
+		if (x.b.isZero()) {
+			if (Fp::squareRoot(t1, x.a)) {
+				y.a = t1;
+				y.b.clear();
+			} else {
+				if (!Fp::squareRoot(t1, -x.a)) throw cybozu::Exception("Fp2T:squareRoot:internal error1") << x;
+				y.a.clear();
+				y.b = t1;
+			}
+			return true;
+		}
+		Fp::sqr(t1, x.a);
+		Fp::sqr(t2, x.b);
+		t1 += t2; // c^2 + d^2
+		if (!Fp::squareRoot(t1, t1)) return false;
+		Fp::add(t2, x.a, t1);
+		Fp::divBy2(t2, t2);
+		if (!Fp::squareRoot(t2, t2)) {
+			Fp::sub(t2, x.a, t1);
+			Fp::divBy2(t2, t2);
+			if (!Fp::squareRoot(t2, t2)) throw cybozu::Exception("Fp2T:squareRoot:internal error2") << x;
+		}
+		y.a = t2;
+		t2 += t2;
+		Fp::inv(t2, t2);
+		Fp::mul(y.b, x.b, t2);
+		return true;
 	}
 
 	static const Fp& getXi_a() { return xi_a_; }
diff --git a/include/mcl/gmp_util.hpp b/include/mcl/gmp_util.hpp
index 673da5a..8d1cab3 100644
--- a/include/mcl/gmp_util.hpp
+++ b/include/mcl/gmp_util.hpp
@@ -1,4 +1,5 @@
 #pragma once
+#include <iostream>
 /**
 	@file
 	@brief util function for gmp
@@ -340,6 +341,10 @@ public:
 	bool get(mpz_class& x, const mpz_class& a) const
 	{
 		if (!isPrime) throw cybozu::Exception("SquareRoot:get:not prime") << p;
+		if (a == 0) {
+			x = 0;
+			return true;
+		}
 		if (gmp::legendre(a, p) < 0) return false;
 		if (r == 1) {
 			// (p + 1) / 4 = (q + 1) / 2
diff --git a/test/bn_test.cpp b/test/bn_test.cpp
index 640fc2c..9728ba8 100644
--- a/test/bn_test.cpp
+++ b/test/bn_test.cpp
@@ -108,6 +108,20 @@ CYBOZU_TEST_AUTO(MapTo)
 	CYBOZU_TEST_EXCEPTION(mapTo.calcG1(g, -mapTo.c1), cybozu::Exception);
 }
 
+CYBOZU_TEST_AUTO(stream)
+{
+	const TestSet& ts = g_testSetTbl[0];
+	G2 Q0(Fp2(ts.g2.aa, ts.g2.ab), Fp2(ts.g2.ba, ts.g2.bb));
+	G2::setCompressedExpression();
+	G2 Q;
+	for (int i = 0; i < 10; i++) {
+		G2 R;
+		R.setStr(Q.getStr());
+		CYBOZU_TEST_EQUAL(Q, R);
+		G2::add(Q, Q, Q0);
+	}
+}
+
 int main(int argc, char *argv[])
 	try
 {
diff --git a/test/fp_tower_test.cpp b/test/fp_tower_test.cpp
index 0a32f08..8267cec 100644
--- a/test/fp_tower_test.cpp
+++ b/test/fp_tower_test.cpp
@@ -120,6 +120,26 @@ void testFp2()
 	inv(y, x);
 	y *= x;
 	CYBOZU_TEST_EQUAL(y, 1);
+
+	// square root
+	for (int i = 0; i < 3; i++) {
+		x.a = i * i + i * 2;
+		x.b = i;
+		Fp2::sqr(y, x);
+		CYBOZU_TEST_ASSERT(Fp2::squareRoot(z, y));
+		CYBOZU_TEST_EQUAL(z * z, y);
+		CYBOZU_TEST_ASSERT(Fp2::squareRoot(y, y));
+		CYBOZU_TEST_EQUAL(z * z, y * y);
+		x.b = 0;
+		Fp2::sqr(y, x);
+		CYBOZU_TEST_ASSERT(Fp2::squareRoot(z, y));
+		CYBOZU_TEST_EQUAL(z * z, y);
+		x.a = 0;
+		x.b = i * i + i * 3;
+		Fp2::sqr(y, x);
+		CYBOZU_TEST_ASSERT(Fp2::squareRoot(z, y));
+		CYBOZU_TEST_EQUAL(z * z, y);
+	}
 }
 
 void testFp6sqr(const Fp2& a, const Fp2& b, const Fp2& c, const Fp6& x)
diff --git a/test/sq_test.cpp b/test/sq_test.cpp
index 6174be6..4c386d2 100644
--- a/test/sq_test.cpp
+++ b/test/sq_test.cpp
@@ -9,7 +9,7 @@ CYBOZU_TEST_AUTO(sqrt)
 	for (size_t i = 0; i < CYBOZU_NUM_OF_ARRAY(tbl); i++) {
 		const mpz_class p = tbl[i];
 		sq.set(p);
-		for (mpz_class a = 1; a < p; a++) {
+		for (mpz_class a = 0; a < p; a++) {
 			mpz_class x;
 			if (sq.get(x, a)) {
 				mpz_class y;