sidef-scripts/Math/square_root_modulo_n_tonelli-shanks.sf

#!/usr/bin/ruby

# Daniel "Trizen" Șuteu
# Date: 28 October 2017
# https://github.com/trizen

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

# Defined for any values of `a` and `n` for which `kronecker(a, n) = 1`.

# When `kronecker(a, n) != 1`, for example:
#
#   a = 472
#   n = 972
#
# which represents:
#   x^2 = 472 (mod 972)
#
# this algorithm is not able to find a solution, although there exist four solutions:
#   x = {38, 448, 524, 934}

# Code inspired from:
#   https://github.com/Magtheridon96/Square-Root-Modulo-N

func tonelli(n, p) {

    var q = p-1
    var s = valuation(q, 2)

    s == 1 && return powmod(n, (p + 1) >> 2, p)

    q >>= s

    var c = powmod(2 ..^ p -> first {|z| legendre(z, p) == -1 }, q, p)
    var r = powmod(n, (q + 1) >> 1, p)
    var t = powmod(n, q, p)

    while (!p.divides(t - 1)) {

        var b = 1
        var t2 = (t*t % p)

        for i in (1 ..^ s) {
            if (p.divides(t2 - 1)) {
                b = powmod(c, 1 << (s - i - 1), p)
                s = i
                break
            }
            t2 = (t2*t2 % p)
        }

        r = (r*b % p)
        c = (b*b % p)
        t = (t*c % p)
    }

    return r
}

func sqrt_mod_n(a, n) is cached {

    kronecker(a, n) == 1 || return []

    a = (a % n)

    if ((n & (n - 1)) == 0) {       # n is a power of 2

        if (a % 8 == 1) {

            var k = n.valuation(2)

            k == 1 && return [1]
            k == 2 && return [1, 3]
            k == 3 && return [1, 3, 5, 7]

            if (a == 1) {
                return [1, (n>>1) - 1, (n>>1) + 1, n - 1]
            }

            return gather {
                for s in (sqrt_mod_n(a, n >> 1)) {
                    var i = (((s*s - a) >> (k - 1)) % 2)
                    var r = (s + (i << (k - 2)))
                    take(r, n - r)
                }
            }.uniq.sort
        }

        return []
    }

    if (n.is_prime) {               # n is a prime
        return gather {
            var r = tonelli(a, n)
            take(r, n - r)
        }.sort
    }

    var pe = n.factor_exp           # factorize `n` into prime powers

    if (pe.len == 1) {              # `n` is an odd prime power

        var p = pe[0][0]
        var k = pe[0][1]

        kronecker(a, p) == 1 || return []

        var roots = gather {
            var r = tonelli(a, p)
            take(r, n - r)
        }

        var pk = p
        var pi = p*p

        (k-1).times {
            var x = roots[0]
            var y = (invmod(2, pk) * invmod(x, pk))

            roots[0] = ((pi + x - y*(x*x - a + pi)) % pi)
            roots[1] = (pi - roots[0])

            pk *= p
            pi *= p
        }

        return roots.sort
    }

    var solutions = []

    for p,e in (pe) {
        var m = p**e
        var r = sqrt_mod_n(a, m)
        solutions.append(r.map {|r0| [r0, m] })
    }

    gather {
        solutions.cartesian {|*a|
            take(Math.chinese(a...))
        }
    }.uniq.sort
}

for n in (1..1000) {

    var a = irand(2, n)
    var solutions = sqrt_mod_n(a, n) || next

    say "x^2 = #{a} (mod #{n}); x = {#{solutions.join(', ')}}"
}

__END__
x^2 = 455 (mod 838); x = {413, 425}
x^2 = 425 (mod 842); x = {419, 423}
x^2 =  97 (mod 849); x = {83, 200, 649, 766}
x^2 = 202 (mod 859); x = {283, 576}
x^2 = 551 (mod 862); x = {219, 643}
x^2 = 742 (mod 869); x = {128, 425, 444, 741}
x^2 = 628 (mod 871); x = {206, 340, 531, 665}
x^2 = 607 (mod 877); x = {153, 724}
x^2 = 126 (mod 887); x = {112, 775}
x^2 = 841 (mod 905); x = {29, 391, 514, 876}
x^2 = 310 (mod 911); x = {76, 835}
x^2 =  83 (mod 919); x = {177, 742}
x^2 = 441 (mod 920); x = {21, 71, 159, 209, 251, 301, 389, 439, 481, 531, 619, 669, 711, 761, 849, 899}
x^2 = 576 (mod 923); x = {24, 379, 544, 899}
x^2 =  99 (mod 929); x = {429, 500}
x^2 = 884 (mod 937); x = {126, 811}
x^2 =  88 (mod 941); x = {111, 830}
x^2 = 875 (mod 949); x = {119, 392, 557, 830}
x^2 = 836 (mod 953); x = {115, 838}
x^2 = 415 (mod 959); x = {276, 409, 550, 683}
x^2 = 450 (mod 961); x = {779, 182}
x^2 = 289 (mod 974); x = {17, 957}
x^2 = 821 (mod 977); x = {306, 671}
x^2 = 439 (mod 983); x = {173, 810}
x^2 = 574 (mod 991); x = {430, 561}
x^2 = 289 (mod 992); x = {17, 79, 417, 479, 513, 575, 913, 975}
x^2 = 464 (mod 995); x = {128, 327, 668, 867}