Function sqrtModPrime

src/compute-engine/numerics/diophantine.ts:427–469

* Tonelli–Shanks: a single square root of `a` modulo an ODD PRIME `p`, in * `[0, p)`, or `null` if `a` is a quadratic non-residue. The other root is * `p − r`.

(a: bigint, p: bigint)

 * `p − r`.
 */
function sqrtModPrime(a: bigint, p: bigint): bigint | null {
  a = ((a % p) + p) % p;
  if (a === 0n) return 0n;
  if (p === 2n) return a & 1n;
  // Euler's criterion.
  if (modPow(a, (p - 1n) / 2n, p) !== 1n) return null;
  if (p % 4n === 3n) return modPow(a, (p + 1n) / 4n, p);

  // Factor p − 1 = q·2^s with q odd.
  let q = p - 1n;
  let s = 0n;
  while (q % 2n === 0n) {
    q /= 2n;
    s += 1n;
  }
  // Find a quadratic non-residue z.
  let z = 2n;
  while (modPow(z, (p - 1n) / 2n, p) !== p - 1n) z += 1n;

  let m = s;
  let cc = modPow(z, q, p);
  let t = modPow(a, q, p);
  let r = modPow(a, (q + 1n) / 2n, p);
  let guard = 0;
  while (t !== 1n) {
    if (++guard > MAX_ITERATIONS) throw new DiophantineBudgetError();
    // Least i, 0 < i < m, with t^(2^i) = 1.
    let i = 0n;
    let t2 = t;
    while (t2 !== 1n) {
      t2 = (t2 * t2) % p;
      i += 1n;
      if (i === m) return null; // should not happen for a QR
    }
    let b = cc;
    for (let j = 0n; j < m - i - 1n; j++) b = (b * b) % p;
    m = i;
    cc = (b * b) % p;
    t = (t * cc) % p;
    r = (r * b) % p;
  }
  return r;
}

/**
 * All roots of `x² ≡ a (mod pᵏ)` for an ODD prime `p`, `k ≥ 1`, sorted