Python calculate modular roots using libnum and tonelli-shanks

modroots.py

import libnum
from typing import List, Optional


def modular_square_roots(a: int, p: int) -> Optional[List[int]]:
    """Return all x such that x² ≡ a (mod p), or None if none exist."""
    if not libnum.has_sqrtmod_prime_power(a, p):
        return None
    return list(libnum.sqrtmod_prime_power(a, p))


def modular_cube_roots(a: int, p: int) -> List[int]:
    """
    Compute all x such that x³ ≡ a (mod p).
    Python3 port of Rolandb’s StackOverflow cubic-root routine.
    """
    def _find_nontrivial_rth_root(r: int) -> int:
        t = p - 2
        while pow(t, (p - 1) // r, p) == 1:
            t -= 1
        return pow(t, (p - 1) // r, p)

    def _all_conjugates(root: int) -> List[int]:
        g = _find_nontrivial_rth_root(3)
        return [root, (root * g) % p, (root * g * g) % p]

    a_mod = a % p

    # trivial and single‐root cases
    if p in (2, 3) or a_mod == 0:
        return [a_mod]
    if p % 3 == 2:
        return [pow(a_mod, ((2 * p) - 1) // 3, p)]
    if pow(a_mod, (p - 1) // 3, p) > 1:
        return []

    # small‐prime shortcuts
    for mod9, exp in ((4, ((2 * p) + 1) // 9),
                      (7, (p + 2) // 9)):
        if p % 9 == mod9:
            r = pow(a_mod, exp, p)
            return _all_conjugates(r) if pow(r, 3, p) == a_mod else []

    if p % 27 in (10, 19):
        offset = 7 if p % 27 == 10 else 8
        r = pow(a_mod, ((2 * p) + offset) // 27, p)
        nth9 = _find_nontrivial_rth_root(9)
        for _ in range(9):
            if pow(r, 3, p) == a_mod:
                return _all_conjugates(r)
            r = (r * nth9) % p
        return []

    # fallback to general Tonelli‐Shanks style
    return _tonelli_shanks_cubic(a_mod, p, all_roots=True)


def _tonelli_shanks_cubic(a: int, p: int, all_roots: bool = False) -> List[int]:
    """General cubic‐root via Tonelli–Shanks; returns one or all roots."""
    def _find_nontrivial_rth_root(r: int) -> int:
        t = p - 2
        while pow(t, (p - 1) // r, p) == 1:
            t -= 1
        return pow(t, (p - 1) // r, p)

    def _all_conjugates(root: int) -> List[int]:
        g = _find_nontrivial_rth_root(3)
        return [root, (root * g) % p, (root * g * g) % p]

    # reduce, handle trivial/single‐root/no‐root
    a %= p
    if p in (2, 3) or a == 0:
        return [a]
    if p % 3 == 2:
        return [pow(a, ((2 * p) - 1) // 3, p)]
    if pow(a, (p - 1) // 3, p) > 1:
        return []

    # write p-1 = 3^s * t
    s, t = 0, p - 1
    while t % 3 == 0:
        s, t = s + 1, t // 3

    # find a cubic non-residue b
    b = p - 2
    while pow(b, (p - 1) // 3, p) == 1:
        b -= 1

    c = pow(b, t, p)
    r = pow(a, t, p)
    h = 1
    c_inv = pow(c, p - 2, p)

    # loop to adjust
    for i in range(1, s):
        d = pow(r, 3 ** (s - i - 1), p)
        if d == pow(c, 3 ** (s - 1), p):
            h = (h * c_inv) % p
            r = (r * pow(c, 3, p)) % p
        elif d != 1:
            h = (h * pow(c_inv, 2, p)) % p
            r = (r * pow(c_inv, 6, p)) % p
        c_inv = pow(c_inv, 3, p)

    # finalize r
    k = (t - 1) // 3 if (t - 1) % 3 == 0 else (t + 1) // 3
    r = (pow(a, k, p) * h) % p
    if (t - 1) % 3 == 0:
        r = pow(r, p - 2, p)

    if pow(r, 3, p) != a:
        return []
    return _all_conjugates(r) if all_roots else [r]


def modular_fourth_roots(a: int, p: int) -> Optional[List[int]]:
    """Return all x such that x⁴ ≡ a (mod p), or None if impossible."""
    first = modular_square_roots(a, p)
    if first is None:
        return None
    roots = []
    for r in first:
        second = modular_square_roots(r, p)
        if second is None:
            return None
        roots.extend(second)
    return roots


def modular_sixth_roots(a: int, p: int) -> Optional[List[int]]:
    """Return all x such that x⁶ ≡ a (mod p), or None if none exist."""
    cubes = modular_cube_roots(a, p)
    if not cubes:
        return None
    roots = []
    for c in cubes:
        sq = modular_square_roots(c, p)
        if sq is None:
            return None
        roots.extend(sq)
    return roots