authorgravatar for 124872+jedisct1@users.noreply.github.comFrank Denis <124872+jedisct1@users.noreply.github.com> 2025-11-18 16:39:58+01:00
committergravatar for noreply@github.comGitHub <noreply@github.com> 2025-11-18 15:39:58+00:00
log4ea472808437c932f6240c0df5e4d5912d282052
tree9ecb59130b643be94937c84c94d49fe1278bd10d
parent73f863a6fb1e844ec312345b3160e3987fe6586d
signaturebadge-check Signed by PGP key B5690EEEBB952194

Align ML-KEM code with ML-DSA (#25964)

This will facilitate maintainance and code sharing between primitives.

1 files changed, 338 insertions(+), 196 deletions(-)

lib/std/crypto/ml_kem.zig+338-196
......@@ -105,19 +105,20 @@ const crypto = std.crypto;
105105const errors = std.crypto.errors;
106106const math = std.math;
107107const mem = std.mem;
108const RndGen = std.Random.DefaultPrng;
109108const sha3 = crypto.hash.sha3;
110109
111// Q is the parameter q ≡ 3329 = 2¹¹ + 2¹⁰ + 2⁸ + 1.
110const RndGen = std.Random.DefaultPrng;
111
112// Q is the modulus q ≡ 3329 = 2¹¹ + 2¹⁰ + 2⁸ + 1
112113const Q: i16 = 3329;
113114
114// Montgomery R
115// Montgomery R = 2^16 mod Q (for Montgomery multiplication)
115116const R: i32 = 1 << 16;
116117
117// Parameter n, degree of polynomials.
118// N is the degree of polynomials (polynomial ring dimension)
118119const N: usize = 256;
119120
120// Size of "small" vectors used in encryption blinds.
121// eta2 is the size of "small" vectors used in encryption blinds
121122const eta2: u8 = 2;
122123
123124const Params = struct {
......@@ -215,7 +216,7 @@ fn Kyber(comptime p: Params) type {
215216 pub const ciphertext_length = Poly.compressedSize(p.du) * p.k + Poly.compressedSize(p.dv);
216217
217218 const Self = @This();
218 const V = Vec(p.k);
219 const V = PolyVec(p.k);
219220 const M = Mat(p.k);
220221
221222 /// Length (in bytes) of a shared secret.
......@@ -241,7 +242,7 @@ fn Kyber(comptime p: Params) type {
241242 hpk: [h_length]u8, // H(pk)
242243
243244 /// Size of a serialized representation of the key, in bytes.
244 pub const bytes_length = InnerPk.bytes_length;
245 pub const encoded_length = InnerPk.encoded_length;
245246
246247 /// Generates a shared secret, and encapsulates it for the public key.
247248 /// If `seed` is `null`, a random seed is used. This is recommended.
......@@ -289,14 +290,14 @@ fn Kyber(comptime p: Params) type {
289290 }
290291
291292 /// Serializes the key into a byte array.
292 pub fn toBytes(pk: PublicKey) [bytes_length]u8 {
293 pub fn toBytes(pk: PublicKey) [encoded_length]u8 {
293294 return pk.pk.toBytes();
294295 }
295296
296297 /// Deserializes the key from a byte array.
297 pub fn fromBytes(buf: *const [bytes_length]u8) errors.NonCanonicalError!PublicKey {
298 pub fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!PublicKey {
298299 var ret: PublicKey = undefined;
299 ret.pk = try InnerPk.fromBytes(buf[0..InnerPk.bytes_length]);
300 ret.pk = try InnerPk.fromBytes(buf[0..InnerPk.encoded_length]);
300301 sha3.Sha3_256.hash(buf, &ret.hpk, .{});
301302 return ret;
302303 }
......@@ -310,8 +311,8 @@ fn Kyber(comptime p: Params) type {
310311 z: [shared_length]u8,
311312
312313 /// Size of a serialized representation of the key, in bytes.
313 pub const bytes_length: usize =
314 InnerSk.bytes_length + InnerPk.bytes_length + h_length + shared_length;
314 pub const encoded_length: usize =
315 InnerSk.encoded_length + InnerPk.encoded_length + h_length + shared_length;
315316
316317 /// Decapsulates the shared secret within ct using the private key.
317318 pub fn decaps(sk: SecretKey, ct: *const [ciphertext_length]u8) ![shared_length]u8 {
......@@ -346,18 +347,18 @@ fn Kyber(comptime p: Params) type {
346347 }
347348
348349 /// Serializes the key into a byte array.
349 pub fn toBytes(sk: SecretKey) [bytes_length]u8 {
350 pub fn toBytes(sk: SecretKey) [encoded_length]u8 {
350351 return sk.sk.toBytes() ++ sk.pk.toBytes() ++ sk.hpk ++ sk.z;
351352 }
352353
353354 /// Deserializes the key from a byte array.
354 pub fn fromBytes(buf: *const [bytes_length]u8) errors.NonCanonicalError!SecretKey {
355 pub fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!SecretKey {
355356 var ret: SecretKey = undefined;
356357 comptime var s: usize = 0;
357 ret.sk = InnerSk.fromBytes(buf[s .. s + InnerSk.bytes_length]);
358 s += InnerSk.bytes_length;
359 ret.pk = try InnerPk.fromBytes(buf[s .. s + InnerPk.bytes_length]);
360 s += InnerPk.bytes_length;
358 ret.sk = InnerSk.fromBytes(buf[s .. s + InnerSk.encoded_length]);
359 s += InnerSk.encoded_length;
360 ret.pk = try InnerPk.fromBytes(buf[s .. s + InnerPk.encoded_length]);
361 s += InnerPk.encoded_length;
361362 ret.hpk = buf[s..][0..h_length].*;
362363 s += h_length;
363364 ret.z = buf[s..][0..shared_length].*;
......@@ -418,7 +419,7 @@ fn Kyber(comptime p: Params) type {
418419 // Cached values
419420 aT: M,
420421
421 const bytes_length = V.bytes_length + 32;
422 const encoded_length = V.encoded_length + 32;
422423
423424 fn encrypt(
424425 pk: InnerPk,
......@@ -436,7 +437,7 @@ fn Kyber(comptime p: Params) type {
436437 // Note that coefficients of r are bounded by q and those of Aᵀ
437438 // are bounded by 4.5q and so their product is bounded by 2¹⁵q
438439 // as required for multiplication.
439 u.ps[i] = pk.aT.vs[i].dotHat(rh);
440 u.ps[i] = pk.aT.rows[i].dotHat(rh);
440441 }
441442
442443 // Aᵀ and r were not in Montgomery form, so the Montgomery
......@@ -451,14 +452,14 @@ fn Kyber(comptime p: Params) type {
451452 return u.compress(p.du) ++ v.compress(p.dv);
452453 }
453454
454 fn toBytes(pk: InnerPk) [bytes_length]u8 {
455 fn toBytes(pk: InnerPk) [encoded_length]u8 {
455456 return pk.th.toBytes() ++ pk.rho;
456457 }
457458
458 fn fromBytes(buf: *const [bytes_length]u8) errors.NonCanonicalError!InnerPk {
459 fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!InnerPk {
459460 var ret: InnerPk = undefined;
460461
461 const th_bytes = buf[0..V.bytes_length];
462 const th_bytes = buf[0..V.encoded_length];
462463 ret.th = V.fromBytes(th_bytes).normalize();
463464
464465 if (p.ml_kem) {
......@@ -468,7 +469,7 @@ fn Kyber(comptime p: Params) type {
468469 }
469470 }
470471
471 ret.rho = buf[V.bytes_length..bytes_length].*;
472 ret.rho = buf[V.encoded_length..encoded_length].*;
472473 ret.aT = M.uniform(ret.rho, true);
473474 return ret;
474475 }
......@@ -477,7 +478,7 @@ fn Kyber(comptime p: Params) type {
477478 // Private key of the inner PKE
478479 const InnerSk = struct {
479480 sh: V, // NTT(s), normalized
480 const bytes_length = V.bytes_length;
481 const encoded_length = V.encoded_length;
481482
482483 fn decrypt(sk: InnerSk, ct: *const [ciphertext_length]u8) [inner_plaintext_length]u8 {
483484 const u = V.decompress(p.du, ct[0..comptime V.compressedSize(p.du)]);
......@@ -491,11 +492,11 @@ fn Kyber(comptime p: Params) type {
491492 .normalize().compress(1);
492493 }
493494
494 fn toBytes(sk: InnerSk) [bytes_length]u8 {
495 fn toBytes(sk: InnerSk) [encoded_length]u8 {
495496 return sk.sh.toBytes();
496497 }
497498
498 fn fromBytes(buf: *const [bytes_length]u8) InnerSk {
499 fn fromBytes(buf: *const [encoded_length]u8) InnerSk {
499500 var ret: InnerSk = undefined;
500501 ret.sh = V.fromBytes(buf).normalize();
501502 return ret;
......@@ -516,7 +517,7 @@ fn Kyber(comptime p: Params) type {
516517 // Sample secret vector s.
517518 sk.sh = V.noise(p.eta1, 0, sigma).ntt().normalize();
518519
519 const eh = Vec(p.k).noise(p.eta1, p.k, sigma).ntt(); // sample blind e.
520 const eh = PolyVec(p.k).noise(p.eta1, p.k, sigma).ntt(); // sample blind e.
520521 var th: V = undefined;
521522
522523 // Next, we compute t = A s + e.
......@@ -528,7 +529,7 @@ fn Kyber(comptime p: Params) type {
528529 // multiplications in the inner product added a factor R⁻¹ which
529530 // we'll cancel out with toMont(). This will also ensure the
530531 // coefficients of th are bounded in absolute value by q.
531 th.ps[i] = pk.aT.vs[i].dotHat(sk.sh).toMont();
532 th.ps[i] = pk.aT.rows[i].dotHat(sk.sh).toMont();
532533 }
533534
534535 pk.th = th.add(eh).normalize(); // bounded by 8q
......@@ -565,7 +566,6 @@ const zetas = computeZetas();
565566// not enough, the other coefficient is reduced as well.
566567//
567568// This is actually optimal, as proven in https://eprint.iacr.org/2020/1377.pdf
568// TODO generate comptime?
569569const inv_ntt_reductions = [_]i16{
570570 -1, // after layer 1
571571 -1, // after layer 2
......@@ -634,31 +634,8 @@ 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.
641fn 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
649fn 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.
654fn lcm(a: anytype, b: @TypeOf(a)) @TypeOf(a) {
655 const r = eea(a, b);
656 return a * b / r.gcd;
657}
658
659// Invert modulo p.
660637fn invertMod(a: anytype, p: @TypeOf(a)) @TypeOf(a) {
661 const r = eea(a, p);
638 const r = extendedEuclidean(@TypeOf(a), a, p);
662639 assert(r.gcd == 1);
663640 return r.x;
664641}
......@@ -788,31 +765,12 @@ test "Test csubq" {
788765 }
789766}
790767
791// Compute a^s mod p.
792fn mpow(a: anytype, s: @TypeOf(a), p: @TypeOf(a)) @TypeOf(a) {
793 var ret: @TypeOf(a) = 1;
794 var s2 = s;
795 var a2 = a;
796
797 while (true) {
798 if (s2 & 1 == 1) {
799 ret = @mod(ret * a2, p);
800 }
801 s2 >>= 1;
802 if (s2 == 0) {
803 break;
804 }
805 a2 = @mod(a2 * a2, p);
806 }
807 return ret;
808}
809
810768// Computes zetas table used by ntt and invNTT.
811769fn computeZetas() [128]i16 {
812770 @setEvalBranchQuota(10000);
813771 var ret: [128]i16 = undefined;
814772 for (&ret, 0..) |*r, i| {
815 const t = @as(i16, @intCast(mpow(@as(i32, zeta), @bitReverse(@as(u7, @intCast(i))), Q)));
773 const t = @as(i16, @intCast(modularPow(i32, zeta, @bitReverse(@as(u7, @intCast(i))), Q)));
816774 r.* = csubq(feBarrettReduce(feToMont(t)));
817775 }
818776 return ret;
......@@ -828,9 +786,10 @@ fn computeZetas() [128]i16 {
828786const Poly = struct {
829787 cs: [N]i16,
830788
831 const bytes_length = N / 2 * 3;
789 const encoded_length = N / 2 * 3;
832790 const zero: Poly = .{ .cs = .{0} ** N };
833791
792 // Add two polynomials (coefficients not normalized)
834793 fn add(a: Poly, b: Poly) Poly {
835794 var ret: Poly = undefined;
836795 for (0..N) |i| {
......@@ -839,6 +798,7 @@ const Poly = struct {
839798 return ret;
840799 }
841800
801 // Subtract two polynomials (coefficients not normalized)
842802 fn sub(a: Poly, b: Poly) Poly {
843803 var ret: Poly = undefined;
844804 for (0..N) |i| {
......@@ -847,25 +807,6 @@ const Poly = struct {
847807 return ret;
848808 }
849809
850 // For testing, generates a random polynomial with for each
851 // coefficient |x| ≤ q.
852 fn randAbsLeqQ(rnd: anytype) Poly {
853 var ret: Poly = undefined;
854 for (0..N) |i| {
855 ret.cs[i] = rnd.random().intRangeAtMost(i16, -Q, Q);
856 }
857 return ret;
858 }
859
860 // For testing, generates a random normalized polynomial.
861 fn randNormalized(rnd: anytype) Poly {
862 var ret: Poly = undefined;
863 for (0..N) |i| {
864 ret.cs[i] = rnd.random().intRangeLessThan(i16, 0, Q);
865 }
866 return ret;
867 }
868
869810 // Executes a forward "NTT" on p.
870811 //
871812 // Assumes the coefficients are in absolute value ≤q. The resulting
......@@ -1054,7 +995,7 @@ const Poly = struct {
1054995 var in_off: usize = 0;
1055996 var out_off: usize = 0;
1056997
1057 const batch_size: usize = comptime lcm(@as(i16, d), 8);
998 const batch_size: usize = comptime math.lcm(d, 8);
1058999 const in_batch_size: usize = comptime batch_size / d;
10591000 const out_batch_size: usize = comptime batch_size / 8;
10601001
......@@ -1118,7 +1059,7 @@ const Poly = struct {
11181059 var in_off: usize = 0;
11191060 var out_off: usize = 0;
11201061
1121 const batch_size: usize = comptime lcm(@as(i16, d), 8);
1062 const batch_size: usize = comptime math.lcm(d, 8);
11221063 const in_batch_size: usize = comptime batch_size / 8;
11231064 const out_batch_size: usize = comptime batch_size / d;
11241065
......@@ -1275,53 +1216,23 @@ const Poly = struct {
12751216 return ret;
12761217 }
12771218
1278 // Sample p uniformly from the given seed and x and y coordinates.
12791219 fn uniform(seed: [32]u8, x: u8, y: u8) Poly {
1280 var h = sha3.Shake128.init(.{});
1281 const suffix: [2]u8 = .{ x, y };
1282 h.update(&seed);
1283 h.update(&suffix);
1284
1285 const buf_len = sha3.Shake128.block_length; // rate SHAKE-128
1286 var buf: [buf_len]u8 = undefined;
1287
1288 var ret: Poly = undefined;
1289 var i: usize = 0; // index into ret.cs
1290 outer: while (true) {
1291 h.squeeze(&buf);
1292
1293 var j: usize = 0; // index into buf
1294 while (j < buf_len) : (j += 3) {
1295 const b0 = @as(u16, buf[j]);
1296 const b1 = @as(u16, buf[j + 1]);
1297 const b2 = @as(u16, buf[j + 2]);
1298
1299 const ts: [2]u16 = .{
1300 b0 | ((b1 & 0xf) << 8),
1301 (b1 >> 4) | (b2 << 4),
1302 };
1303
1304 inline for (ts) |t| {
1305 if (t < Q) {
1306 ret.cs[i] = @as(i16, @intCast(t));
1307 i += 1;
1308
1309 if (i == N) {
1310 break :outer;
1311 }
1312 }
1313 }
1314 }
1315 }
1316
1317 return ret;
1220 const domain_sep: [2]u8 = .{ x, y };
1221 return sampleUniformRejection(
1222 Poly,
1223 Q,
1224 12,
1225 N,
1226 &seed,
1227 &domain_sep,
1228 );
13181229 }
13191230
13201231 // Packs p.
13211232 //
13221233 // Assumes p is normalized (and not just Barrett reduced).
1323 fn toBytes(p: Poly) [bytes_length]u8 {
1324 var ret: [bytes_length]u8 = undefined;
1234 fn toBytes(p: Poly) [encoded_length]u8 {
1235 var ret: [encoded_length]u8 = undefined;
13251236 for (0..comptime N / 2) |i| {
13261237 const t0 = @as(u16, @intCast(p.cs[2 * i]));
13271238 const t1 = @as(u16, @intCast(p.cs[2 * i + 1]));
......@@ -1335,7 +1246,7 @@ const Poly = struct {
13351246 // Unpacks a Poly from buf.
13361247 //
13371248 // p will not be normalized; instead 0 ≤ p[i] < 4096.
1338 fn fromBytes(buf: *const [bytes_length]u8) Poly {
1249 fn fromBytes(buf: *const [encoded_length]u8) Poly {
13391250 var ret: Poly = undefined;
13401251 for (0..comptime N / 2) |i| {
13411252 const b0 = @as(i16, buf[3 * i]);
......@@ -1348,71 +1259,65 @@ const Poly = struct {
13481259 }
13491260};
13501261
1351// A vector of K polynomials.
1352fn Vec(comptime K: u8) type {
1262// A vector of k polynomials.
1263fn PolyVec(comptime k: u8) type {
13531264 return struct {
1354 ps: [K]Poly,
1265 ps: [k]Poly,
13551266
13561267 const Self = @This();
1357 const bytes_length = K * Poly.bytes_length;
1268 const encoded_length = k * Poly.encoded_length;
13581269
13591270 fn compressedSize(comptime d: u8) usize {
1360 return Poly.compressedSize(d) * K;
1271 return Poly.compressedSize(d) * k;
13611272 }
13621273
1363 fn ntt(a: Self) Self {
1274 /// Apply unary operation to each polynomial
1275 fn map(v: Self, comptime op: fn (Poly) Poly) Self {
13641276 var ret: Self = undefined;
1365 for (0..K) |i| {
1366 ret.ps[i] = a.ps[i].ntt();
1277 inline for (0..k) |i| {
1278 ret.ps[i] = op(v.ps[i]);
13671279 }
13681280 return ret;
13691281 }
13701282
1371 fn invNTT(a: Self) Self {
1283 /// Apply binary operation pairwise
1284 fn mapBinary(a: Self, b: Self, comptime op: fn (Poly, Poly) Poly) Self {
13721285 var ret: Self = undefined;
1373 for (0..K) |i| {
1374 ret.ps[i] = a.ps[i].invNTT();
1286 inline for (0..k) |i| {
1287 ret.ps[i] = op(a.ps[i], b.ps[i]);
13751288 }
13761289 return ret;
13771290 }
13781291
1379 fn normalize(a: Self) Self {
1380 var ret: Self = undefined;
1381 for (0..K) |i| {
1382 ret.ps[i] = a.ps[i].normalize();
1383 }
1384 return ret;
1292 fn ntt(v: Self) Self {
1293 return map(v, Poly.ntt);
13851294 }
13861295
1387 fn barrettReduce(a: Self) Self {
1388 var ret: Self = undefined;
1389 for (0..K) |i| {
1390 ret.ps[i] = a.ps[i].barrettReduce();
1391 }
1392 return ret;
1296 fn invNTT(v: Self) Self {
1297 return map(v, Poly.invNTT);
1298 }
1299
1300 fn normalize(v: Self) Self {
1301 return map(v, Poly.normalize);
1302 }
1303
1304 fn barrettReduce(v: Self) Self {
1305 return map(v, Poly.barrettReduce);
13931306 }
13941307
13951308 fn add(a: Self, b: Self) Self {
1396 var ret: Self = undefined;
1397 for (0..K) |i| {
1398 ret.ps[i] = a.ps[i].add(b.ps[i]);
1399 }
1400 return ret;
1309 return mapBinary(a, b, Poly.add);
14011310 }
14021311
14031312 fn sub(a: Self, b: Self) Self {
1404 var ret: Self = undefined;
1405 for (0..K) |i| {
1406 ret.ps[i] = a.ps[i].sub(b.ps[i]);
1407 }
1408 return ret;
1313 return mapBinary(a, b, Poly.sub);
14091314 }
14101315
14111316 // Samples v[i] from centered binomial distribution with the given η,
14121317 // seed and nonce+i.
14131318 fn noise(comptime eta: u8, nonce: u8, seed: *const [32]u8) Self {
14141319 var ret: Self = undefined;
1415 for (0..K) |i| {
1320 for (0..k) |i| {
14161321 ret.ps[i] = Poly.noise(eta, nonce + @as(u8, @intCast(i)), seed);
14171322 }
14181323 return ret;
......@@ -1428,7 +1333,7 @@ fn Vec(comptime K: u8) type {
14281333 // of the Montgomery factor.
14291334 fn dotHat(a: Self, b: Self) Poly {
14301335 var ret: Poly = Poly.zero;
1431 for (0..K) |i| {
1336 for (0..k) |i| {
14321337 ret = ret.add(a.ps[i].mulHat(b.ps[i]));
14331338 }
14341339 return ret;
......@@ -1437,7 +1342,7 @@ fn Vec(comptime K: u8) type {
14371342 fn compress(v: Self, comptime d: u8) [compressedSize(d)]u8 {
14381343 const cs = comptime Poly.compressedSize(d);
14391344 var ret: [compressedSize(d)]u8 = undefined;
1440 inline for (0..K) |i| {
1345 inline for (0..k) |i| {
14411346 ret[i * cs .. (i + 1) * cs].* = v.ps[i].compress(d);
14421347 }
14431348 return ret;
......@@ -1446,27 +1351,27 @@ fn Vec(comptime K: u8) type {
14461351 fn decompress(comptime d: u8, buf: *const [compressedSize(d)]u8) Self {
14471352 const cs = comptime Poly.compressedSize(d);
14481353 var ret: Self = undefined;
1449 inline for (0..K) |i| {
1354 inline for (0..k) |i| {
14501355 ret.ps[i] = Poly.decompress(d, buf[i * cs .. (i + 1) * cs]);
14511356 }
14521357 return ret;
14531358 }
14541359
14551360 /// Serializes the key into a byte array.
1456 fn toBytes(v: Self) [bytes_length]u8 {
1457 var ret: [bytes_length]u8 = undefined;
1458 inline for (0..K) |i| {
1459 ret[i * Poly.bytes_length .. (i + 1) * Poly.bytes_length].* = v.ps[i].toBytes();
1361 fn toBytes(v: Self) [encoded_length]u8 {
1362 var ret: [encoded_length]u8 = undefined;
1363 inline for (0..k) |i| {
1364 ret[i * Poly.encoded_length .. (i + 1) * Poly.encoded_length].* = v.ps[i].toBytes();
14601365 }
14611366 return ret;
14621367 }
14631368
14641369 /// Deserializes the key from a byte array.
1465 fn fromBytes(buf: *const [bytes_length]u8) Self {
1370 fn fromBytes(buf: *const [encoded_length]u8) Self {
14661371 var ret: Self = undefined;
1467 inline for (0..K) |i| {
1372 inline for (0..k) |i| {
14681373 ret.ps[i] = Poly.fromBytes(
1469 buf[i * Poly.bytes_length .. (i + 1) * Poly.bytes_length],
1374 buf[i * Poly.encoded_length .. (i + 1) * Poly.encoded_length],
14701375 );
14711376 }
14721377 return ret;
......@@ -1474,19 +1379,19 @@ fn Vec(comptime K: u8) type {
14741379 };
14751380}
14761381
1477// A matrix of K vectors
1478fn Mat(comptime K: u8) type {
1382// A matrix of k vectors
1383fn Mat(comptime k: u8) type {
14791384 return struct {
14801385 const Self = @This();
1481 vs: [K]Vec(K),
1386 rows: [k]PolyVec(k),
14821387
14831388 fn uniform(seed: [32]u8, comptime transposed: bool) Self {
14841389 var ret: Self = undefined;
14851390 var i: u8 = 0;
1486 while (i < K) : (i += 1) {
1391 while (i < k) : (i += 1) {
14871392 var j: u8 = 0;
1488 while (j < K) : (j += 1) {
1489 ret.vs[i].ps[j] = Poly.uniform(
1393 while (j < k) : (j += 1) {
1394 ret.rows[i].ps[j] = Poly.uniform(
14901395 seed,
14911396 if (transposed) i else j,
14921397 if (transposed) j else i,
......@@ -1499,9 +1404,9 @@ fn Mat(comptime K: u8) type {
14991404 // Returns transpose of A
15001405 fn transpose(m: Self) Self {
15011406 var ret: Self = undefined;
1502 for (0..K) |i| {
1503 for (0..K) |j| {
1504 ret.vs[i].ps[j] = m.vs[j].ps[i];
1407 for (0..k) |i| {
1408 for (0..k) |j| {
1409 ret.rows[i].ps[j] = m.rows[j].ps[i];
15051410 }
15061411 }
15071412 return ret;
......@@ -1522,12 +1427,30 @@ fn cmov(comptime len: usize, dst: *[len]u8, src: [len]u8, b: u1) void {
15221427 }
15231428}
15241429
1430// Test helper: generates a random polynomial with each coefficient |x| ≤ q
1431fn randPolyAbsLeqQ(rnd: anytype) Poly {
1432 var ret: Poly = undefined;
1433 for (0..N) |i| {
1434 ret.cs[i] = rnd.random().intRangeAtMost(i16, -Q, Q);
1435 }
1436 return ret;
1437}
1438
1439// Test helper: generates a random normalized polynomial
1440fn randPolyNormalized(rnd: anytype) Poly {
1441 var ret: Poly = undefined;
1442 for (0..N) |i| {
1443 ret.cs[i] = rnd.random().intRangeLessThan(i16, 0, Q);
1444 }
1445 return ret;
1446}
1447
15251448test "MulHat" {
15261449 var rnd = RndGen.init(0);
15271450
15281451 for (0..100) |_| {
1529 const a = Poly.randAbsLeqQ(&rnd);
1530 const b = Poly.randAbsLeqQ(&rnd);
1452 const a = randPolyAbsLeqQ(&rnd);
1453 const b = randPolyAbsLeqQ(&rnd);
15311454
15321455 const p2 = a.ntt().mulHat(b.ntt()).barrettReduce().invNTT().normalize();
15331456 var p: Poly = undefined;
......@@ -1557,7 +1480,7 @@ test "NTT" {
15571480 var rnd = RndGen.init(0);
15581481
15591482 for (0..1000) |_| {
1560 var p = Poly.randAbsLeqQ(&rnd);
1483 var p = randPolyAbsLeqQ(&rnd);
15611484 const q = p.toMont().normalize();
15621485 p = p.ntt();
15631486
......@@ -1580,7 +1503,7 @@ test "Compression" {
15801503 var rnd = RndGen.init(0);
15811504 inline for (.{ 1, 4, 5, 10, 11 }) |d| {
15821505 for (0..1000) |_| {
1583 const p = Poly.randNormalized(&rnd);
1506 const p = randPolyNormalized(&rnd);
15841507 const pp = p.compress(d);
15851508 const pq = Poly.decompress(d, &pp).compress(d);
15861509 try testing.expectEqual(pp, pq);
......@@ -1671,7 +1594,7 @@ test "Polynomial packing" {
16711594 var rnd = RndGen.init(0);
16721595
16731596 for (0..1000) |_| {
1674 const p = Poly.randNormalized(&rnd);
1597 const p = randPolyNormalized(&rnd);
16751598 try testing.expectEqual(Poly.fromBytes(&p.toBytes()), p);
16761599 }
16771600}
......@@ -1839,3 +1762,222 @@ const NistDRBG = struct {
18391762 return ret;
18401763 }
18411764};
1765
1766/// Extended Euclidian Algorithm
1767/// Only meant to be used on comptime values; correctness matters, performance doesn't.
1768fn extendedEuclidean(comptime T: type, comptime a_: T, comptime b_: T) struct { gcd: T, x: T, y: T } {
1769 var a = a_;
1770 var b = b_;
1771 var x0: T = 1;
1772 var x1: T = 0;
1773 var y0: T = 0;
1774 var y1: T = 1;
1775
1776 while (b != 0) {
1777 const q = @divTrunc(a, b);
1778 const temp_a = a;
1779 a = b;
1780 b = temp_a - q * b;
1781
1782 const temp_x = x0;
1783 x0 = x1;
1784 x1 = temp_x - q * x1;
1785
1786 const temp_y = y0;
1787 y0 = y1;
1788 y1 = temp_y - q * y1;
1789 }
1790
1791 return .{ .gcd = a, .x = x0, .y = y0 };
1792}
1793
1794/// Modular inversion: computes a^(-1) mod p
1795/// Requires gcd(a,p) = 1. The result is normalized to the range [0, p).
1796fn modularInverse(comptime T: type, comptime a: T, comptime p: T) T {
1797 // Use a signed type for EEA computation
1798 const type_info = @typeInfo(T);
1799 const SignedT = if (type_info == .int and type_info.int.signedness == .unsigned)
1800 std.meta.Int(.signed, type_info.int.bits)
1801 else
1802 T;
1803
1804 const a_signed = @as(SignedT, @intCast(a));
1805 const p_signed = @as(SignedT, @intCast(p));
1806
1807 const r = extendedEuclidean(SignedT, a_signed, p_signed);
1808 assert(r.gcd == 1);
1809
1810 // Normalize result to [0, p)
1811 var result = r.x;
1812 while (result < 0) {
1813 result += p_signed;
1814 }
1815
1816 return @intCast(result);
1817}
1818
1819/// Modular exponentiation: computes a^s mod p using square-and-multiply algorithm.
1820fn modularPow(comptime T: type, comptime a: T, s: T, comptime p: T) T {
1821 const type_info = @typeInfo(T);
1822 const bits = type_info.int.bits;
1823 const WideT = std.meta.Int(.unsigned, bits * 2);
1824
1825 var ret: T = 1;
1826 var base: T = a;
1827 var exp = s;
1828
1829 while (exp > 0) {
1830 if (exp & 1 == 1) {
1831 ret = @intCast((@as(WideT, ret) * @as(WideT, base)) % p);
1832 }
1833 base = @intCast((@as(WideT, base) * @as(WideT, base)) % p);
1834 exp >>= 1;
1835 }
1836
1837 return ret;
1838}
1839
1840/// Creates an all-ones or all-zeros mask from a single bit value.
1841/// Returns all 1s (0xFF...FF) if bit == 1, all 0s if bit == 0.
1842fn bitMask(comptime T: type, bit: T) T {
1843 const type_info = @typeInfo(T);
1844 if (type_info != .int or type_info.int.signedness != .unsigned) {
1845 @compileError("bitMask requires an unsigned integer type");
1846 }
1847 return -%bit;
1848}
1849
1850/// Creates a mask from the sign bit of a signed integer.
1851/// Returns all 1s (0xFF...FF) if x < 0, all 0s if x >= 0.
1852fn signMask(comptime T: type, x: T) std.meta.Int(.unsigned, @typeInfo(T).int.bits) {
1853 const type_info = @typeInfo(T);
1854 if (type_info != .int) {
1855 @compileError("signMask requires an integer type");
1856 }
1857
1858 const bits = type_info.int.bits;
1859 const SignedT = std.meta.Int(.signed, bits);
1860
1861 // Convert to signed if needed, arithmetic right shift to propagate sign bit
1862 const x_signed: SignedT = if (type_info.int.signedness == .signed) x else @bitCast(x);
1863 const shifted = x_signed >> (bits - 1);
1864 return @bitCast(shifted);
1865}
1866
1867test "bitMask and signMask helpers" {
1868 try testing.expectEqual(@as(u32, 0x00000000), bitMask(u32, 0));
1869 try testing.expectEqual(@as(u32, 0xFFFFFFFF), bitMask(u32, 1));
1870 try testing.expectEqual(@as(u8, 0x00), bitMask(u8, 0));
1871 try testing.expectEqual(@as(u8, 0xFF), bitMask(u8, 1));
1872 try testing.expectEqual(@as(u64, 0x0000000000000000), bitMask(u64, 0));
1873 try testing.expectEqual(@as(u64, 0xFFFFFFFFFFFFFFFF), bitMask(u64, 1));
1874
1875 try testing.expectEqual(@as(u32, 0xFFFFFFFF), signMask(i32, -1));
1876 try testing.expectEqual(@as(u32, 0xFFFFFFFF), signMask(i32, -100));
1877 try testing.expectEqual(@as(u32, 0x00000000), signMask(i32, 0));
1878 try testing.expectEqual(@as(u32, 0x00000000), signMask(i32, 1));
1879 try testing.expectEqual(@as(u32, 0x00000000), signMask(i32, 100));
1880
1881 try testing.expectEqual(@as(u32, 0xFFFFFFFF), signMask(u32, 0x80000000)); // MSB set
1882 try testing.expectEqual(@as(u32, 0x00000000), signMask(u32, 0x7FFFFFFF)); // MSB clear
1883}
1884
1885/// Montgomery reduction: for input x, returns y where y ≡ x*R^(-1) (mod q).
1886/// This is a generic implementation parameterized by the modulus q, its inverse qInv,
1887/// the Montgomery constant R, and the result bound.
1888///
1889/// For ML-DSA: R = 2^32, returns y < 2q
1890/// For ML-KEM: R = 2^16, returns y in range (-q, q)
1891fn montgomeryReduce(
1892 comptime InT: type,
1893 comptime OutT: type,
1894 comptime q: comptime_int,
1895 comptime qInv: comptime_int,
1896 comptime r_bits: comptime_int,
1897 x: InT,
1898) OutT {
1899 const mask = (@as(InT, 1) << r_bits) - 1;
1900 const m_full = (x *% qInv) & mask;
1901 const m: OutT = @truncate(m_full);
1902
1903 const yR = x -% @as(InT, m) * @as(InT, q);
1904 const y_shifted = @as(std.meta.Int(.unsigned, @typeInfo(InT).Int.bits), @bitCast(yR)) >> r_bits;
1905 return @bitCast(@as(std.meta.Int(.unsigned, @typeInfo(OutT).Int.bits), @truncate(y_shifted)));
1906}
1907
1908/// Uniform sampling using SHAKE-128 with rejection sampling.
1909/// Samples polynomial coefficients uniformly from [0, q) using rejection sampling.
1910///
1911/// Parameters:
1912/// - PolyType: The polynomial type to return
1913/// - q: Modulus
1914/// - bits_per_coef: Number of bits per coefficient (12 or 23)
1915/// - n: Number of coefficients
1916/// - seed: Random seed
1917/// - domain_sep: Domain separation bytes (appended to seed)
1918fn sampleUniformRejection(
1919 comptime PolyType: type,
1920 comptime q: comptime_int,
1921 comptime bits_per_coef: comptime_int,
1922 comptime n: comptime_int,
1923 seed: []const u8,
1924 domain_sep: []const u8,
1925) PolyType {
1926 var h = sha3.Shake128.init(.{});
1927 h.update(seed);
1928 h.update(domain_sep);
1929
1930 const buf_len = sha3.Shake128.block_length; // 168 bytes
1931 var buf: [buf_len]u8 = undefined;
1932
1933 var ret: PolyType = undefined;
1934 var coef_idx: usize = 0;
1935
1936 if (bits_per_coef == 12) {
1937 // ML-KEM path: pack 2 coefficients per 3 bytes (12 bits each)
1938 outer: while (true) {
1939 h.squeeze(&buf);
1940
1941 var j: usize = 0;
1942 while (j < buf_len) : (j += 3) {
1943 const b0 = @as(u16, buf[j]);
1944 const b1 = @as(u16, buf[j + 1]);
1945 const b2 = @as(u16, buf[j + 2]);
1946
1947 const ts: [2]u16 = .{
1948 b0 | ((b1 & 0xf) << 8),
1949 (b1 >> 4) | (b2 << 4),
1950 };
1951
1952 inline for (ts) |t| {
1953 if (t < q) {
1954 ret.cs[coef_idx] = @intCast(t);
1955 coef_idx += 1;
1956 if (coef_idx == n) break :outer;
1957 }
1958 }
1959 }
1960 }
1961 } else if (bits_per_coef == 23) {
1962 // ML-DSA path: 1 coefficient per 3 bytes (23 bits)
1963 while (coef_idx < n) {
1964 h.squeeze(&buf);
1965
1966 var j: usize = 0;
1967 while (j < buf_len and coef_idx < n) : (j += 3) {
1968 const t = (@as(u32, buf[j]) |
1969 (@as(u32, buf[j + 1]) << 8) |
1970 (@as(u32, buf[j + 2]) << 16)) & 0x7fffff;
1971
1972 if (t < q) {
1973 ret.cs[coef_idx] = @intCast(t);
1974 coef_idx += 1;
1975 }
1976 }
1977 }
1978 } else {
1979 @compileError("bits_per_coef must be 12 or 23");
1980 }
1981
1982 return ret;
1983}