Add EGCD.

Fix some comments in GCD.

Make ml_kem use lcm and egcd from std/math.

Author: two-horned

lib/std/crypto/ml_kem.zig

@@ -634,33 +634,11 @@ test "invNTTReductions bounds" {
     }
 }
 
-// Extended euclidean algorithm.
-//
-// For a, b finds x, y such that  x a + y b = gcd(a, b). Used to compute
-// modular inverse.
-fn eea(a: anytype, b: @TypeOf(a)) EeaResult(@TypeOf(a)) {
-    if (a == 0) {
-        return .{ .gcd = b, .x = 0, .y = 1 };
-    }
-    const r = eea(@rem(b, a), a);
-    return .{ .gcd = r.gcd, .x = r.y - @divTrunc(b, a) * r.x, .y = r.x };
-}
-
-fn EeaResult(comptime T: type) type {
-    return struct { gcd: T, x: T, y: T };
-}
-
-// Returns least common multiple of a and b.
-fn lcm(a: anytype, b: @TypeOf(a)) @TypeOf(a) {
-    const r = eea(a, b);
-    return a * b / r.gcd;
-}
-
 // Invert modulo p.
 fn invertMod(a: anytype, p: @TypeOf(a)) @TypeOf(a) {
-    const r = eea(a, p);
-    assert(r.gcd == 1);
-    return r.x;
+    const gcd, const x, _ = std.math.egcd(a, p);
+    assert(gcd == 1);
+    return x;
 }
 
 // Reduce mod q for testing.
@@ -1054,7 +1032,7 @@ const Poly = struct {
         var in_off: usize = 0;
         var out_off: usize = 0;
 
-        const batch_size: usize = comptime lcm(@as(i16, d), 8);
+        const batch_size: usize = comptime std.math.lcm(@as(i16, d), 8);
         const in_batch_size: usize = comptime batch_size / d;
         const out_batch_size: usize = comptime batch_size / 8;
 
@@ -1118,7 +1096,7 @@ const Poly = struct {
         var in_off: usize = 0;
         var out_off: usize = 0;
 
-        const batch_size: usize = comptime lcm(@as(i16, d), 8);
+        const batch_size: usize = comptime std.math.lcm(@as(i16, d), 8);
         const in_batch_size: usize = comptime batch_size / 8;
         const out_batch_size: usize = comptime batch_size / d;
 

lib/std/math/egcd.zig

@@ -0,0 +1,112 @@
+//! Extended Greatest Common Divisor (https://mathworld.wolfram.com/ExtendedGreatestCommonDivisor.html)
+const std = @import("../std.zig");
+
+/// Result type of `egcd`.
+pub fn ExtendedCommonDivisor(S: anytype) type {
+    if (@typeInfo(S) != .int or @typeInfo(S).int.signedness != .signed) {
+        @compileError("`S` must be a signed integer.");
+    }
+
+    return struct {
+        gcd: S,
+        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) ExtendedCommonDivisor(@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));
+
+    // Mantain a = s * x + t * y.
+    var s: S = std.math.sign(a);
+    var t: S = 0;
+
+    // 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;
+    }
+
+    return .{ .gcd = y, .bezout_coeff_1 = u, .bezout_coeff_2 = v };
+}
+
+test {
+    {
+        const a: i32 = 0;
+        const b: i32 = 5;
+        const s, const t = egcd(a, b);
+        const g = std.math.gcd(@abs(a), @abs(b));
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a: i32 = 5;
+        const b: i32 = 0;
+        const s, const t = egcd(a, b);
+        const g = std.math.gcd(@as(u32, @intCast(a)), @as(u32, @intCast(b)));
+        try std.testing.expect(s * a + t * b == g);
+    }
+
+    {
+        const a: i32 = 21;
+        const b: i32 = 15;
+        const s, const t = egcd(a, b);
+        const g = std.math.gcd(@abs(a), @abs(b));
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a: i32 = -21;
+        const b: i32 = 15;
+        const s, const t = egcd(a, b);
+        const g = std.math.gcd(@abs(a), @abs(b));
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a = -21;
+        const b = 15;
+        const s, const t = egcd(a, b);
+        const g = std.math.gcd(@abs(a), @abs(b));
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a = 927372692193078999176;
+        const b = 573147844013817084101;
+        const s, const t = egcd(a, b);
+        const g = std.math.gcd(@abs(a), @abs(b));
+        try std.testing.expect(s * a + t * b == g);
+    }
+    {
+        const a = 453973694165307953197296969697410619233826;
+        const b = 280571172992510140037611932413038677189525;
+        const s, const t = egcd(a, b);
+        const g = std.math.gcd(@abs(a), @abs(b));
+        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)) {