From fbbbae8381db60fbbafc36a7495704a42eadafea Mon Sep 17 00:00:00 2001
From: MITSUNARI Shigeo <herumi@nifty.com>
Date: Thu, 24 Oct 2024 12:07:25 +0900
Subject: [PATCH] add squareRoot

---
 Makefile           |  1 +
 src/value-types.ts | 25 +++++++++++++++++++++++++
 test/bench.js      |  6 ++++++
 test/test-ts.ts    | 26 ++++++++++++++++++++++++++
 4 files changed, 58 insertions(+)

diff --git a/Makefile b/Makefile
index f7b2833..65552e1 100644
--- a/Makefile
+++ b/Makefile
@@ -14,6 +14,7 @@ EMCC_OPT+=-DCYBOZU_MINIMUM_EXCEPTION
 EMCC_OPT+=-s ABORTING_MALLOC=0
 EMCC_OPT+=-s STACK_SIZE=1MB
 EMCC_OPT+=-sEXPORTED_FUNCTIONS=_malloc,_free,stackAlloc,stackSave,stackRestore
+#EMCC_OPT+=-msimd128
 EMCC_OPT+=-flto
 
 all: $(MCL_JS)
diff --git a/src/value-types.ts b/src/value-types.ts
index a276f56..0d3ce3a 100644
--- a/src/value-types.ts
+++ b/src/value-types.ts
@@ -135,6 +135,18 @@ abstract class Common {
     return z
   }
 
+  /** @internal r = func(y, this) and return (true, b) if r = true else (false, null) */
+  _squareRoot (func: (yPos: number, xPos: number) => number): [boolean, any] {
+    const y = new (this.constructor as any)()
+    const stack = mod.stackSave()
+    const xPos = this._sallocAndCopy()
+    const yPos = y._salloc()
+    const r = func(yPos, xPos)
+    y._save(yPos)
+    mod.stackRestore(stack)
+    return r === 0 ? [true, y] : [false, null]
+  }
+
   /** @internal devide Uint32Array a into n and chose the idx-th */
   _getSubArray (idx: number, n: number): Uint32Array {
     const d = this.a_.length / n
@@ -753,6 +765,19 @@ export const inv = <T extends Fp | Fr | GT | Fp2>(x: T): T => {
   throw new Error('inv:bad type')
 }
 
+export const squareRoot = <T extends Fp | Fr | Fp2>(x: T): [boolean, T] => {
+  if (x instanceof Fp) {
+    return x._squareRoot(mod._mclBnFp_squareRoot)
+  }
+  if (x instanceof Fr) {
+    return x._squareRoot(mod._mclBnFr_squareRoot)
+  }
+  if (x instanceof Fp2) {
+    return x._squareRoot(mod._mclBnFp2_squareRoot)
+  }
+  throw new Error('squareRoot:bad type')
+}
+
 export const normalize = <T extends G1 | G2>(x: T): T => {
   if (x instanceof G1) {
     return x._op1(mod._mclBnG1_normalize)
diff --git a/test/bench.js b/test/bench.js
index 1cc15a5..48e9e46 100644
--- a/test/bench.js
+++ b/test/bench.js
@@ -624,6 +624,8 @@ function benchAll () {
     bench('Fp::mul', C2, () => { b = mcl.mul(b, a) })
     bench('Fp::sqr', C2, () => { b = mcl.sqr(b) })
     bench('Fp::inv', C2, () => { b = mcl.inv(b) })
+    b = mcl.sqr(b)
+    bench('Fp::squareRoot', C, () => { mcl.squareRoot(b) })
   }
   {
     const a = new mcl.Fp2()
@@ -636,6 +638,8 @@ function benchAll () {
     bench('Fp2::mul', C2, () => { b = mcl.mul(b, a) })
     bench('Fp2::sqr', C2, () => { b = mcl.sqr(b) })
     bench('Fp2::inv', C2, () => { b = mcl.inv(b) })
+    b = mcl.sqr(b)
+    bench('Fp2::squareRoot', C, () => { mcl.squareRoot(b) })
   }
   {
     let b = new mcl.Fr()
@@ -646,6 +650,8 @@ function benchAll () {
     bench('Fr::mul', C2, () => { b = mcl.mul(b, a) })
     bench('Fr::sqr', C2, () => { b = mcl.sqr(b) })
     bench('Fr::inv', C2, () => { b = mcl.inv(b) })
+    b = mcl.sqr(b)
+    bench('Fr::squareRoot', C, () => { mcl.squareRoot(b) })
   }
 
   {
diff --git a/test/test-ts.ts b/test/test-ts.ts
index b955076..9decea6 100644
--- a/test/test-ts.ts
+++ b/test/test-ts.ts
@@ -129,6 +129,29 @@ function invVecTest(cstr: any): void {
 //  bench('invVecTime', 1000, () => { w = mcl.invVec(w) })
 }
 
+function squareRootTest1 (x, expectBool = undefined): boolean {
+  const [t, y] = mcl.squareRoot(x)
+  if (expectBool !== undefined) {
+    assert(t === expectBool)
+  }
+  if (t) {
+    assert(mcl.sqr(y).isEqual(x))
+  } else {
+    assert(y === null)
+  }
+  return t
+}
+
+function squareRootTest(x: any): void {
+  const n = 100
+  let d = x.clone()
+  for (let i = 0; i < n; i++) {
+    const dd = mcl.sqr(d)
+    squareRootTest1(dd, true)
+    d = dd
+  }
+}
+
 function FrTest() {
   const a = new mcl.Fr()
   a.setInt(5)
@@ -189,6 +212,7 @@ function FrTest() {
   }
   powTest(mcl.Fr)
   invVecTest(mcl.Fr)
+  squareRootTest(a)
 }
 
 function FpTest() {
@@ -237,6 +261,7 @@ function FpTest() {
   }
   powTest(mcl.Fp)
   invVecTest(mcl.Fp)
+  squareRootTest(a)
 }
 
 function Fp2Test() {
@@ -307,6 +332,7 @@ function Fp2Test() {
   a.setInt(-5)
   z.set_b(mcl.div(a, b))
   assert(mcl.inv(x).isEqual(z))
+  squareRootTest(a)
 }
 
 function normalizeVecTest(cstr: any): void {