#### Square Roots Modulo a Prime

This code requires the Modular Arithmetic package.

```/* sqrt_prime(A, P, R) is true if P is a prime number, A is a quadratic      */
/*   residue modulo P, and +R and -R are the square roots of A modulo P.     */
sqrt_prime(A, 2, R):-!,
R is A mod 2.
sqrt_prime(A, P, 0):-
A mod P =:= 0, !.
sqrt_prime(A, P, R):-
P mod 4 =:= 3, !,
sqrt_3_4(A, P, R).
sqrt_prime(A, P, R):-
P mod 8 =:= 5, !,
sqrt_5_8(A, P, R).
sqrt_prime(A, P, R):-
P mod 8 =:= 1,
sqrt_1_8(A, P, R).

/* sqrt_3_4(A, P, R) is true if P is a prime number congruent to 3 modulo 4, */
/*   A is a quadratic residue modulo P, and +R and -R are the square roots   */
/*   of A modulo P.                                                          */
/* e.g. sqrt_3_4(9, 32719, X) gives X=32716                                  */
sqrt_3_4(A0, P, R):-
P mod 4 =:= 3,
is_prime(P),
A is A0 mod P,
P_1_4 is (P + 1) // 4,
powmod(A, P_1_4, P, R), /* r=a^((p+1)/4) mod p                 */
multmod(R, R, P, A).    /* Ensure that r is a square root of a */

/* sqrt_5_8(A, P, R) is true if P is a prime number congruent to 5 modulo 8, */
/*   A is a quadratic residue modulo P, and +R and -R are the square roots   */
/*   of A modulo P.                                                          */
/* e.g. sqrt_5_8(9, 32717, X) gives X=32714                                  */
sqrt_5_8(A0, P, R):-
P mod 8 =:= 5,
is_prime(P),
A is A0 mod P,
P_5_8 is (P - 5) // 8,
addmod(A, A, P, A2),     /* Atkin's algorithm:                  */
powmod(A2, P_5_8, P, V), /* v=(2a)^((p-5)/8) mod p              */
multmod(A2, V, P, I0),
multmod(I0, V, P, I),    /* i=(2av^2) mod p                     */
subtmod(I, 1, P, I_1),
multmod(A, V, P, AV),
multmod(AV, I_1, P, R),  /* r=(av(i-1)) mod p                   */
multmod(R, R, P, A).     /* Ensure that r is a square root of a */

/* sqrt_1_8(A, P, R) is true if P is a prime number congruent to 1 modulo 8, */
/*   A is a quadratic residue modulo P, and +R and -R are the square roots   */
/*   of A modulo P.                                                          */
/* Shank's algorithm [Knuth Exercise 4.6.2-15]: Find e and q such that       */
/*   p=(2^e)q+1 and e>2. Choose x at random in the range 1 < x < p, and set  */
/*   z=x^q mod p. If z^(2^(e-1)) mod p=1, repeat this step with a different  */
/*   value of x. Otherwise initialize y=z, s=e, x=a^((q-1)/2) mod p,         */
/*   r=ax mod p and w=ax^2 mod p, and repeat the following until w=1, when   */
/*   r is the answer:                                                        */
/*     Find the smallest k such that w^(2^k) mod p=1. If k=s there is no     */
/*     solution. Otherwise set y=y^(2^(s-k)), s=k, r=ry^(2^(s-k-1)) and      */
/*     w=wy^(2^(s-k)).                                                       */
/* e.g. sqrt_1_8(9, 32713, X) gives X=32710                                  */
sqrt_1_8(A0, P, R):-
P mod 8 =:= 1,
is_prime(P),
A is A0 mod P,
P1 is P - 1,
sqrt_1_8_e_q(P1, Q, 0, E),   /* p=(2^e)q+1                                */
sqrt_1_8_z(2, P, E, Q, Z),   /* z=x^q mod p such that z^(2^(e-1)) mod p=1 */
Q1 is (Q - 1) // 2,
powmod(A, Q1, P, X),         /* x=a^((q-1)/2) mod p                       */
multmod(A, X, P, R0),        /* r=ax mod p                                */
multmod(R0, X, P, W),        /* w=ax^2 mod p                              */
sqrt_1_8_iterate(W, P, X, Z, E, R0, R).

sqrt_1_8_iterate(1, _, _, _, _, R, R):-!.
sqrt_1_8_iterate(W, P, X, Y, S, R0, R):-
sqrt_1_8_k(1, S, W, P, K),   /* Smallest k < s such that w^(2^k) mod p=1 */
S_K is S - K,
S_K_1 is S - K - 1,
powmod(2, S_K, P, S_K_2),
powmod(2, S_K_1, P, S_K_1_2),
powmod(Y, S_K_2, P, Y1),     /* y1=y^(2^(s-k))    */
S1 = K,                      /* s1=k              */
powmod(Y, S_K_1_2, P, T),
multmod(R0, T, P, R1),       /* r1=ry^(2^(s-k-1)) */
multmod(W, Y1, P, W1),       /* w1=wy^(2^(s-k))   */
sqrt_1_8_iterate(W1, P, X, Y1, S1, R1, R).

/* sqrt_1_8_k(1, S, W, P, K) is true if k is the smallest integer less than  */
/*   s such that w^(2^k) mod p=1.                                            */
sqrt_1_8_k(K, S, W, P, K):-
K < S,
powmod(2, K, P, T),
powmod(W, T, P, 1), !.
sqrt_1_8_k(I, S, W, P, K):-
I < S,
I1 is I + 1,
sqrt_1_8_k(I1, S, W, P, K).

/* sqrt_1_8_e_q(Q0, Q, 0, E) is true if q0 = (2^e)q.                         */
/* e.g. sqrt_1_8_e_q(88, 11, 0, 3).                                          */
sqrt_1_8_e_q(Q, Q, E, E):-
Q mod 2 =:= 1, !.
sqrt_1_8_e_q(Q0, Q, E0, E):-
Q1 is Q0 // 2,
E1 is E0 + 1,
sqrt_1_8_e_q(Q1, Q, E1, E).

/* sqrt_1_8_z(2, P, E, Q, Z) is true if x is the smallest integer [it should */
/*   be a random integer] such that z=x^q mod p and z^(2^(e-1)) mod p != 1.  */
sqrt_1_8_z(X, P, E, Q, Z):-
powmod(X, Q, P, Z),
E_1 is E - 1,
powmod(2, E_1, P, T),
powmod(Z, T, P, S),
S =\= 1, !.
sqrt_1_8_z(X, P, E, Q, Z):-
X1 is X + 1,
sqrt_1_8_z(X1, P, E, Q, Z).

/* is_prime(P) is true if P is a prime number.                               */
is_prime(2):-!.
is_prime(P):-
P > 2,
P mod 2 =:= 1,
SqrtP is sqrt(P),
is_prime_1(3, P, SqrtP).

is_prime_1(I, _, SqrtP):-
I > SqrtP, !.
is_prime_1(I, P, SqrtP):-
P mod I > 0,
I1 is I + 2,
is_prime_1(I1, P, SqrtP).

/* sqrt_prime_squared(A, P, R1, R2) is true if P is an odd prime number, A   */
/*   is a quadratic residue modulo P, and +/-R1 and +/-R2 are the square     */
/*   roots of A modulo P^2.                                                  */
/* e.g. sqrt_prime_squared(39, 179, R1, R2) gives R1=10000, R2=22041         */
/*   i.e. (10000*10000) mod (179*179) =:= 39                                 */
/*    and (22041*22041) mod (179*179) =:= 39                                 */
sqrt_prime_squared(A, P, R1, R2):-
P =< 32767,
A > 0,
P_2 is P * P,
A < P_2,
AmodP is A mod P,
sqrt_prime(AmodP, P, H1),
TwoH1 is (H1 + H1) mod P,
inversemod(TwoH1, P, Inv),
H12 is H1 * H1,
sqrt_prime_squared_1(A, H12, P, Temp),
H2P is ((Inv * (Temp mod P)) mod P) * P,
addmod(H1, H2P, P_2, R1),
R2 is P_2 - R1,
multmod(R1, R1, P_2, A), % Sanity check
multmod(R2, R2, P_2, A). % Sanity check

sqrt_prime_squared_1(A, H12, P, H2P):-
A >= H12, !,
H2P is (A - H12) // P.
sqrt_prime_squared_1(A, H12, P, H2P):-
H2P is (P + (A - H12) // P).

/* sqrt_prime_power(N, P, K, X) is true if P is an odd prime number, K is a  */
/*   positive integer, N is a quadratic residue modulo P, and X is a square  */
/*   root of N modulo P^K.                                                   */
/* e.g. sqrt_prime_power(7, 3, 9, X) gives X=11110                           */
/*   that is, (11110*11110) mod (3^9) = 7                                    */
sqrt_prime_power(N, P, K, X):-
K > 0,
NmodP is N mod P,
sqrt_prime(NmodP, P, T),
sqrt_prime_power_1(K, N, P, P, T, X),
% Avoid incorrect results due to integer overflow or rounding errors:
PK is int(P^K + 0.5),
multmod(X, X, PK, M),
M =:= N mod PK.

sqrt_prime_power_1(1, _, _, _, T, T):-!.
sqrt_prime_power_1(I, N, P, Pk, T, X):-
V0 is (N - T * T) // Pk,
sqrt_prime_power_2(V0, Pk, V),
addmod(T, T, P, U),
W is V mod P,
divmod(W, U, P, C),
I1 is I - 1,
Pk1 is Pk * P,
T1 is T + C * Pk,
sqrt_prime_power_1(I1, N, P, Pk1, T1, X).

sqrt_prime_power_2(S, Pk, V):-
S < 0, !,
V is Pk + S.
sqrt_prime_power_2(S, _, S).

/*
* Sample application: Cornacchia's Algorithm
*/

% cornacchia(D, N, X, Y) is true if (X,Y) is the solution of x^2+Dy^2=N where
%   0 < D < N, and N is prime.
% e.g. cornacchia(33, 30529, X, Y) gives X=152, Y=15
cornacchia(D, N, X, Y):-
0 < D, D < N,
is_prime(N),
N_D is N - D,
sqrt_prime(N_D, N, R),
cornacchia_1(N, R, N, X),
N_XX is N - X*X,
0 =:= N_XX mod D,
NXXD is N_XX // D,
Y is int(sqrt(NXXD)),
Y*Y =:= NXXD,
X*X + D*Y*Y =:= N. % Sanity check

% e.g. cornacchia_1(30529, 16272, 30529, X) gives X=152
cornacchia_1(_, R, N, X):-
R*R =< N, !,
X = R.
cornacchia_1(R0, R1, N, X):-
R2 is R0 mod R1,
cornacchia_1(R1, R2, N, X).
```