Skip to content

Add std.math.egcd (Extended GCD). Micro-optimize GCD. #25533

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 2 commits into
base: master
Choose a base branch
from
+268 −44
Draft

Add std.math.egcd (Extended GCD). Micro-optimize GCD. #25533

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
91139db
Microptimization.
two-horned Sep 27, 2025
ba5f824
Add EGCD.
two-horned Oct 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 4 additions & 26 deletions lib/std/crypto/ml_kem.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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);
const r = std.math.egcd(a, p);
assert(r.gcd == 1);
return r.x;
return r.bezout_coeff_1;
}

// Reduce mod q for testing.
Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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;

Expand Down
1 change: 1 addition & 0 deletions lib/std/math.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down
241 changes: 241 additions & 0 deletions lib/std/math/egcd.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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.
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this still a work in progress?

Copy link
Author

@two-horned two-horned Oct 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

One can avoid integer widening by implementing a custom function for terms of the form (a - b * c) / d. Something like this:

inline fn mul_div_thing(a: anytype, b: anytype, c: anytype, d: anytype) @TypeOf(b, c, d) {
    @setEvalBranchQuota(10_000_000);
    const S = @TypeOf(b, c, d);
    const sign_ = -1 * sign(b) * sign(c) * sign(d);
    var x = @abs(b);
    const y = @abs(c);
    const z = @abs(d);

    var acc = a;
    var counter: S = 0;
    if (x == 0) return sign_;
    while (x > 0) : (x -= 1) {
        acc += y;
        while (acc >= z) {
            acc -= z;
            counter += 1;
        }
    }
    return counter * sign_;
}

As you can see, this solution is highly inefficient because I am not doing long multiplication / division. If I implement long multiplication / division, this algorithm becomes much more complicated, so I didn't even bother really. Still, I want to encourage looking for something that doesn't require widening the integer type, that's why I left the comment.

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);
}
}
40 changes: 22 additions & 18 deletions lib/std/math/gcd.zig
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,20 +1,22 @@
//! 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)) {
// 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");
@compileError("`a` and `b` must be unsigned 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 (a == 0) return b;
Expand All @@ -26,25 +28,27 @@ pub fn gcd(a: anytype, b: anytype) @TypeOf(a, b) {
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);
x = @shrExact(x, @intCast(xz));
y = @shrExact(y, @intCast(yz));

var y_minus_x = y -% x;
while (y_minus_x != 0) : (y_minus_x = y -% x) {
const copy_x = x;
const zeros = @ctz(y_minus_x);
const carry = x < y;
x -%= y;
if (carry) {
x = y_minus_x;
y = copy_x;
}
x = @shrExact(x, @intCast(zeros));
}
return y << @intCast(shift);

return @shlExact(y, @intCast(shift));
}

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);
Expand Down