Skip to content

Commit ba5f824

Browse files
two-hornedtwo-horned
committed
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.
1 parent 91139db commit ba5f824

File tree

4 files changed

+249
-29
lines changed

4 files changed

+249
-29
lines changed

‎lib/std/crypto/ml_kem.zig‎

Lines changed: 4 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -634,33 +634,11 @@ test "invNTTReductions bounds" {
634634
}
635635
}
636636

637-
// Extended euclidean algorithm.
638-
//
639-
// For a, b finds x, y such that x a + y b = gcd(a, b). Used to compute
640-
// modular inverse.
641-
fn eea(a: anytype, b: @TypeOf(a)) EeaResult(@TypeOf(a)) {
642-
if (a == 0) {
643-
return .{ .gcd = b, .x = 0, .y = 1 };
644-
}
645-
const r = eea(@rem(b, a), a);
646-
return .{ .gcd = r.gcd, .x = r.y - @divTrunc(b, a) * r.x, .y = r.x };
647-
}
648-
649-
fn EeaResult(comptime T: type) type {
650-
return struct { gcd: T, x: T, y: T };
651-
}
652-
653-
// Returns least common multiple of a and b.
654-
fn lcm(a: anytype, b: @TypeOf(a)) @TypeOf(a) {
655-
const r = eea(a, b);
656-
return a * b / r.gcd;
657-
}
658-
659637
// Invert modulo p.
660638
fn invertMod(a: anytype, p: @TypeOf(a)) @TypeOf(a) {
661-
const r = eea(a, p);
639+
const r = std.math.egcd(a, p);
662640
assert(r.gcd == 1);
663-
return r.x;
641+
return r.bezout_coeff_1;
664642
}
665643

666644
// Reduce mod q for testing.
@@ -1054,7 +1032,7 @@ const Poly = struct {
10541032
var in_off: usize = 0;
10551033
var out_off: usize = 0;
10561034

1057-
const batch_size: usize = comptime lcm(@as(i16, d), 8);
1035+
const batch_size: usize = comptime std.math.lcm(@as(i16, d), 8);
10581036
const in_batch_size: usize = comptime batch_size / d;
10591037
const out_batch_size: usize = comptime batch_size / 8;
10601038

@@ -1118,7 +1096,7 @@ const Poly = struct {
11181096
var in_off: usize = 0;
11191097
var out_off: usize = 0;
11201098

1121-
const batch_size: usize = comptime lcm(@as(i16, d), 8);
1099+
const batch_size: usize = comptime std.math.lcm(@as(i16, d), 8);
11221100
const in_batch_size: usize = comptime batch_size / 8;
11231101
const out_batch_size: usize = comptime batch_size / d;
11241102

‎lib/std/math.zig‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -238,6 +238,7 @@ pub const sinh = @import("math/sinh.zig").sinh;
238238
pub const cosh = @import("math/cosh.zig").cosh;
239239
pub const tanh = @import("math/tanh.zig").tanh;
240240
pub const gcd = @import("math/gcd.zig").gcd;
241+
pub const egcd = @import("math/egcd.zig").egcd;
241242
pub const lcm = @import("math/lcm.zig").lcm;
242243
pub const gamma = @import("math/gamma.zig").gamma;
243244
pub const lgamma = @import("math/gamma.zig").lgamma;

‎lib/std/math/egcd.zig‎

Lines changed: 241 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,241 @@
1+
//! Extended Greatest Common Divisor (https://mathworld.wolfram.com/ExtendedGreatestCommonDivisor.html)
2+
const std = @import("../std.zig");
3+
4+
/// Result type of `egcd`.
5+
pub fn ExtendedGreatestCommonDivisor(S: anytype) type {
6+
const N = switch (S) {
7+
comptime_int => comptime_int,
8+
else => |T| std.meta.Int(.unsigned, @bitSizeOf(T)),
9+
};
10+
11+
return struct {
12+
gcd: N,
13+
bezout_coeff_1: S,
14+
bezout_coeff_2: S,
15+
};
16+
}
17+
18+
/// Returns the Extended Greatest Common Divisor (EGCD) of two signed integers (`a` and `b`) which are not both zero.
19+
pub fn egcd(a: anytype, b: anytype) ExtendedGreatestCommonDivisor(@TypeOf(a, b)) {
20+
const S = switch (@TypeOf(a, b)) {
21+
comptime_int => b: {
22+
const n = @max(@abs(a), @abs(b));
23+
break :b std.math.IntFittingRange(-n, n);
24+
},
25+
else => |T| T,
26+
};
27+
if (@typeInfo(S) != .int or @typeInfo(S).int.signedness != .signed) {
28+
@compileError("`a` and `b` must be signed integers");
29+
}
30+
31+
std.debug.assert(a != 0 or b != 0);
32+
33+
if (a == 0) return .{ .gcd = @abs(b), .bezout_coeff_1 = 0, .bezout_coeff_2 = std.math.sign(b) };
34+
if (b == 0) return .{ .gcd = @abs(a), .bezout_coeff_1 = std.math.sign(a), .bezout_coeff_2 = 0 };
35+
36+
const other: S, const odd: S, const shift, const switch_coeff = b: {
37+
const xz = @ctz(@as(S, a));
38+
const yz = @ctz(@as(S, b));
39+
break :b if (xz < yz) .{ b, a, xz, true } else .{ a, b, yz, false };
40+
};
41+
const toinv = @shrExact(other, @intCast(shift));
42+
const ctrl = @shrExact(odd, @intCast(shift)); // Invariant: |s|, |t|, |ctrl| < |MIN_OF(S)|
43+
const half_ctrl = 1 + @shrExact(ctrl - 1, 1);
44+
const abs_ctrl = @abs(ctrl);
45+
46+
var s: S = std.math.sign(toinv);
47+
var t: S = 0;
48+
49+
var x = @abs(toinv);
50+
var y = abs_ctrl;
51+
52+
{
53+
const xz = @ctz(x);
54+
x = @shrExact(x, @intCast(xz));
55+
for (0..xz) |_| {
56+
const half_s = s >> 1;
57+
if (s & 1 == 0)
58+
s = half_s
59+
else
60+
s = half_s + half_ctrl;
61+
}
62+
}
63+
64+
var y_minus_x = y -% x;
65+
while (y_minus_x != 0) : (y_minus_x = y -% x) {
66+
const t_minus_s = t - s;
67+
const copy_x = x;
68+
const copy_s = s;
69+
const xz = @ctz(y_minus_x);
70+
71+
s -= t;
72+
const carry = x < y;
73+
x -%= y;
74+
if (carry) {
75+
x = y_minus_x;
76+
y = copy_x;
77+
s = t_minus_s;
78+
t = copy_s;
79+
}
80+
x = @shrExact(x, @intCast(xz));
81+
for (0..xz) |_| {
82+
const half_s = s >> 1;
83+
if (s & 1 == 0)
84+
s = half_s
85+
else
86+
s = half_s + half_ctrl;
87+
}
88+
89+
if (s < 0) s = @intCast(abs_ctrl - @abs(s));
90+
}
91+
92+
// Using integer widening is only a temporary solution.
93+
const W = std.meta.Int(.signed, @bitSizeOf(S) * 2);
94+
t = @intCast(@divExact(y - @as(W, s) * toinv, ctrl));
95+
const final_s, const final_t = if (switch_coeff) .{ t, s } else .{ s, t };
96+
return .{
97+
.gcd = @shlExact(y, @intCast(shift)),
98+
.bezout_coeff_1 = final_s,
99+
.bezout_coeff_2 = final_t,
100+
};
101+
}
102+
103+
test {
104+
{
105+
const a: i2 = 0;
106+
const b: i2 = 1;
107+
const r = egcd(a, b);
108+
const g = r.gcd;
109+
const s: i2 = r.bezout_coeff_1;
110+
const t: i2 = r.bezout_coeff_2;
111+
try std.testing.expect(s * a + t * b == g);
112+
}
113+
{
114+
const a: i8 = -128;
115+
const b: i8 = 127;
116+
const r = egcd(a, b);
117+
const g = r.gcd;
118+
const s: i16 = r.bezout_coeff_1;
119+
const t: i16 = r.bezout_coeff_2;
120+
try std.testing.expect(s * a + t * b == g);
121+
}
122+
{
123+
const a: i16 = -32768;
124+
const b: i16 = -32768;
125+
const r = egcd(a, b);
126+
const g = r.gcd;
127+
const s: i32 = r.bezout_coeff_1;
128+
const t: i32 = r.bezout_coeff_2;
129+
try std.testing.expect(s * a + t * b == g);
130+
}
131+
{
132+
const a: i32 = 128;
133+
const b: i32 = 112;
134+
const r = egcd(a, b);
135+
const g = r.gcd;
136+
const s: i64 = r.bezout_coeff_1;
137+
const t: i64 = r.bezout_coeff_2;
138+
try std.testing.expect(s * a + t * b == g);
139+
}
140+
{
141+
const a: i32 = 4 * 89;
142+
const b: i32 = 2 * 17;
143+
const r = egcd(a, b);
144+
const g = r.gcd;
145+
const s: i64 = r.bezout_coeff_1;
146+
const t: i64 = r.bezout_coeff_2;
147+
try std.testing.expect(s * a + t * b == g);
148+
}
149+
{
150+
const a: i8 = 127;
151+
const b: i8 = 126;
152+
const r = egcd(a, b);
153+
const g = r.gcd;
154+
const s: i16 = r.bezout_coeff_1;
155+
const t: i16 = r.bezout_coeff_2;
156+
try std.testing.expect(s * a + t * b == g);
157+
}
158+
{
159+
const a: i4 = -8;
160+
const b: i4 = 1;
161+
const r = egcd(a, b);
162+
const g = r.gcd;
163+
const s = r.bezout_coeff_1;
164+
const t = r.bezout_coeff_2;
165+
try std.testing.expect(s * a + t * b == g);
166+
}
167+
{
168+
const a: i4 = -8;
169+
const b: i4 = 5;
170+
const r = egcd(a, b);
171+
const g = r.gcd;
172+
// Avoid overflow in assert.
173+
const s: i8 = r.bezout_coeff_1;
174+
const t: i8 = r.bezout_coeff_2;
175+
try std.testing.expect(s * a + t * b == g);
176+
}
177+
{
178+
const a: i32 = 0;
179+
const b: i32 = 5;
180+
const r = egcd(a, b);
181+
const g = r.gcd;
182+
const s = r.bezout_coeff_1;
183+
const t = r.bezout_coeff_2;
184+
try std.testing.expect(s * a + t * b == g);
185+
}
186+
{
187+
const a: i32 = 5;
188+
const b: i32 = 0;
189+
const r = egcd(a, b);
190+
const g = r.gcd;
191+
const s = r.bezout_coeff_1;
192+
const t = r.bezout_coeff_2;
193+
try std.testing.expect(s * a + t * b == g);
194+
}
195+
196+
{
197+
const a: i32 = 21;
198+
const b: i32 = 15;
199+
const r = egcd(a, b);
200+
const g = r.gcd;
201+
const s = r.bezout_coeff_1;
202+
const t = r.bezout_coeff_2;
203+
try std.testing.expect(s * a + t * b == g);
204+
}
205+
{
206+
const a: i32 = -21;
207+
const b: i32 = 15;
208+
const r = egcd(a, b);
209+
const g = r.gcd;
210+
const s = r.bezout_coeff_1;
211+
const t = r.bezout_coeff_2;
212+
try std.testing.expect(s * a + t * b == g);
213+
}
214+
{
215+
const a = -21;
216+
const b = 15;
217+
const r = egcd(a, b);
218+
const g = r.gcd;
219+
const s = r.bezout_coeff_1;
220+
const t = r.bezout_coeff_2;
221+
try std.testing.expect(s * a + t * b == g);
222+
}
223+
{
224+
const a = 927372692193078999176;
225+
const b = 573147844013817084101;
226+
const r = egcd(a, b);
227+
const g = r.gcd;
228+
const s = r.bezout_coeff_1;
229+
const t = r.bezout_coeff_2;
230+
try std.testing.expect(s * a + t * b == g);
231+
}
232+
{
233+
const a = 453973694165307953197296969697410619233826;
234+
const b = 280571172992510140037611932413038677189525;
235+
const r = egcd(a, b);
236+
const g = r.gcd;
237+
const s = r.bezout_coeff_1;
238+
const t = r.bezout_coeff_2;
239+
try std.testing.expect(s * a + t * b == g);
240+
}
241+
}

‎lib/std/math/gcd.zig‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
1-
//! Greatest common divisor (https://mathworld.wolfram.com/GreatestCommonDivisor.html)
2-
const std = @import("std");
1+
//! Greatest Common Divisor (https://mathworld.wolfram.com/GreatestCommonDivisor.html)
2+
const std = @import("../std.zig");
33

4-
/// Returns the greatest common divisor (GCD) of two unsigned integers (`a` and `b`) which are not both zero.
4+
/// Returns the Greatest Common Divisor (GCD) of two unsigned integers (`a` and `b`) which are not both zero.
55
/// For example, the GCD of `8` and `12` is `4`, that is, `gcd(8, 12) == 4`.
66
pub fn gcd(a: anytype, b: anytype) @TypeOf(a, b) {
77
const N = switch (@TypeOf(a, b)) {

0 commit comments

Comments
 (0)