Add EGCD.

Fix some comments in GCD.

Make ml_kem use lcm and egcd from std/math.

Fix name.

Add egcd function.

Don't destructure.

Use binary gcd and make overflow safe.

Force inlining, use ctz to reduce dependency in loop.

Avoid integer overflow for temporary value.

Add test against previous overflow capability.

More optimization friendly expression.

Fix egcd for even numbers.

Minvalue causes crash.

Remove helper function. Fix casting issues.

Use shift instead division (to support i2) and avoid overflow of temp results.

Author: two-horned

lib/std/crypto/ml_kem.zig

@@ -636,7 +636,7 @@ test "invNTTReductions bounds" {
 fn invertMod(a: anytype, p: @TypeOf(a)) @TypeOf(a) {
     const r = extendedEuclidean(@TypeOf(a), a, p);
     assert(r.gcd == 1);
-    return r.x;
+    return r.bezout_coeff_1;
 }
 
 // Reduce mod q for testing.

lib/std/math.zig

@@ -238,6 +238,7 @@ pub const sinh = @import("math/sinh.zig").sinh;
 pub const cosh = @import("math/cosh.zig").cosh;
 pub const tanh = @import("math/tanh.zig").tanh;
 pub const gcd = @import("math/gcd.zig").gcd;
+pub const egcd = @import("math/egcd.zig").egcd;
 pub const lcm = @import("math/lcm.zig").lcm;
 pub const gamma = @import("math/gamma.zig").gamma;
 pub const lgamma = @import("math/gamma.zig").lgamma;

lib/std/math/egcd.zig

@@ -0,0 +1,241 @@
+//! Extended Greatest Common Divisor (https://mathworld.wolfram.com/ExtendedGreatestCommonDivisor.html)
+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: N,
+        bezout_coeff_1: S,
+        bezout_coeff_2: S,
+    };
+}
+
+/// 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)) {
+        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);
+
+    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 };
+
+    const other: S, const odd: S, const shift, const switch_coeff = b: {
+        const xz = @ctz(@as(S, a));
+        const yz = @ctz(@as(S, b));
+        break :b if (xz < yz) .{ b, a, xz, true } else .{ a, b, yz, false };
+    };
+    const toinv = @shrExact(other, @intCast(shift));
+    const ctrl = @shrExact(odd, @intCast(shift)); // Invariant: |s|, |t|, |ctrl| < |MIN_OF(S)|
+    const half_ctrl = 1 + @shrExact(ctrl - 1, 1);
+    const abs_ctrl = @abs(ctrl);
+
+    var s: S = std.math.sign(toinv);
+    var t: S = 0;
+
+    var x = @abs(toinv);
+    var y = abs_ctrl;
+
+    {
+        const xz = @ctz(x);
+        x = @shrExact(x, @intCast(xz));
+        for (0..xz) |_| {
+            const half_s = s >> 1;
+            if (s & 1 == 0)
+                s = half_s
+            else
+                s = half_s + half_ctrl;
+        }
+    }
+
+    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;
+        const xz = @ctz(y_minus_x);
+
+        s -= t;
+        const carry = x < y;
+        x -%= y;
+        if (carry) {
+            x = y_minus_x;
+            y = copy_x;
+            s = t_minus_s;
+            t = copy_s;
+        }
+        x = @shrExact(x, @intCast(xz));
+        for (0..xz) |_| {
+            const half_s = s >> 1;
+            if (s & 1 == 0)
+                s = half_s
+            else
+                s = half_s + half_ctrl;
+        }
+
+        if (s < 0) s = @intCast(abs_ctrl - @abs(s));
+    }
+
+    // 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));
+    const final_s, const final_t = if (switch_coeff) .{ t, s } else .{ s, t };
+    return .{
+        .gcd = @shlExact(y, @intCast(shift)),
+        .bezout_coeff_1 = final_s,
+        .bezout_coeff_2 = final_t,
+    };
+}
+
+test {
+    {
+        const a: i2 = 0;
+        const b: i2 = 1;
+        const r = egcd(a, b);
+        const g = r.gcd;
+        const s: i2 = r.bezout_coeff_1;
+        const t: i2 = r.bezout_coeff_2;
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a: i8 = -128;
+        const b: i8 = 127;
+        const r = egcd(a, b);
+        const g = r.gcd;
+        const s: i16 = r.bezout_coeff_1;
+        const t: i16 = r.bezout_coeff_2;
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a: i16 = -32768;
+        const b: i16 = -32768;
+        const r = egcd(a, b);
+        const g = r.gcd;
+        const s: i32 = r.bezout_coeff_1;
+        const t: i32 = r.bezout_coeff_2;
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a: i32 = 128;
+        const b: i32 = 112;
+        const r = egcd(a, b);
+        const g = r.gcd;
+        const s: i64 = r.bezout_coeff_1;
+        const t: i64 = r.bezout_coeff_2;
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a: i32 = 4 * 89;
+        const b: i32 = 2 * 17;
+        const r = egcd(a, b);
+        const g = r.gcd;
+        const s: i64 = r.bezout_coeff_1;
+        const t: i64 = r.bezout_coeff_2;
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a: i8 = 127;
+        const b: i8 = 126;
+        const r = egcd(a, b);
+        const g = r.gcd;
+        const s: i16 = r.bezout_coeff_1;
+        const t: i16 = r.bezout_coeff_2;
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        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;
+        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: i32 = 5;
+        const b: i32 = 0;
+        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: i32 = 21;
+        const b: i32 = 15;
+        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: i32 = -21;
+        const b: i32 = 15;
+        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 = -21;
+        const b = 15;
+        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 = 927372692193078999176;
+        const b = 573147844013817084101;
+        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 = 453973694165307953197296969697410619233826;
+        const b = 280571172992510140037611932413038677189525;
+        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);
+    }
+}

lib/std/math/gcd.zig

@@ -1,7 +1,7 @@
-//! Greatest common divisor (https://mathworld.wolfram.com/GreatestCommonDivisor.html)
-const std = @import("std");
+//! Greatest Common Divisor (https://mathworld.wolfram.com/GreatestCommonDivisor.html)
+const std = @import("../std.zig");
 
-/// Returns the greatest common divisor (GCD) of two unsigned integers (`a` and `b`) which are not both zero.
+/// Returns the Greatest Common Divisor (GCD) of two unsigned integers (`a` and `b`) which are not both zero.
 /// For example, the GCD of `8` and `12` is `4`, that is, `gcd(8, 12) == 4`.
 pub fn gcd(a: anytype, b: anytype) @TypeOf(a, b) {
     const N = switch (@TypeOf(a, b)) {