From b2d60f00fabfc569fab1a858f2d7bf373c20eda7 Mon Sep 17 00:00:00 2001
From: Fri3dNstuff <102751849+Fri3dNstuff@users.noreply.github.com>
Date: Tue, 24 Sep 2024 03:15:57 +0300
Subject: [PATCH] std.math: change gcd's implementation to use Stein's
 algorithm instead of Euclid's (#21077)

---
 lib/std/math/gcd.zig | 59 ++++++++++++++++++++++++++------------------
 1 file changed, 35 insertions(+), 24 deletions(-)

diff --git a/lib/std/math/gcd.zig b/lib/std/math/gcd.zig
index 36ba8e3614ce..16ca7846f19a 100644
--- a/lib/std/math/gcd.zig
+++ b/lib/std/math/gcd.zig
@@ -1,41 +1,50 @@
 //! Greatest common divisor (https://mathworld.wolfram.com/GreatestCommonDivisor.html)
 const std = @import("std");
-const expectEqual = std.testing.expectEqual;
 
-/// 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.
+/// 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) {
-
-    // only unsigned integers are allowed and not both must be zero
-    comptime switch (@typeInfo(@TypeOf(a, b))) {
-        .int => |int| std.debug.assert(int.signedness == .unsigned),
-        .comptime_int => {
-            std.debug.assert(a >= 0);
-            std.debug.assert(b >= 0);
-        },
-        else => unreachable,
+    const N = switch (@TypeOf(a, b)) {
+        // convert comptime_int to some sized int type for @ctz
+        comptime_int => std.math.IntFittingRange(@min(a, b), @max(a, b)),
+        else => |T| T,
     };
+    if (@typeInfo(N) != .int or @typeInfo(N).int.signedness != .unsigned) {
+        @compileError("`a` and `b` must be usigned integers");
+    }
+
+    // using an optimised form of Stein's algorithm:
+    // https://en.wikipedia.org/wiki/Binary_GCD_algorithm
     std.debug.assert(a != 0 or b != 0);
 
-    // if one of them is zero, the other is returned
     if (a == 0) return b;
     if (b == 0) return a;
 
-    // init vars
-    var x: @TypeOf(a, b) = a;
-    var y: @TypeOf(a, b) = b;
-    var m: @TypeOf(a, b) = a;
+    var x: N = a;
+    var y: N = b;
 
-    // using the Euclidean algorithm (https://mathworld.wolfram.com/EuclideanAlgorithm.html)
-    while (y != 0) {
-        m = x % y;
-        x = y;
-        y = m;
+    const xz = @ctz(x);
+    const yz = @ctz(y);
+    const shift = @min(xz, yz);
+    x >>= @intCast(xz);
+    y >>= @intCast(yz);
+
+    var diff = y -% x;
+    while (diff != 0) : (diff = y -% x) {
+        // ctz is invariant under negation, we
+        // put it here to ease data dependencies,
+        // makes the CPU happy.
+        const zeros = @ctz(diff);
+        if (x > y) diff = -%diff;
+        y = @min(x, y);
+        x = diff >> @intCast(zeros);
     }
-    return x;
+    return y << @intCast(shift);
 }
 
-test "gcd" {
+test gcd {
+    const expectEqual = std.testing.expectEqual;
+
     try expectEqual(gcd(0, 5), 5);
     try expectEqual(gcd(5, 0), 5);
     try expectEqual(gcd(8, 12), 4);
@@ -45,4 +54,6 @@ test "gcd" {
     try expectEqual(gcd(49865, 69811), 9973);
     try expectEqual(gcd(300_000, 2_300_000), 100_000);
     try expectEqual(gcd(90000000_000_000_000_000_000, 2), 2);
+    try expectEqual(gcd(@as(u80, 90000000_000_000_000_000_000), 2), 2);
+    try expectEqual(gcd(300_000, @as(u32, 2_300_000)), 100_000);
 }