Tonelli-Shanks algorithm implementation of prime modular square root

I did an implementation of the Tonelli-Shanks algorithm as defined on Wikipedia. I put it here for review and sharing purpose.

Legendre Symbol implementation:

def legendre_symbol(a, p):
    """
    Legendre symbol
    Define if a is a quadratic residue modulo odd prime
    http://en.wikipedia.org/wiki/Legendre_symbol
    """
    ls = pow(a, (p - 1)/2, p)
    if ls == p - 1:
        return -1
    return ls

Prime modular square root (I just renamed the solution variable R to x and n to a):

def prime_mod_sqrt(a, p):
    """
    Square root modulo prime number
    Solve the equation
        x^2 = a mod p
    and return list of x solution
    http://en.wikipedia.org/wiki/Tonelli-Shanks_algorithm
    """
    a %= p

    # Simple case
    if a == 0:
        return [0]
    if p == 2:
        return [a]

    # Check solution existence on odd prime
    if legendre_symbol(a, p) != 1:
        return []

    # Simple case
    if p % 4 == 3:
        x = pow(a, (p + 1)/4, p)
        return [x, p-x]

    # Factor p-1 on the form q * 2^s (with Q odd)
    q, s = p - 1, 0
    while q % 2 == 0:
        s += 1
        q //= 2

    # Select a z which is a quadratic non resudue modulo p
    z = 1
    while legendre_symbol(z, p) != -1:
        z += 1
    c = pow(z, q, p)

    # Search for a solution
    x = pow(a, (q + 1)/2, p)
    t = pow(a, q, p)
    m = s
    while t != 1:
        # Find the lowest i such that t^(2^i) = 1
        i, e = 0, 2
        for i in xrange(1, m):
            if pow(t, e, p) == 1:
                break
            e *= 2

        # Update next value to iterate
        b = pow(c, 2**(m - i - 1), p)
        x = (x * b) % p
        t = (t * b * b) % p
        c = (b * b) % p
        m = i

    return [x, p-x]

If you have any optimization or have found any error, please report it.

Good job! I don't have a lot to comment on in this code. You have written straightforward clear code whose only complexity stems directly from the complexity of the operation it is performing. It would be good to include some of your external commentary (such as the renames of R and n) in the code itself to make it easier for someone to follow the documentation on wikipedia. You may want to include some of that documentation as well.

For reference, the rest of this review assumes that the code functions correctly; I don't have my math head on tonight.

There appears to be one case of redundant code, unless m can ever be 1, resulting in an empty range and thus no reassignment of i. Otherwise you can skip the assignment to i in the following:

i, e = 0, 2
for i in xrange(1, m):
    ...

There are a number of small strength-reduction optimizations you might consider, but in Python their impact is likely to be minimized - definitely profile before heading too deeply down the optimization path. For example in the following while loop:

# Factor p-1 on the form q * 2^s (with Q odd)
q, s = p - 1, 0
while q % 2 == 0:
    s += 1
    q //= 2

Both operations on q can be reduced. The modulus can be rewritten as a binary and q & 1, and the division as a binary shift q >>= 1. Alternately, you can use divmod to perform both operations at once.

Similarly, 2**(m - i - 1) is identical to 1 << (m - i - 1) for non-negative exponents.

StackExchange Code Review Q#43210, answer score: 8