Modular arithmetic and fast exponentiation

import math
import random

from modpow import inverse_mod, pow_mod

P = 10**9 + 7


def test_small_powers():
    """small cases, including a base larger than the modulus and exponent 0"""
    assert pow_mod(7, 45, 1000) == 807
    assert pow_mod(2, 10, 1000) == 24
    assert pow_mod(10, 3, 7) == 6
    assert pow_mod(5, 0, 7) == 1
    assert pow_mod(0, 0, 7) == 1
    assert pow_mod(5, 0, 1) == 0
    assert pow_mod(0, 5, 7) == 0


def test_exponents_near_ten_to_the_eighteen():
    """exponents up to 10^18, a prime modulus and a modulus near 10^18"""
    assert pow_mod(2, 10**18, P) == 719476260
    assert pow_mod(3, 10**18 - 1, 10**18 + 9) == pow(3, 10**18 - 1, 10**18 + 9)
    assert pow_mod(10**18 + 8, 2, 10**18 + 9) == 1


def test_random_powers():
    """500 random cases compared with Python's pow"""
    rng = random.Random(21)
    for _ in range(500):
        a, b, m = rng.randrange(0, 10**18), rng.randrange(0, 10**18), rng.randrange(1, 10**18)
        assert pow_mod(a, b, m) == pow(a, b, m), (a, b, m)


def test_small_inverses():
    """inverses that exist, and two that do not"""
    assert inverse_mod(3, 7) == 5
    assert inverse_mod(10, 7) == 5
    assert inverse_mod(7, P) == 142857144
    assert inverse_mod(2, 9) == 5
    assert inverse_mod(1, 2) == 1
    assert inverse_mod(2, 10) is None
    assert inverse_mod(6, 9) is None
    assert inverse_mod(0, 5) is None


def test_every_small_modulus():
    """every a below m for every m from 2 to 60, against a search"""
    for m in range(2, 61):
        for a in range(m):
            found = next((x for x in range(m) if a * x % m == 1), None)
            assert inverse_mod(a, m) == found, (a, m)


def test_moduli_that_are_not_prime():
    """Fermat's a^(m - 2) fails here; the extended algorithm does not"""
    m = 10**18
    for a in (999_999_999_999_999_999, 123_456_789_012_345_677, 3):
        x = inverse_mod(a, m)
        assert x is not None and 0 <= x < m and a * x % m == 1
    assert inverse_mod(10**9, m) is None
    rng = random.Random(8)
    for _ in range(300):
        m = rng.randrange(2, 10**12)
        a = rng.randrange(0, m)
        x = inverse_mod(a, m)
        if math.gcd(a, m) == 1:
            assert x == pow(a, -1, m), (a, m)
        else:
            assert x is None, (a, m)

import { test, assert } from 'mysmartcopilot:test';
import { inverseMod, powMod } from './modpow.mjs';

const P = 1000000007n;

/** Exact a^b mod m with BigInt, by plain repeated squaring of a different shape (right to left over a bit string). */
function referencePow(a, b, m) {
  let result = 1n % m;
  const bits = b.toString(2);
  for (const bit of bits) {
    result = (result * result) % m;
    if (bit === '1') result = (result * a) % m;
  }
  return result;
}

/** A small seeded generator of BigInts below limit (a 64-bit linear congruential generator). */
function generator(seed) {
  let state = BigInt(seed);
  return (limit) => {
    state = (state * 6364136223846793005n + 1442695040888963407n) % (1n << 64n);
    return state % limit;
  };
}

function gcd(a, b) {
  while (b) [a, b] = [b, a % b];
  return a;
}

test('small cases, including a base larger than the modulus and exponent 0', () => {
  assert.equal(powMod(7n, 45n, 1000n), 807n);
  assert.equal(powMod(2n, 10n, 1000n), 24n);
  assert.equal(powMod(10n, 3n, 7n), 6n);
  assert.equal(powMod(5n, 0n, 7n), 1n);
  assert.equal(powMod(0n, 0n, 7n), 1n);
  assert.equal(powMod(5n, 0n, 1n), 0n);
  assert.equal(powMod(0n, 5n, 7n), 0n);
});

test('exponents up to 10^18, a prime modulus and a modulus near 10^18', () => {
  assert.equal(powMod(2n, 10n ** 18n, P), 719476260n);
  assert.equal(powMod(10n ** 18n + 8n, 2n, 10n ** 18n + 9n), 1n);
});

test('300 random cases compared with exact BigInt arithmetic', () => {
  const next = generator(21);
  for (let i = 0; i < 300; i++) {
    const a = next(10n ** 18n);
    const b = next(10n ** 18n);
    const m = next(10n ** 18n - 1n) + 1n;
    assert.equal(powMod(a, b, m), referencePow(a % m, b, m), `${a} ${b} ${m}`);
  }
});

test('inverses that exist, and two that do not', () => {
  assert.equal(inverseMod(3n, 7n), 5n);
  assert.equal(inverseMod(10n, 7n), 5n);
  assert.equal(inverseMod(7n, P), 142857144n);
  assert.equal(inverseMod(2n, 9n), 5n);
  assert.equal(inverseMod(1n, 2n), 1n);
  assert.equal(inverseMod(2n, 10n), null);
  assert.equal(inverseMod(6n, 9n), null);
  assert.equal(inverseMod(0n, 5n), null);
});

test('every a below m for every m from 2 to 60, against a search', () => {
  for (let m = 2n; m <= 60n; m++) {
    for (let a = 0n; a < m; a++) {
      let found = null;
      for (let x = 0n; x < m; x++) if ((a * x) % m === 1n) { found = x; break; }
      assert.equal(inverseMod(a, m), found, `${a} ${m}`);
    }
  }
});

test("Fermat's a^(m - 2) fails here; the extended algorithm does not", () => {
  const m = 10n ** 18n;
  for (const a of [999999999999999999n, 123456789012345677n, 3n]) {
    const x = inverseMod(a, m);
    assert.ok(x !== null && x >= 0n && x < m && (a * x) % m === 1n, `${a}`);
  }
  assert.equal(inverseMod(10n ** 9n, m), null);
  const next = generator(8);
  for (let i = 0; i < 200; i++) {
    const mod = next(10n ** 12n - 2n) + 2n;
    const a = next(mod);
    const x = inverseMod(a, mod);
    if (gcd(a, mod) === 1n) assert.ok(x !== null && (a * x) % mod === 1n && x < mod, `${a} ${mod}`);
    else assert.equal(x, null, `${a} ${mod}`);
  }
});