129 lines
3.2 KiB
Zig
129 lines
3.2 KiB
Zig
const std = @import("std");
|
|
const stdout = std.io.getStdOut().writer();
|
|
|
|
const Pair = struct {
|
|
n: u64,
|
|
p: u64,
|
|
};
|
|
|
|
const Solution = struct {
|
|
root1: u64,
|
|
root2: u64,
|
|
is_square: bool,
|
|
};
|
|
|
|
fn multiplyModulus(a: u64, b: u64, modulus: u64) u64 {
|
|
var a_mod = a % modulus;
|
|
var b_mod = b % modulus;
|
|
|
|
if (b_mod < a_mod) {
|
|
const temp = a_mod;
|
|
a_mod = b_mod;
|
|
b_mod = temp;
|
|
}
|
|
|
|
var result: u64 = 0;
|
|
var a_temp = a_mod;
|
|
var b_temp = b_mod;
|
|
|
|
while (a_temp > 0) {
|
|
if (a_temp % 2 == 1) {
|
|
result = (result + b_temp) % modulus;
|
|
}
|
|
b_temp = (b_temp << 1) % modulus;
|
|
a_temp >>= 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
fn powerModulus(base: u64, exponent: u64, modulus: u64) u64 {
|
|
if (modulus == 1) {
|
|
return 0;
|
|
}
|
|
|
|
var base_mod = base % modulus;
|
|
var result: u64 = 1;
|
|
var exp = exponent;
|
|
|
|
while (exp > 0) {
|
|
if ((exp & 1) == 1) {
|
|
result = multiplyModulus(result, base_mod, modulus);
|
|
}
|
|
base_mod = multiplyModulus(base_mod, base_mod, modulus);
|
|
exp >>= 1;
|
|
}
|
|
return result;
|
|
}
|
|
|
|
fn legendre(a: u64, p: u64) u64 {
|
|
return powerModulus(a, (p - 1) / 2, p);
|
|
}
|
|
|
|
fn tonelliShanks(n: u64, p: u64) Solution {
|
|
if (legendre(n, p) != 1) {
|
|
return Solution{ .root1 = 0, .root2 = 0, .is_square = false };
|
|
}
|
|
|
|
// Factor out powers of 2 from p - 1
|
|
var q = p - 1;
|
|
var s: u64 = 0;
|
|
while (q % 2 == 0) {
|
|
q /= 2;
|
|
s += 1;
|
|
}
|
|
|
|
if (s == 1) {
|
|
const result = powerModulus(n, (p + 1) / 4, p);
|
|
return Solution{ .root1 = result, .root2 = p - result, .is_square = true };
|
|
}
|
|
|
|
// Find a non-square z such as ( z | p ) = -1
|
|
var z: u64 = 2;
|
|
while (legendre(z, p) != p - 1) {
|
|
z += 1;
|
|
}
|
|
|
|
var c = powerModulus(z, q, p);
|
|
var t = powerModulus(n, q, p);
|
|
var m = s;
|
|
var result = powerModulus(n, (q + 1) >> 1, p);
|
|
|
|
while (t != 1) {
|
|
var i: u64 = 1;
|
|
var z_temp = multiplyModulus(t, t, p);
|
|
while (z_temp != 1 and i < m - 1) {
|
|
i += 1;
|
|
z_temp = multiplyModulus(z_temp, z_temp, p);
|
|
}
|
|
const b = powerModulus(c, @as(u64, 1) << @as( u6, @intCast(m - i - 1)), p);
|
|
c = multiplyModulus(b, b, p);
|
|
t = multiplyModulus(t, c, p);
|
|
m = i;
|
|
result = multiplyModulus(result, b, p);
|
|
}
|
|
|
|
return Solution{ .root1 = result, .root2 = p - result, .is_square = true };
|
|
}
|
|
|
|
pub fn main() !void {
|
|
const tests = [_]Pair{
|
|
.{ .n = 10, .p = 13 },
|
|
.{ .n = 56, .p = 101 },
|
|
.{ .n = 1030, .p = 1009 },
|
|
.{ .n = 1032, .p = 1009 },
|
|
.{ .n = 44402, .p = 100049 },
|
|
.{ .n = 665820697, .p = 1000000009 },
|
|
.{ .n = 881398088036, .p = 1000000000039 },
|
|
};
|
|
|
|
for (tests) |my_test| {
|
|
const solution = tonelliShanks(my_test.n, my_test.p);
|
|
try stdout.print("n = {}, p = {}", .{ my_test.n, my_test.p });
|
|
if (solution.is_square) {
|
|
try stdout.print(" has solutions: {} and {}\n\n", .{ solution.root1, solution.root2 });
|
|
} else {
|
|
try stdout.print(" has no solutions because n is not a square modulo p\n\n", .{});
|
|
}
|
|
}
|
|
}
|