modular_square_root_all_solutions.sf

#!/usr/bin/ruby

# Find all solutions to the quadratic congruence:
#   x^2 = a (mod n)

# Based on algorithm by Hugo van der Sanden:
#   https://github.com/danaj/Math-Prime-Util/pull/55

func sqrtmod_all(a, n) {

    n = -n if (n < 0)

    n == 0 && return []
    n == 1 && return [0]

    a = (a % n)

    func sqrtmod_pk(a, p, k) {
        var pk = p**k

        if (p `divides` a) {

            if (pk `divides` a) {
                var low = p**(k >> 1)
                var high = (k.is_odd ? (low*p) : low)
                return (^low -> map {|i| high * i })
            }

            var a2 = a/p
            return [] if !(p `divides` a2)
            var pj = pk/p
            return __FUNC__(a2/p, p, k-2).map {|q|
                ^p -> map {|i| q*p + i*pj }...
            }
        }

        var q = sqrtmod(a, pk)

        q.is_nan && return []

        return [q, pk-q] if (p != 2)
        return [q]       if (k == 1)
        return [q, pk-q] if (k == 2)

        var pj = p**(k-1)
        var q2 = (q * (pj-1))%pk

        return [q, pk-q, q2, pk-q2]
    }

    var roots = []

    n.factor_map {|p,k|
        sqrtmod_pk(a, p, k).map {|r| [r, p**k] }
    }.cartesian {|*a|
        roots << Math.chinese(a...)
    }

    return roots.sort
}

say sqrtmod_all(1, 8)           #=> [1, 3, 5, 7]
say sqrtmod_all(120, 5045)      #=> [1165, 3880]
say sqrtmod_all(4095, 8469)     #=> [1110, 1713, 3933, 4536, 6756, 7359]

# Run some tests

assert_eq(
    sqrtmod_all(-1, 13**18 * 5**7)
    %n(633398078861605286438568 2308322911594648160422943 6477255756527023177780182 8152180589260066051764557)
)

assert_eq(
    sqrtmod_all(2466, 5967),
    [120 237 426 543 1446 1563 1752 1869 2109 2226 2415 2532 3435 3552 3741 3858 4098 4215 4404 4521 5424 5541 5730 5847]
)

assert_eq(
    sqrtmod_all(7281, 9954),
    %n(1233 1611 1707 2085 4551 4929 5025 5403 7869 8247 8343 8721)
)
assert_eq(
    sqrtmod_all(1701, 6300),
    %n[399, 651, 1449, 1701, 2499, 2751, 3549, 3801, 4599, 4851, 5649, 5901]
)

assert_eq(
    sqrtmod_all(306, 810),
    %n[66, 96, 174, 204, 336, 366, 444, 474, 606, 636, 714, 744]
)

assert_eq(
    sqrtmod_all(2754, 6561),
    %n[126, 603, 855, 1332, 1584, 2061, 2313, 2790, 3042, 3519, 3771, 4248, 4500, 4977, 5229, 5706, 5958, 6435]
)

assert_eq(
    sqrtmod_all(17640, 48465),
    %n[2865, 7905, 8250, 13290, 19020, 24060, 24405, 29445, 35175, 40215, 40560, 45600]
)

for n in (0..20), a in (^n) {

    var roots = sqrtmod_all(a, n)

    assert(roots.all {|r|
        r**2 % n == a
    }, "sqrtmod(#{a}, #{n}) = #{roots}")

    assert_eq(
        ^n -> grep {|k| mulmod(k, k, n) == a },
        roots
    )
}