# derived from: http://eli.thegreenplace.net/2009/03/07/computing-square-roots-in-python
# Compute the Legendre symbol a|p using Euler's criterion.
# p is a prime, a is relatively prime to p (if p divides a,
# then a|p = 0)
legendresymbol(a, p, r) {
pm1 = p-1;
mod(p) r = a^(pm1>>1);
if(r == pm1)
r = -1;
}
# Find a quadratic residue (mod p) of 'a'. p must be an
# odd prime.
#
# Solve the congruence of the form:
# x^2 = a (mod p)
# And returns x. Node that p - x is also a root.
#
# 0 is returned if no square root exists for these
# a and p.
#
# The Tonelli-Shanks algorithm is used (except
# for some simple cases in which the solution is known
# from an identity).
msqrt(a, p, r) {
if(legendresymbol(a, p) != 1)
r = 0;
else if(a == 0)
r = 0;
else if(p == 2)
r = a;
else if(p%4 == 3){
e = p+1 >> 2;
mod(p) r = a^e;
} else {
# Partition p-1 to s * 2^e for an odd s (i.e.
# reduce all the powers of 2 from p-1)
s = p-1;
e = 0;
while(s%2 == 0){
s = s >> 1;
e = e + 1;
}
# Find some 'n' with a legendre symbol n|p = -1.
# Shouldn't take long.
n = 2;
while(legendresymbol(n, p) != -1)
n = n + 1;
# x is a guess of the square root that gets better
# with each iteration.
# b is the "fudge factor" - by now much we're off
# with the guess. The invariant x^2 == a*b (mod p)
# is maintained throughout the loop.
# g is used for successive powers of n to update
# both a and b
# e is the exponent - decreases with each update
mod(p){
x = a^(s+1 >> 1);
b = a^s;
g = n^s;
}
while(1==1){
t = b;
m = 0;
while(m < e){
if(t == 1)
break;
t = t*t % p;
m = m + 1;
}
if(m == 0){
r = x;
break;
}
t = 2^(e-m-1);
mod(p){
gs = g^t;
g = gs*gs;
x = x*gs;
b = b*g;
}
e = m;
}
}
}
# modular inverse square-root
misqrt(a, p, r) {
if((p % 4) == 3){
e = ((p-3)>>2);
mod(p) r = a^e;
} else {
r = msqrt(a, p);
if(r != 0)
mod(p) r = 1/r;
}
}