Use binary gcd and make overflow safe.

Author: two-horned

lib/std/math/egcd.zig

@@ -3,58 +3,117 @@ const std = @import("../std.zig");
 
 /// Result type of `egcd`.
 pub fn ExtendedGreatestCommonDivisor(S: anytype) type {
+    const N = switch (S) {
+        comptime_int => comptime_int,
+        else => |T| std.meta.Int(.unsigned, @bitSizeOf(T)),
+    };
+
     return struct {
-        gcd: S,
+        gcd: N,
         bezout_coeff_1: S,
         bezout_coeff_2: S,
     };
 }
 
+fn egcd_helper(other: anytype, odd: anytype, shift: anytype) [3]@TypeOf(other, odd) {
+    const S = @TypeOf(other, odd);
+    const toinv = @shrExact(other, @intCast(shift));
+    const ctrl = @shrExact(odd, @intCast(shift));
+
+    var s: S = std.math.sign(toinv);
+    var t: S = 0;
+
+    var x = @abs(toinv);
+    var y = @abs(ctrl);
+
+    while (x & 1 == 0) {
+        x = @shrExact(x, 1);
+        s = @shrExact(if (s & 1 == 0) s else s + ctrl, 1);
+    }
+
+    var y_minus_x = y -% x;
+    while (y_minus_x != 0) : (y_minus_x = y -% x) {
+        const t_minus_s = t - s;
+        const copy_x = x;
+        const copy_s = s;
+
+        s -= t;
+        const carry = x < y;
+        x -%= y;
+        if (carry) {
+            x = y_minus_x;
+            y = copy_x;
+            s = t_minus_s;
+            t = copy_s;
+        }
+        while (x & 1 == 0) {
+            x = @shrExact(x, 1);
+            s = @shrExact(if (s & 1 == 0) s else s + ctrl, 1);
+        }
+    }
+
+    y = @shlExact(y, @intCast(shift));
+    s = @shlExact(s, @intCast(shift));
+    // Using integer widening is only a temporary solution.
+    const W = std.meta.Int(.signed, @bitSizeOf(S) * 2);
+    t = @intCast(@divExact(y - @as(W, s) * toinv, ctrl));
+    return .{ @bitCast(y), s, t };
+}
+
 /// Returns the Extended Greatest Common Divisor (EGCD) of two signed integers (`a` and `b`) which are not both zero.
 pub fn egcd(a: anytype, b: anytype) ExtendedGreatestCommonDivisor(@TypeOf(a, b)) {
     const S = switch (@TypeOf(a, b)) {
-        // convert comptime_int to some sized int type for @ctz
         comptime_int => b: {
             const n = @max(@abs(a), @abs(b));
             break :b std.math.IntFittingRange(-n, n);
         },
         else => |T| T,
     };
-
     if (@typeInfo(S) != .int or @typeInfo(S).int.signedness != .signed) {
         @compileError("`a` and `b` must be signed integers");
     }
 
     std.debug.assert(a != 0 or b != 0);
 
-    var x: S = @intCast(@abs(a));
-    var y: S = @intCast(@abs(b));
+    if (a == 0) return .{ .gcd = @abs(b), .bezout_coeff_1 = 0, .bezout_coeff_2 = std.math.sign(b) };
+    if (b == 0) return .{ .gcd = @abs(a), .bezout_coeff_1 = std.math.sign(a), .bezout_coeff_2 = 0 };
 
-    // Mantain a = s * x + t * y.
-    var s: S = std.math.sign(a);
-    var t: S = 0;
+    const x: S = a;
+    const y: S = b;
 
-    // Mantain b = u * x + v * y.
-    var u: S = 0;
-    var v: S = std.math.sign(b);
-
-    while (x != 0) {
-        const q = @divTrunc(y, x);
-        const old_x = x;
-        const old_s = s;
-        const old_t = t;
-        x = y - q * x;
-        s = u - q * s;
-        t = v - q * t;
-        y = old_x;
-        u = old_s;
-        v = old_t;
-    }
+    const xz = @ctz(x);
+    const yz = @ctz(y);
+    const shift = @min(xz, yz);
 
-    return .{ .gcd = y, .bezout_coeff_1 = u, .bezout_coeff_2 = v };
+    if (xz < yz) {
+        const gcd, const t, const s = egcd_helper(y, x, shift);
+        return .{ .gcd = @intCast(gcd), .bezout_coeff_1 = s, .bezout_coeff_2 = t };
+    } else {
+        const gcd, const s, const t = egcd_helper(x, y, shift);
+        return .{ .gcd = @intCast(gcd), .bezout_coeff_1 = s, .bezout_coeff_2 = t };
+    }
 }
 
 test {
+    {
+        const a: i4 = -8;
+        const b: i4 = 1;
+        const r = egcd(a, b);
+        const g = r.gcd;
+        const s = r.bezout_coeff_1;
+        const t = r.bezout_coeff_2;
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a: i4 = -8;
+        const b: i4 = 5;
+        const r = egcd(a, b);
+        const g = r.gcd;
+        // Avoid overflow in assert.
+        const s: i8 = r.bezout_coeff_1;
+        const t: i8 = r.bezout_coeff_2;
+        try std.testing.expect(s * a + t * b == g);
+    }
     {
         const a: i32 = 0;
         const b: i32 = 5;