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;...@@ -105,19 +105,20 @@ const crypto = std.crypto;
105const errors = std.crypto.errors;105const errors = std.crypto.errors;
106const math = std.math;106const math = std.math;
107const mem = std.mem;107const mem = std.mem;
108const RndGen = std.Random.DefaultPrng;
109const sha3 = crypto.hash.sha3;108const 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
112const Q: i16 = 3329;113const Q: i16 = 3329;
113114
114// Montgomery R115// Montgomery R = 2^16 mod Q (for Montgomery multiplication)
115const R: i32 = 1 << 16;116const R: i32 = 1 << 16;
116117
117// Parameter n, degree of polynomials.118// N is the degree of polynomials (polynomial ring dimension)
118const N: usize = 256;119const 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
121const eta2: u8 = 2;122const eta2: u8 = 2;
122123
123const Params = struct {124const Params = struct {
...@@ -215,7 +216,7 @@ fn Kyber(comptime p: Params) type {...@@ -215,7 +216,7 @@ fn Kyber(comptime p: Params) type {
215 pub const ciphertext_length = Poly.compressedSize(p.du) * p.k + Poly.compressedSize(p.dv);216 pub const ciphertext_length = Poly.compressedSize(p.du) * p.k + Poly.compressedSize(p.dv);
216217
217 const Self = @This();218 const Self = @This();
218 const V = Vec(p.k);219 const V = PolyVec(p.k);
219 const M = Mat(p.k);220 const M = Mat(p.k);
220221
221 /// Length (in bytes) of a shared secret.222 /// Length (in bytes) of a shared secret.
...@@ -241,7 +242,7 @@ fn Kyber(comptime p: Params) type {...@@ -241,7 +242,7 @@ fn Kyber(comptime p: Params) type {
241 hpk: [h_length]u8, // H(pk)242 hpk: [h_length]u8, // H(pk)
242243
243 /// Size of a serialized representation of the key, in bytes.244 /// 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
246 /// Generates a shared secret, and encapsulates it for the public key.247 /// Generates a shared secret, and encapsulates it for the public key.
247 /// If `seed` is `null`, a random seed is used. This is recommended.248 /// If `seed` is `null`, a random seed is used. This is recommended.
...@@ -289,14 +290,14 @@ fn Kyber(comptime p: Params) type {...@@ -289,14 +290,14 @@ fn Kyber(comptime p: Params) type {
289 }290 }
290291
291 /// Serializes the key into a byte array.292 /// 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 {
293 return pk.pk.toBytes();294 return pk.pk.toBytes();
294 }295 }
295296
296 /// Deserializes the key from a byte array.297 /// 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 {
298 var ret: PublicKey = undefined;299 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]);
300 sha3.Sha3_256.hash(buf, &ret.hpk, .{});301 sha3.Sha3_256.hash(buf, &ret.hpk, .{});
301 return ret;302 return ret;
302 }303 }
...@@ -310,8 +311,8 @@ fn Kyber(comptime p: Params) type {...@@ -310,8 +311,8 @@ fn Kyber(comptime p: Params) type {
310 z: [shared_length]u8,311 z: [shared_length]u8,
311312
312 /// Size of a serialized representation of the key, in bytes.313 /// Size of a serialized representation of the key, in bytes.
313 pub const bytes_length: usize =314 pub const encoded_length: usize =
314 InnerSk.bytes_length + InnerPk.bytes_length + h_length + shared_length;315 InnerSk.encoded_length + InnerPk.encoded_length + h_length + shared_length;
315316
316 /// Decapsulates the shared secret within ct using the private key.317 /// Decapsulates the shared secret within ct using the private key.
317 pub fn decaps(sk: SecretKey, ct: *const [ciphertext_length]u8) ![shared_length]u8 {318 pub fn decaps(sk: SecretKey, ct: *const [ciphertext_length]u8) ![shared_length]u8 {
...@@ -346,18 +347,18 @@ fn Kyber(comptime p: Params) type {...@@ -346,18 +347,18 @@ fn Kyber(comptime p: Params) type {
346 }347 }
347348
348 /// Serializes the key into a byte array.349 /// 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 {
350 return sk.sk.toBytes() ++ sk.pk.toBytes() ++ sk.hpk ++ sk.z;351 return sk.sk.toBytes() ++ sk.pk.toBytes() ++ sk.hpk ++ sk.z;
351 }352 }
352353
353 /// Deserializes the key from a byte array.354 /// 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 {
355 var ret: SecretKey = undefined;356 var ret: SecretKey = undefined;
356 comptime var s: usize = 0;357 comptime var s: usize = 0;
357 ret.sk = InnerSk.fromBytes(buf[s .. s + InnerSk.bytes_length]);358 ret.sk = InnerSk.fromBytes(buf[s .. s + InnerSk.encoded_length]);
358 s += InnerSk.bytes_length;359 s += InnerSk.encoded_length;
359 ret.pk = try InnerPk.fromBytes(buf[s .. s + InnerPk.bytes_length]);360 ret.pk = try InnerPk.fromBytes(buf[s .. s + InnerPk.encoded_length]);
360 s += InnerPk.bytes_length;361 s += InnerPk.encoded_length;
361 ret.hpk = buf[s..][0..h_length].*;362 ret.hpk = buf[s..][0..h_length].*;
362 s += h_length;363 s += h_length;
363 ret.z = buf[s..][0..shared_length].*;364 ret.z = buf[s..][0..shared_length].*;
...@@ -418,7 +419,7 @@ fn Kyber(comptime p: Params) type {...@@ -418,7 +419,7 @@ fn Kyber(comptime p: Params) type {
418 // Cached values419 // Cached values
419 aT: M,420 aT: M,
420421
421 const bytes_length = V.bytes_length + 32;422 const encoded_length = V.encoded_length + 32;
422423
423 fn encrypt(424 fn encrypt(
424 pk: InnerPk,425 pk: InnerPk,
...@@ -436,7 +437,7 @@ fn Kyber(comptime p: Params) type {...@@ -436,7 +437,7 @@ fn Kyber(comptime p: Params) type {
436 // Note that coefficients of r are bounded by q and those of Aᵀ437 // Note that coefficients of r are bounded by q and those of Aᵀ
437 // are bounded by 4.5q and so their product is bounded by 2¹⁵q438 // are bounded by 4.5q and so their product is bounded by 2¹⁵q
438 // as required for multiplication.439 // as required for multiplication.
439 u.ps[i] = pk.aT.vs[i].dotHat(rh);440 u.ps[i] = pk.aT.rows[i].dotHat(rh);
440 }441 }
441442
442 // Aᵀ and r were not in Montgomery form, so the Montgomery443 // Aᵀ and r were not in Montgomery form, so the Montgomery
...@@ -451,14 +452,14 @@ fn Kyber(comptime p: Params) type {...@@ -451,14 +452,14 @@ fn Kyber(comptime p: Params) type {
451 return u.compress(p.du) ++ v.compress(p.dv);452 return u.compress(p.du) ++ v.compress(p.dv);
452 }453 }
453454
454 fn toBytes(pk: InnerPk) [bytes_length]u8 {455 fn toBytes(pk: InnerPk) [encoded_length]u8 {
455 return pk.th.toBytes() ++ pk.rho;456 return pk.th.toBytes() ++ pk.rho;
456 }457 }
457458
458 fn fromBytes(buf: *const [bytes_length]u8) errors.NonCanonicalError!InnerPk {459 fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!InnerPk {
459 var ret: InnerPk = undefined;460 var ret: InnerPk = undefined;
460461
461 const th_bytes = buf[0..V.bytes_length];462 const th_bytes = buf[0..V.encoded_length];
462 ret.th = V.fromBytes(th_bytes).normalize();463 ret.th = V.fromBytes(th_bytes).normalize();
463464
464 if (p.ml_kem) {465 if (p.ml_kem) {
...@@ -468,7 +469,7 @@ fn Kyber(comptime p: Params) type {...@@ -468,7 +469,7 @@ fn Kyber(comptime p: Params) type {
468 }469 }
469 }470 }
470471
471 ret.rho = buf[V.bytes_length..bytes_length].*;472 ret.rho = buf[V.encoded_length..encoded_length].*;
472 ret.aT = M.uniform(ret.rho, true);473 ret.aT = M.uniform(ret.rho, true);
473 return ret;474 return ret;
474 }475 }
...@@ -477,7 +478,7 @@ fn Kyber(comptime p: Params) type {...@@ -477,7 +478,7 @@ fn Kyber(comptime p: Params) type {
477 // Private key of the inner PKE478 // Private key of the inner PKE
478 const InnerSk = struct {479 const InnerSk = struct {
479 sh: V, // NTT(s), normalized480 sh: V, // NTT(s), normalized
480 const bytes_length = V.bytes_length;481 const encoded_length = V.encoded_length;
481482
482 fn decrypt(sk: InnerSk, ct: *const [ciphertext_length]u8) [inner_plaintext_length]u8 {483 fn decrypt(sk: InnerSk, ct: *const [ciphertext_length]u8) [inner_plaintext_length]u8 {
483 const u = V.decompress(p.du, ct[0..comptime V.compressedSize(p.du)]);484 const u = V.decompress(p.du, ct[0..comptime V.compressedSize(p.du)]);
...@@ -491,11 +492,11 @@ fn Kyber(comptime p: Params) type {...@@ -491,11 +492,11 @@ fn Kyber(comptime p: Params) type {
491 .normalize().compress(1);492 .normalize().compress(1);
492 }493 }
493494
494 fn toBytes(sk: InnerSk) [bytes_length]u8 {495 fn toBytes(sk: InnerSk) [encoded_length]u8 {
495 return sk.sh.toBytes();496 return sk.sh.toBytes();
496 }497 }
497498
498 fn fromBytes(buf: *const [bytes_length]u8) InnerSk {499 fn fromBytes(buf: *const [encoded_length]u8) InnerSk {
499 var ret: InnerSk = undefined;500 var ret: InnerSk = undefined;
500 ret.sh = V.fromBytes(buf).normalize();501 ret.sh = V.fromBytes(buf).normalize();
501 return ret;502 return ret;
...@@ -516,7 +517,7 @@ fn Kyber(comptime p: Params) type {...@@ -516,7 +517,7 @@ fn Kyber(comptime p: Params) type {
516 // Sample secret vector s.517 // Sample secret vector s.
517 sk.sh = V.noise(p.eta1, 0, sigma).ntt().normalize();518 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.
520 var th: V = undefined;521 var th: V = undefined;
521522
522 // Next, we compute t = A s + e.523 // Next, we compute t = A s + e.
...@@ -528,7 +529,7 @@ fn Kyber(comptime p: Params) type {...@@ -528,7 +529,7 @@ fn Kyber(comptime p: Params) type {
528 // multiplications in the inner product added a factor R⁻¹ which529 // multiplications in the inner product added a factor R⁻¹ which
529 // we'll cancel out with toMont(). This will also ensure the530 // we'll cancel out with toMont(). This will also ensure the
530 // coefficients of th are bounded in absolute value by q.531 // 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();
532 }533 }
533534
534 pk.th = th.add(eh).normalize(); // bounded by 8q535 pk.th = th.add(eh).normalize(); // bounded by 8q
...@@ -565,7 +566,6 @@ const zetas = computeZetas();...@@ -565,7 +566,6 @@ const zetas = computeZetas();
565// not enough, the other coefficient is reduced as well.566// not enough, the other coefficient is reduced as well.
566//567//
567// This is actually optimal, as proven in https://eprint.iacr.org/2020/1377.pdf568// This is actually optimal, as proven in https://eprint.iacr.org/2020/1377.pdf
568// TODO generate comptime?
569const inv_ntt_reductions = [_]i16{569const inv_ntt_reductions = [_]i16{
570 -1, // after layer 1570 -1, // after layer 1
571 -1, // after layer 2571 -1, // after layer 2
...@@ -634,31 +634,8 @@ test "invNTTReductions bounds" {...@@ -634,31 +634,8 @@ test "invNTTReductions bounds" {
634 }634 }
635}635}
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.
660fn invertMod(a: anytype, p: @TypeOf(a)) @TypeOf(a) {637fn invertMod(a: anytype, p: @TypeOf(a)) @TypeOf(a) {
661 const r = eea(a, p);638 const r = extendedEuclidean(@TypeOf(a), a, p);
662 assert(r.gcd == 1);639 assert(r.gcd == 1);
663 return r.x;640 return r.x;
664}641}
...@@ -788,31 +765,12 @@ test "Test csubq" {...@@ -788,31 +765,12 @@ test "Test csubq" {
788 }765 }
789}766}
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
810// Computes zetas table used by ntt and invNTT.768// Computes zetas table used by ntt and invNTT.
811fn computeZetas() [128]i16 {769fn computeZetas() [128]i16 {
812 @setEvalBranchQuota(10000);770 @setEvalBranchQuota(10000);
813 var ret: [128]i16 = undefined;771 var ret: [128]i16 = undefined;
814 for (&ret, 0..) |*r, i| {772 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)));
816 r.* = csubq(feBarrettReduce(feToMont(t)));774 r.* = csubq(feBarrettReduce(feToMont(t)));
817 }775 }
818 return ret;776 return ret;
...@@ -828,9 +786,10 @@ fn computeZetas() [128]i16 {...@@ -828,9 +786,10 @@ fn computeZetas() [128]i16 {
828const Poly = struct {786const Poly = struct {
829 cs: [N]i16,787 cs: [N]i16,
830788
831 const bytes_length = N / 2 * 3;789 const encoded_length = N / 2 * 3;
832 const zero: Poly = .{ .cs = .{0} ** N };790 const zero: Poly = .{ .cs = .{0} ** N };
833791
792 // Add two polynomials (coefficients not normalized)
834 fn add(a: Poly, b: Poly) Poly {793 fn add(a: Poly, b: Poly) Poly {
835 var ret: Poly = undefined;794 var ret: Poly = undefined;
836 for (0..N) |i| {795 for (0..N) |i| {
...@@ -839,6 +798,7 @@ const Poly = struct {...@@ -839,6 +798,7 @@ const Poly = struct {
839 return ret;798 return ret;
840 }799 }
841800
801 // Subtract two polynomials (coefficients not normalized)
842 fn sub(a: Poly, b: Poly) Poly {802 fn sub(a: Poly, b: Poly) Poly {
843 var ret: Poly = undefined;803 var ret: Poly = undefined;
844 for (0..N) |i| {804 for (0..N) |i| {
...@@ -847,25 +807,6 @@ const Poly = struct {...@@ -847,25 +807,6 @@ const Poly = struct {
847 return ret;807 return ret;
848 }808 }
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
869 // Executes a forward "NTT" on p.810 // Executes a forward "NTT" on p.
870 //811 //
871 // Assumes the coefficients are in absolute value ≤q. The resulting812 // Assumes the coefficients are in absolute value ≤q. The resulting
...@@ -1054,7 +995,7 @@ const Poly = struct {...@@ -1054,7 +995,7 @@ const Poly = struct {
1054 var in_off: usize = 0;995 var in_off: usize = 0;
1055 var out_off: usize = 0;996 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);
1058 const in_batch_size: usize = comptime batch_size / d;999 const in_batch_size: usize = comptime batch_size / d;
1059 const out_batch_size: usize = comptime batch_size / 8;1000 const out_batch_size: usize = comptime batch_size / 8;
10601001
...@@ -1118,7 +1059,7 @@ const Poly = struct {...@@ -1118,7 +1059,7 @@ const Poly = struct {
1118 var in_off: usize = 0;1059 var in_off: usize = 0;
1119 var out_off: usize = 0;1060 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);
1122 const in_batch_size: usize = comptime batch_size / 8;1063 const in_batch_size: usize = comptime batch_size / 8;
1123 const out_batch_size: usize = comptime batch_size / d;1064 const out_batch_size: usize = comptime batch_size / d;
11241065
...@@ -1275,53 +1216,23 @@ const Poly = struct {...@@ -1275,53 +1216,23 @@ const Poly = struct {
1275 return ret;1216 return ret;
1276 }1217 }
12771218
1278 // Sample p uniformly from the given seed and x and y coordinates.
1279 fn uniform(seed: [32]u8, x: u8, y: u8) Poly {1219 fn uniform(seed: [32]u8, x: u8, y: u8) Poly {
1280 var h = sha3.Shake128.init(.{});1220 const domain_sep: [2]u8 = .{ x, y };
1281 const suffix: [2]u8 = .{ x, y };1221 return sampleUniformRejection(
1282 h.update(&seed);1222 Poly,
1283 h.update(&suffix);1223 Q,
12841224 12,
1285 const buf_len = sha3.Shake128.block_length; // rate SHAKE-1281225 N,
1286 var buf: [buf_len]u8 = undefined;1226 &seed,
12871227 &domain_sep,
1288 var ret: Poly = undefined;1228 );
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;
1318 }1229 }
13191230
1320 // Packs p.1231 // Packs p.
1321 //1232 //
1322 // Assumes p is normalized (and not just Barrett reduced).1233 // Assumes p is normalized (and not just Barrett reduced).
1323 fn toBytes(p: Poly) [bytes_length]u8 {1234 fn toBytes(p: Poly) [encoded_length]u8 {
1324 var ret: [bytes_length]u8 = undefined;1235 var ret: [encoded_length]u8 = undefined;
1325 for (0..comptime N / 2) |i| {1236 for (0..comptime N / 2) |i| {
1326 const t0 = @as(u16, @intCast(p.cs[2 * i]));1237 const t0 = @as(u16, @intCast(p.cs[2 * i]));
1327 const t1 = @as(u16, @intCast(p.cs[2 * i + 1]));1238 const t1 = @as(u16, @intCast(p.cs[2 * i + 1]));
...@@ -1335,7 +1246,7 @@ const Poly = struct {...@@ -1335,7 +1246,7 @@ const Poly = struct {
1335 // Unpacks a Poly from buf.1246 // Unpacks a Poly from buf.
1336 //1247 //
1337 // p will not be normalized; instead 0 ≤ p[i] < 4096.1248 // 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 {
1339 var ret: Poly = undefined;1250 var ret: Poly = undefined;
1340 for (0..comptime N / 2) |i| {1251 for (0..comptime N / 2) |i| {
1341 const b0 = @as(i16, buf[3 * i]);1252 const b0 = @as(i16, buf[3 * i]);
...@@ -1348,71 +1259,65 @@ const Poly = struct {...@@ -1348,71 +1259,65 @@ const Poly = struct {
1348 }1259 }
1349};1260};
13501261
1351// A vector of K polynomials.1262// A vector of k polynomials.
1352fn Vec(comptime K: u8) type {1263fn PolyVec(comptime k: u8) type {
1353 return struct {1264 return struct {
1354 ps: [K]Poly,1265 ps: [k]Poly,
13551266
1356 const Self = @This();1267 const Self = @This();
1357 const bytes_length = K * Poly.bytes_length;1268 const encoded_length = k * Poly.encoded_length;
13581269
1359 fn compressedSize(comptime d: u8) usize {1270 fn compressedSize(comptime d: u8) usize {
1360 return Poly.compressedSize(d) * K;1271 return Poly.compressedSize(d) * k;
1361 }1272 }
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 {
1364 var ret: Self = undefined;1276 var ret: Self = undefined;
1365 for (0..K) |i| {1277 inline for (0..k) |i| {
1366 ret.ps[i] = a.ps[i].ntt();1278 ret.ps[i] = op(v.ps[i]);
1367 }1279 }
1368 return ret;1280 return ret;
1369 }1281 }
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 {
1372 var ret: Self = undefined;1285 var ret: Self = undefined;
1373 for (0..K) |i| {1286 inline for (0..k) |i| {
1374 ret.ps[i] = a.ps[i].invNTT();1287 ret.ps[i] = op(a.ps[i], b.ps[i]);
1375 }1288 }
1376 return ret;1289 return ret;
1377 }1290 }
13781291
1379 fn normalize(a: Self) Self {1292 fn ntt(v: Self) Self {
1380 var ret: Self = undefined;1293 return map(v, Poly.ntt);
1381 for (0..K) |i| {
1382 ret.ps[i] = a.ps[i].normalize();
1383 }
1384 return ret;
1385 }1294 }
13861295
1387 fn barrettReduce(a: Self) Self {1296 fn invNTT(v: Self) Self {
1388 var ret: Self = undefined;1297 return map(v, Poly.invNTT);
1389 for (0..K) |i| {1298 }
1390 ret.ps[i] = a.ps[i].barrettReduce();1299
1391 }1300 fn normalize(v: Self) Self {
1392 return ret;1301 return map(v, Poly.normalize);
1302 }
1303
1304 fn barrettReduce(v: Self) Self {
1305 return map(v, Poly.barrettReduce);
1393 }1306 }
13941307
1395 fn add(a: Self, b: Self) Self {1308 fn add(a: Self, b: Self) Self {
1396 var ret: Self = undefined;1309 return mapBinary(a, b, Poly.add);
1397 for (0..K) |i| {
1398 ret.ps[i] = a.ps[i].add(b.ps[i]);
1399 }
1400 return ret;
1401 }1310 }
14021311
1403 fn sub(a: Self, b: Self) Self {1312 fn sub(a: Self, b: Self) Self {
1404 var ret: Self = undefined;1313 return mapBinary(a, b, Poly.sub);
1405 for (0..K) |i| {
1406 ret.ps[i] = a.ps[i].sub(b.ps[i]);
1407 }
1408 return ret;
1409 }1314 }
14101315
1411 // Samples v[i] from centered binomial distribution with the given η,1316 // Samples v[i] from centered binomial distribution with the given η,
1412 // seed and nonce+i.1317 // seed and nonce+i.
1413 fn noise(comptime eta: u8, nonce: u8, seed: *const [32]u8) Self {1318 fn noise(comptime eta: u8, nonce: u8, seed: *const [32]u8) Self {
1414 var ret: Self = undefined;1319 var ret: Self = undefined;
1415 for (0..K) |i| {1320 for (0..k) |i| {
1416 ret.ps[i] = Poly.noise(eta, nonce + @as(u8, @intCast(i)), seed);1321 ret.ps[i] = Poly.noise(eta, nonce + @as(u8, @intCast(i)), seed);
1417 }1322 }
1418 return ret;1323 return ret;
...@@ -1428,7 +1333,7 @@ fn Vec(comptime K: u8) type {...@@ -1428,7 +1333,7 @@ fn Vec(comptime K: u8) type {
1428 // of the Montgomery factor.1333 // of the Montgomery factor.
1429 fn dotHat(a: Self, b: Self) Poly {1334 fn dotHat(a: Self, b: Self) Poly {
1430 var ret: Poly = Poly.zero;1335 var ret: Poly = Poly.zero;
1431 for (0..K) |i| {1336 for (0..k) |i| {
1432 ret = ret.add(a.ps[i].mulHat(b.ps[i]));1337 ret = ret.add(a.ps[i].mulHat(b.ps[i]));
1433 }1338 }
1434 return ret;1339 return ret;
...@@ -1437,7 +1342,7 @@ fn Vec(comptime K: u8) type {...@@ -1437,7 +1342,7 @@ fn Vec(comptime K: u8) type {
1437 fn compress(v: Self, comptime d: u8) [compressedSize(d)]u8 {1342 fn compress(v: Self, comptime d: u8) [compressedSize(d)]u8 {
1438 const cs = comptime Poly.compressedSize(d);1343 const cs = comptime Poly.compressedSize(d);
1439 var ret: [compressedSize(d)]u8 = undefined;1344 var ret: [compressedSize(d)]u8 = undefined;
1440 inline for (0..K) |i| {1345 inline for (0..k) |i| {
1441 ret[i * cs .. (i + 1) * cs].* = v.ps[i].compress(d);1346 ret[i * cs .. (i + 1) * cs].* = v.ps[i].compress(d);
1442 }1347 }
1443 return ret;1348 return ret;
...@@ -1446,27 +1351,27 @@ fn Vec(comptime K: u8) type {...@@ -1446,27 +1351,27 @@ fn Vec(comptime K: u8) type {
1446 fn decompress(comptime d: u8, buf: *const [compressedSize(d)]u8) Self {1351 fn decompress(comptime d: u8, buf: *const [compressedSize(d)]u8) Self {
1447 const cs = comptime Poly.compressedSize(d);1352 const cs = comptime Poly.compressedSize(d);
1448 var ret: Self = undefined;1353 var ret: Self = undefined;
1449 inline for (0..K) |i| {1354 inline for (0..k) |i| {
1450 ret.ps[i] = Poly.decompress(d, buf[i * cs .. (i + 1) * cs]);1355 ret.ps[i] = Poly.decompress(d, buf[i * cs .. (i + 1) * cs]);
1451 }1356 }
1452 return ret;1357 return ret;
1453 }1358 }
14541359
1455 /// Serializes the key into a byte array.1360 /// Serializes the key into a byte array.
1456 fn toBytes(v: Self) [bytes_length]u8 {1361 fn toBytes(v: Self) [encoded_length]u8 {
1457 var ret: [bytes_length]u8 = undefined;1362 var ret: [encoded_length]u8 = undefined;
1458 inline for (0..K) |i| {1363 inline for (0..k) |i| {
1459 ret[i * Poly.bytes_length .. (i + 1) * Poly.bytes_length].* = v.ps[i].toBytes();1364 ret[i * Poly.encoded_length .. (i + 1) * Poly.encoded_length].* = v.ps[i].toBytes();
1460 }1365 }
1461 return ret;1366 return ret;
1462 }1367 }
14631368
1464 /// Deserializes the key from a byte array.1369 /// 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 {
1466 var ret: Self = undefined;1371 var ret: Self = undefined;
1467 inline for (0..K) |i| {1372 inline for (0..k) |i| {
1468 ret.ps[i] = Poly.fromBytes(1373 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],
1470 );1375 );
1471 }1376 }
1472 return ret;1377 return ret;
...@@ -1474,19 +1379,19 @@ fn Vec(comptime K: u8) type {...@@ -1474,19 +1379,19 @@ fn Vec(comptime K: u8) type {
1474 };1379 };
1475}1380}
14761381
1477// A matrix of K vectors1382// A matrix of k vectors
1478fn Mat(comptime K: u8) type {1383fn Mat(comptime k: u8) type {
1479 return struct {1384 return struct {
1480 const Self = @This();1385 const Self = @This();
1481 vs: [K]Vec(K),1386 rows: [k]PolyVec(k),
14821387
1483 fn uniform(seed: [32]u8, comptime transposed: bool) Self {1388 fn uniform(seed: [32]u8, comptime transposed: bool) Self {
1484 var ret: Self = undefined;1389 var ret: Self = undefined;
1485 var i: u8 = 0;1390 var i: u8 = 0;
1486 while (i < K) : (i += 1) {1391 while (i < k) : (i += 1) {
1487 var j: u8 = 0;1392 var j: u8 = 0;
1488 while (j < K) : (j += 1) {1393 while (j < k) : (j += 1) {
1489 ret.vs[i].ps[j] = Poly.uniform(1394 ret.rows[i].ps[j] = Poly.uniform(
1490 seed,1395 seed,
1491 if (transposed) i else j,1396 if (transposed) i else j,
1492 if (transposed) j else i,1397 if (transposed) j else i,
...@@ -1499,9 +1404,9 @@ fn Mat(comptime K: u8) type {...@@ -1499,9 +1404,9 @@ fn Mat(comptime K: u8) type {
1499 // Returns transpose of A1404 // Returns transpose of A
1500 fn transpose(m: Self) Self {1405 fn transpose(m: Self) Self {
1501 var ret: Self = undefined;1406 var ret: Self = undefined;
1502 for (0..K) |i| {1407 for (0..k) |i| {
1503 for (0..K) |j| {1408 for (0..k) |j| {
1504 ret.vs[i].ps[j] = m.vs[j].ps[i];1409 ret.rows[i].ps[j] = m.rows[j].ps[i];
1505 }1410 }
1506 }1411 }
1507 return ret;1412 return ret;
...@@ -1522,12 +1427,30 @@ fn cmov(comptime len: usize, dst: *[len]u8, src: [len]u8, b: u1) void {...@@ -1522,12 +1427,30 @@ fn cmov(comptime len: usize, dst: *[len]u8, src: [len]u8, b: u1) void {
1522 }1427 }
1523}1428}
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
1525test "MulHat" {1448test "MulHat" {
1526 var rnd = RndGen.init(0);1449 var rnd = RndGen.init(0);
15271450
1528 for (0..100) |_| {1451 for (0..100) |_| {
1529 const a = Poly.randAbsLeqQ(&rnd);1452 const a = randPolyAbsLeqQ(&rnd);
1530 const b = Poly.randAbsLeqQ(&rnd);1453 const b = randPolyAbsLeqQ(&rnd);
15311454
1532 const p2 = a.ntt().mulHat(b.ntt()).barrettReduce().invNTT().normalize();1455 const p2 = a.ntt().mulHat(b.ntt()).barrettReduce().invNTT().normalize();
1533 var p: Poly = undefined;1456 var p: Poly = undefined;
...@@ -1557,7 +1480,7 @@ test "NTT" {...@@ -1557,7 +1480,7 @@ test "NTT" {
1557 var rnd = RndGen.init(0);1480 var rnd = RndGen.init(0);
15581481
1559 for (0..1000) |_| {1482 for (0..1000) |_| {
1560 var p = Poly.randAbsLeqQ(&rnd);1483 var p = randPolyAbsLeqQ(&rnd);
1561 const q = p.toMont().normalize();1484 const q = p.toMont().normalize();
1562 p = p.ntt();1485 p = p.ntt();
15631486
...@@ -1580,7 +1503,7 @@ test "Compression" {...@@ -1580,7 +1503,7 @@ test "Compression" {
1580 var rnd = RndGen.init(0);1503 var rnd = RndGen.init(0);
1581 inline for (.{ 1, 4, 5, 10, 11 }) |d| {1504 inline for (.{ 1, 4, 5, 10, 11 }) |d| {
1582 for (0..1000) |_| {1505 for (0..1000) |_| {
1583 const p = Poly.randNormalized(&rnd);1506 const p = randPolyNormalized(&rnd);
1584 const pp = p.compress(d);1507 const pp = p.compress(d);
1585 const pq = Poly.decompress(d, &pp).compress(d);1508 const pq = Poly.decompress(d, &pp).compress(d);
1586 try testing.expectEqual(pp, pq);1509 try testing.expectEqual(pp, pq);
...@@ -1671,7 +1594,7 @@ test "Polynomial packing" {...@@ -1671,7 +1594,7 @@ test "Polynomial packing" {
1671 var rnd = RndGen.init(0);1594 var rnd = RndGen.init(0);
16721595
1673 for (0..1000) |_| {1596 for (0..1000) |_| {
1674 const p = Poly.randNormalized(&rnd);1597 const p = randPolyNormalized(&rnd);
1675 try testing.expectEqual(Poly.fromBytes(&p.toBytes()), p);1598 try testing.expectEqual(Poly.fromBytes(&p.toBytes()), p);
1676 }1599 }
1677}1600}
...@@ -1839,3 +1762,222 @@ const NistDRBG = struct {...@@ -1839,3 +1762,222 @@ const NistDRBG = struct {
1839 return ret;1762 return ret;
1840 }1763 }
1841};1764};
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}