Fix sqrt

src/abstract/modular.ts

@@ -81,7 +81,7 @@ export function tonelliShanks(P: bigint) {
 
   let Q: bigint, S: number, Z: bigint;
   // Step 1: By factoring out powers of 2 from p - 1,
-  // find q and s such that p - 1 = q2s with q odd
+  // find q and s such that p - 1 = q*(2^s) with q odd
   for (Q = P - _1n, S = 0; Q % _2n === _0n; Q /= _2n, S++);
 
   // Step 2: Select a non-square z such that (z | p) ≡ -1 and set c ≡ zq
@@ -100,22 +100,24 @@ export function tonelliShanks(P: bigint) {
   // Slow-path
   const Q1div2 = (Q + _1n) / _2n;
   return function tonelliSlow<T>(Fp: Field<T>, n: T): T {
-    // Step 0: Check that n is indeed a square: (n | p) must be ≡ 1
-    if (Fp.pow(n, legendreC) !== Fp.ONE) throw new Error('Cannot find square root');
+    // Step 0: Check that n is indeed a square: (n | p) should not be ≡ -1
+    if (Fp.pow(n, legendreC) === Fp.negate(Fp.ONE)) throw new Error('Cannot find square root');
     let r = S;
-    let g = Fp.pow(Fp.create(Z as any as T), Q); // will update both x and b
+    // TODO: will fail at Fp2/etc
+    let g = Fp.pow(Fp.mul(Fp.ONE, Z), Q); // will update both x and b
     let x = Fp.pow(n, Q1div2); // first guess at the square root
     let b = Fp.pow(n, Q); // first guess at the fudge factor
 
-    let t2: typeof Fp.ZERO;
-    while (!Fp.equals(Fp.sub(b, Fp.ONE), Fp.ZERO)) {
-      t2 = Fp.square(b);
-      let m;
-      for (m = 1; m < r; m++) {
+    while (!Fp.equals(b, Fp.ONE)) {
+      if (Fp.equals(b, Fp.ZERO)) return Fp.ZERO; // https://en.wikipedia.org/wiki/Tonelli%E2%80%93Shanks_algorithm (4. If t = 0, return r = 0)
+      // Find m such b^(2^m)==1
+      let m = 1;
+      for (let t2 = Fp.square(b); m < r; m++) {
         if (Fp.equals(t2, Fp.ONE)) break;
         t2 = Fp.square(t2); // t2 *= t2
       }
-      let ge = Fp.pow(g, BigInt(1 << (r - m - 1))); // ge = 2^(r-m-1)
+      // NOTE: r-m-1 can be bigger than 32, need to convert to bigint before shift, otherwise there will be overflow
+      const ge = Fp.pow(g, _1n << BigInt(r - m - 1)); // ge = 2^(r-m-1)
       g = Fp.square(ge); // g = ge * ge
       x = Fp.mul(x, ge); // x *= ge
       b = Fp.mul(b, g); // b *= g

test/basic.test.js

@@ -138,6 +138,9 @@ for (const c in FIELDS) {
         })
       );
     });
+    should(`${name} negate(0)`, () => {
+      deepStrictEqual(Fp.negate(Fp.ZERO), Fp.ZERO);
+    });
 
     should(`${name} multiply/commutativity`, () => {
       fc.assert(
@@ -201,9 +204,27 @@ for (const c in FIELDS) {
         })
       );
     });
+
+    should(`${name} square(0)`, () => {
+      deepStrictEqual(Fp.square(Fp.ZERO), Fp.ZERO);
+      deepStrictEqual(Fp.mul(Fp.ZERO, Fp.ZERO), Fp.ZERO);
+    });
+
+    should(`${name} square(1)`, () => {
+      deepStrictEqual(Fp.square(Fp.ONE), Fp.ONE);
+      deepStrictEqual(Fp.mul(Fp.ONE, Fp.ONE), Fp.ONE);
+    });
+
+    should(`${name} square(-1)`, () => {
+      const minus1 = Fp.negate(Fp.ONE);
+      deepStrictEqual(Fp.square(minus1), Fp.ONE);
+      deepStrictEqual(Fp.mul(minus1, minus1), Fp.ONE);
+    });
+
     const isSquare = mod.FpIsSquare(Fp);
+    // Not implemented
+    if (Fp !== bls12_381.CURVE.Fp12) {
       should(`${name} multiply/sqrt`, () => {
-      if (Fp === bls12_381.CURVE.Fp12) return; // Not implemented
         fc.assert(
           fc.property(FC_BIGINT, (num) => {
             const a = create(num);
@@ -221,6 +242,17 @@ for (const c in FIELDS) {
         );
       });
 
+      should(`${name} sqrt(0)`, () => {
+        deepStrictEqual(Fp.sqrt(Fp.ZERO), Fp.ZERO);
+        const sqrt1 = Fp.sqrt(Fp.ONE);
+        deepStrictEqual(
+          Fp.equals(sqrt1, Fp.ONE) || Fp.equals(sqrt1, Fp.negate(Fp.ONE)),
+          true,
+          'sqrt(1) = 1 or -1'
+        );
+      });
+    }
+
     should(`${name} div/division by one equality`, () => {
       fc.assert(
         fc.property(FC_BIGINT, (num) => {
@@ -560,6 +592,22 @@ for (const name in CURVES) {
     });
   }
 }
+
+should('Secp224k1 sqrt bug', () => {
+  const { Fp } = secp224r1.CURVE;
+  const sqrtMinus1 = Fp.sqrt(-1n);
+  // Verified against sage
+  deepStrictEqual(
+    sqrtMinus1,
+    23621584063597419797792593680131996961517196803742576047493035507225n
+  );
+  deepStrictEqual(
+    Fp.negate(sqrtMinus1),
+    3338362603553219996874421406887633712040719456283732096017030791656n
+  );
+  deepStrictEqual(Fp.square(sqrtMinus1), Fp.create(-1n));
+});
+
 // ESM is broken.
 import url from 'url';
 if (import.meta.url === url.pathToFileURL(process.argv[1]).href) {