gcd.zig

const builtin = @import("builtin");
const std = @import("std");

const assert = std.debug.assert;
const expect = std.testing.expect;
const expectEqual = std.testing.expectEqual;
const expectError = std.testing.expectError;

const Ti = i64;
const Tu = u64;
const Tu_ = u63;

fn gcd(p: anytype, q: anytype) @TypeOf(p, q) {
    assert(p >= 1);
    assert(q >= 1);

    // convert comptime_int to a sized integer within this function so that @ctz will work
    const T = switch (@TypeOf(p, q)) {
        comptime_int => std.math.IntFittingRange(@min(p, q), @max(p, q)),
        else => |U| U,
    };

    switch (@typeInfo(T)) {
        .Int => |info| {
            const unsigned = info.signedness == .unsigned;
            // If we can fit into an i64 or i128 we can use an even faster algorithm
            // https://en.algorithmica.org/hpc/algorithms/gcd/
            const n: ?comptime_int = if (info.bits <= if (unsigned) 63 else 64)
                64
            else if (info.bits <= if (unsigned) 127 else 128) 128 else null;
            if (n) |bits| {
                const iN = @Type(.{ .Int = .{ .signedness = .signed, .bits = bits } });
                const uN = @Type(.{ .Int = .{ .signedness = .unsigned, .bits = bits } });

                var u: iN = p;
                var v: iN = q;

                var uz: uN = @intCast(@ctz(u));
                const vz: uN = @intCast(@ctz(v));
                const shift = @min(uz, vz);
                v >>= @intCast(vz);

                while (true) {
                    u >>= @intCast(uz);
                    const diff = v - u;
                    if (diff == 0) break;
                    uz = @intCast(@ctz(diff));
                    v = @min(u, v);
                    u = @intCast(@abs(diff));
                }

                const result: T = @intCast(v << @intCast(shift));
                assert(result > 0);
                return result;
            } else {
                var u: T = undefined;
                var v: T = undefined;
                if (p < q) {
                    u = q;
                    v = p;
                } else {
                    u = p;
                    v = q;
                }
                assert(v <= u);
                u %= v;
                if (u == 0) return v;

                const zu = @ctz(u);
                const zv = @ctz(v);
                const shift = @min(zu, zv);
                u >>= @intCast(zu);
                v >>= @intCast(zv);

                while (true) {
                    const diff = u -% v;
                    if (u > v) {
                        u = v;
                        v = diff;
                    } else {
                        v -= u;
                    }
                    if (diff != 0) v >>= @intCast(@ctz(diff));
                    if (v == 0) break;
                }

                const result = u << @intCast(shift);
                assert(result > 0);
                return result;
            }
        },
        else => {
            var a = p;
            var b = q;
            var c: T = undefined;

            while (b != 0) {
                c = b;
                b = @mod(a, b);
                a = c;
            }

            assert(a > 0);
            return a;
        },
    }
}

fn gcd64(p: Tu, q: Tu) @TypeOf(p, q) {
    assert(p >= 1);
    assert(q >= 1);

    var u: Ti = @intCast(p);
    var v: Ti = @intCast(q);

    var uz: Tu = @intCast(@ctz(u));
    const vz: Tu = @intCast(@ctz(v));
    const shift = @min(uz, vz);
    v >>= @intCast(vz);

    while (true) {
        u >>= @intCast(uz);
        const diff = v - u;
        if (diff == 0) break;
        uz = @intCast(@ctz(diff));
        v = @min(u, v);
        u = @intCast(@abs(diff));
    }

    const result: Tu = @intCast(v << @intCast(shift));
    assert(result > 0);
    return result;
}

fn gcdLemire(p: Tu, q: Tu) @TypeOf(p, q) {
    assert(p >= 1);
    assert(q >= 1);

    var u: Tu = undefined;
    var v: Tu = undefined;
    if (p < q) {
        u = q;
        v = p;
    } else {
        u = p;
        v = q;
    }
    assert(v <= u);
    u %= v;
    if (u == 0) return v;

    const zu = @ctz(u);
    const zv = @ctz(v);
    const shift = @min(zu, zv);
    u >>= @intCast(zu);
    v >>= @intCast(zv);

    while (true) {
        const diff = u -% v;
        if (u > v) {
            u = v;
            v = diff;
        } else {
            v -= u;
        }
        if (diff != 0) v >>= @intCast(@ctz(diff));
        if (v == 0) break;
    }

    const result = u << @intCast(shift);
    assert(result > 0);
    return result;
}

fn gcdEuclid(p: anytype, q: anytype) @TypeOf(p, q) {
    assert(p >= 1);
    assert(q >= 1);

    var a = p;
    var b = q;
    var c: @TypeOf(p, q) = undefined;

    while (b != 0) {
        c = b;
        b = @mod(a, b);
        a = c;
    }

    assert(a > 0);
    return a;
}

test "gcd" {
    const seed = 0x1234578;
    var prng = std.Random.DefaultPrng.init(seed);
    var random = prng.random();

    var as: usize = 0;
    var bs: usize = 0;
    var cs: usize = 0;
    var ds: usize = 0;

    for (0..100000) |_| {
        const a = random.int(Tu_);
        const b = random.int(Tu_);
        if (a == 0 or b == 0) continue;

        var timer = try std.time.Timer.start();
        const zz = gcd(a, b);
        ds += timer.lap();
        const x = gcd64(a, b);
        as += timer.lap();
        const y = gcdLemire(a, b);
        bs += timer.lap();
        const z = gcdEuclid(a, b);
        cs += timer.lap();

        try expectEqual(x, y);
        try expectEqual(x, z);
        try expectEqual(x, zz);
    }

    std.debug.print("{d} {d} {d} {d}\n", .{ ds, as, bs, cs });
}