use precomputed value of SquareRoot for BN254 and BLS12_381

4 files changed, 139 insertions, 18 deletions

diff --git a/include/mcl/gmp_util.hpp b/include/mcl/gmp_util.hpp
index 4ae95a3..6db9b4a 100644
--- a/include/mcl/gmp_util.hpp
+++ b/include/mcl/gmp_util.hpp
@@ -622,6 +622,7 @@ inline void getRandPrime(mpz_class& z, size_t bitSize, fp::RandGen rg = fp::Rand
 	Tonelli-Shanks
 */
 class SquareRoot {
+	bool isPrecomputed_;
 	bool isPrime;
 	mpz_class p;
 	mpz_class g;
@@ -629,10 +630,78 @@ class SquareRoot {
 	mpz_class q; // p - 1 = 2^r q
 	mpz_class s; // s = g^q
 	mpz_class q_add_1_div_2;
+	struct Tbl {
+		const char *p;
+		const char *g;
+		int r;
+		const char *q;
+		const char *s;
+		const char *q_add_1_div_2;
+	};
+	bool setIfPrecomputed(const mpz_class& p_)
+	{
+		static const Tbl tbl[] = {
+			{ // BN254.p
+				"2523648240000001ba344d80000000086121000000000013a700000000000013",
+				"2",
+				1,
+				"1291b24120000000dd1a26c0000000043090800000000009d380000000000009",
+				"2523648240000001ba344d80000000086121000000000013a700000000000012",
+				"948d920900000006e8d1360000000021848400000000004e9c0000000000005",
+			},
+			{ // BN254.r
+				"2523648240000001ba344d8000000007ff9f800000000010a10000000000000d",
+				"2",
+				2,
+				"948d920900000006e8d136000000001ffe7e000000000042840000000000003",
+				"9366c4800000000555150000000000122400000000000015",
+				"4a46c9048000000374689b000000000fff3f000000000021420000000000002",
+			},
+			{ // BLS12_381,p
+				"1a0111ea397fe69a4b1ba7b6434bacd764774b84f38512bf6730d2a0f6b0f6241eabfffeb153ffffb9feffffffffaaab",
+				"2",
+				1,
+				"d0088f51cbff34d258dd3db21a5d66bb23ba5c279c2895fb39869507b587b120f55ffff58a9ffffdcff7fffffffd555",
+				"1a0111ea397fe69a4b1ba7b6434bacd764774b84f38512bf6730d2a0f6b0f6241eabfffeb153ffffb9feffffffffaaaa",
+				"680447a8e5ff9a692c6e9ed90d2eb35d91dd2e13ce144afd9cc34a83dac3d8907aaffffac54ffffee7fbfffffffeaab",
+			},
+			{ // BLS12_381.r
+				"73eda753299d7d483339d80809a1d80553bda402fffe5bfeffffffff00000001",
+				"5",
+				32,
+				"73eda753299d7d483339d80809a1d80553bda402fffe5bfeffffffff",
+				"212d79e5b416b6f0fd56dc8d168d6c0c4024ff270b3e0941b788f500b912f1f",
+				"39f6d3a994cebea4199cec0404d0ec02a9ded2017fff2dff80000000",
+			},
+		};
+		for (size_t i = 0; i < CYBOZU_NUM_OF_ARRAY(tbl); i++) {
+			mpz_class targetPrime;
+			bool b;
+			mcl::gmp::setStr(&b, targetPrime, tbl[i].p, 16);
+			if (!b) continue;
+			if (targetPrime != p_) continue;
+			isPrime = true;
+			p = p_;
+			mcl::gmp::setStr(&b, g, tbl[i].g, 16);
+			if (!b) continue;
+			r = tbl[i].r;
+			mcl::gmp::setStr(&b, q, tbl[i].q, 16);
+			if (!b) continue;
+			mcl::gmp::setStr(&b, s, tbl[i].s, 16);
+			if (!b) continue;
+			mcl::gmp::setStr(&b, q_add_1_div_2, tbl[i].q_add_1_div_2, 16);
+			if (!b) continue;
+			isPrecomputed_ = true;
+			return true;
+		}
+		return false;
+	}
 public:
 	SquareRoot() { clear(); }
+	bool isPrecomputed() const { return isPrecomputed_; }
 	void clear()
 	{
+		isPrecomputed_ = false;
 		isPrime = false;
 		p = 0;
 		g = 0;
@@ -641,8 +710,23 @@ public:
 		s = 0;
 		q_add_1_div_2 = 0;
 	}
-	void set(bool *pb, const mpz_class& _p)
+#ifndef MCL_USE_VINT
+	void dump() const
+	{
+		printf("\"%s\",\n", mcl::gmp::getStr(p, 16).c_str());
+		printf("\"%s\",\n", mcl::gmp::getStr(g, 16).c_str());
+		printf("%d,\n", r);
+		printf("\"%s\",\n", mcl::gmp::getStr(q, 16).c_str());
+		printf("\"%s\",\n", mcl::gmp::getStr(s, 16).c_str());
+		printf("\"%s\",\n", mcl::gmp::getStr(q_add_1_div_2, 16).c_str());
+	}
+#endif
+	void set(bool *pb, const mpz_class& _p, bool usePrecomputedTable = true)
 	{
+		if (usePrecomputedTable && setIfPrecomputed(_p)) {
+			*pb = true;
+			return;
+		}
 		p = _p;
 		if (p <= 2) {
 			*pb = false;
@@ -761,6 +845,12 @@ public:
 		}
 		return true;
 	}
+	bool operator==(const SquareRoot& rhs) const
+	{
+		return isPrime == rhs.isPrime && p == rhs.p && g == rhs.g && r == rhs.r
+			&& q == rhs.q && s == rhs.s && q_add_1_div_2 == rhs.q_add_1_div_2;
+	}
+	bool operator!=(const SquareRoot& rhs) const { return !operator==(rhs); }
 #ifndef CYBOZU_DONT_USE_EXCEPTION
 	void set(const mpz_class& _p)
 	{
diff --git a/test/bench.hpp b/test/bench.hpp
index 9678d3c..009bad1 100644
--- a/test/bench.hpp
+++ b/test/bench.hpp
@@ -88,3 +88,47 @@ void testBench(const G1& P, const G2& Q)
 	precomputeG2(Qcoeff, Q);
 	CYBOZU_BENCH_C("precomputedML ", C, precomputedMillerLoop, e2, P, Qcoeff);
 }
+
+inline void SquareRootPrecomputeTest(const mpz_class& p)
+{
+	mcl::SquareRoot sq1, sq2;
+	bool b;
+	sq1.set(&b, p, true);
+	CYBOZU_TEST_ASSERT(b);
+	CYBOZU_TEST_ASSERT(sq1.isPrecomputed());
+	sq2.set(&b, p, false);
+	CYBOZU_TEST_ASSERT(sq1 == sq2);
+	if (sq1 != sq2) {
+		puts("dump");
+		puts("sq1");
+		sq1.dump();
+		puts("sq2");
+		sq2.dump();
+		puts("---");
+	}
+}
+
+void testSquareRoot()
+{
+	if (BN::param.cp == mcl::BN254 || BN::param.cp == mcl::BLS12_381) {
+		SquareRootPrecomputeTest(BN::param.p);
+		SquareRootPrecomputeTest(BN::param.r);
+	}
+}
+
+void testLagrange()
+{
+	puts("testLagrange");
+	const int k = 7;
+	Fr c[k], x[k], y[k];
+	for (size_t i = 0; i < k; i++) {
+		c[i].setByCSPRNG();
+		x[i].setByCSPRNG();
+	}
+	for (size_t i = 0; i < k; i++) {
+		mcl::evaluatePolynomial(y[i], c, k, x[i]);
+	}
+	Fr s;
+	mcl::LagrangeInterpolation(s, x, y, k);
+	CYBOZU_TEST_EQUAL(s, c[0]);
+}
diff --git a/test/bn384_test.cpp b/test/bn384_test.cpp
index 3874938..8bb45b2 100644
--- a/test/bn384_test.cpp
+++ b/test/bn384_test.cpp
@@ -13,23 +13,6 @@ mcl::fp::Mode g_mode;
 
 #include "bench.hpp"
 
-void testLagrange()
-{
-	puts("testLagrange");
-	const int k = 7;
-	Fr c[k], x[k], y[k];
-	for (size_t i = 0; i < k; i++) {
-		c[i].setByCSPRNG();
-		x[i].setByCSPRNG();
-	}
-	for (size_t i = 0; i < k; i++) {
-		mcl::evaluatePolynomial(y[i], c, k, x[i]);
-	}
-	Fr s;
-	mcl::LagrangeInterpolation(s, x, y, k);
-	CYBOZU_TEST_EQUAL(s, c[0]);
-}
-
 void testCurve(const mcl::CurveParam& cp)
 {
 	initPairing(cp, g_mode);
@@ -58,6 +41,7 @@ void testCurve(const mcl::CurveParam& cp)
 	GT::pow(e1, e1, a * b);
 	CYBOZU_TEST_EQUAL(e1, e2);
 	testBench(P, Q);
+	testSquareRoot();
 	testLagrange();
 }
 
diff --git a/test/bn512_test.cpp b/test/bn512_test.cpp
index dde8950..8736377 100644
--- a/test/bn512_test.cpp
+++ b/test/bn512_test.cpp
@@ -5,6 +5,7 @@
 #include <cybozu/xorshift.hpp>
 #include <mcl/bn512.hpp>
 #include <mcl/bn.hpp>
+#include <mcl/lagrange.hpp>
 
 using namespace mcl::bn512;
 
@@ -34,6 +35,8 @@ void testCurve(const mcl::CurveParam& cp)
 	GT::pow(e1, e1, a * b);
 	CYBOZU_TEST_EQUAL(e1, e2);
 	testBench(P, Q);
+	testSquareRoot();
+	testLagrange();
 }
 
 CYBOZU_TEST_AUTO(pairing)