1//! Implementation of the IND-CCA2 post-quantum secure key encapsulation mechanism (KEM)
2//! ML-KEM (NIST FIPS-203 publication) and CRYSTALS-Kyber (v3.02/"draft00" CFRG draft).
3//!
4//! The namespace `d00` refers to the version currently implemented, in accordance with the CFRG draft.
5//! The `nist` namespace refers to the FIPS-203 publication.
6//!
7//! Quoting from the CFRG I-D:
8//!
9//! Kyber is not a Diffie-Hellman (DH) style non-interactive key
10//! agreement, but instead, Kyber is a Key Encapsulation Method (KEM).
11//! In essence, a KEM is a Public-Key Encryption (PKE) scheme where the
12//! plaintext cannot be specified, but is generated as a random key as
13//! part of the encryption. A KEM can be transformed into an unrestricted
14//! PKE using HPKE (RFC9180). On its own, a KEM can be used as a key
15//! agreement method in TLS.
16//!
17//! Kyber is an IND-CCA2 secure KEM. It is constructed by applying a
18//! Fujisaki--Okamato style transformation on InnerPKE, which is the
19//! underlying IND-CPA secure Public Key Encryption scheme. We cannot
20//! use InnerPKE directly, as its ciphertexts are malleable.
21//!
22//! ```
23//! F.O. transform
24//! InnerPKE ----------------------> Kyber
25//! IND-CPA IND-CCA2
26//! ```
27//!
28//! Kyber is a lattice-based scheme. More precisely, its security is
29//! based on the learning-with-errors-and-rounding problem in module
30//! lattices (MLWER). The underlying polynomial ring R (defined in
31//! Section 5) is chosen such that multiplication is very fast using the
32//! number theoretic transform (NTT, see Section 5.1.3).
33//!
34//! An InnerPKE private key is a vector _s_ over R of length k which is
35//! _small_ in a particular way. Here k is a security parameter akin to
36//! the size of a prime modulus. For Kyber512, which targets AES-128's
37//! security level, the value of k is 2.
38//!
39//! The public key consists of two values:
40//!
41//! * _A_ a uniformly sampled k by k matrix over R _and_
42//!
43//! * _t = A s + e_, where e is a suitably small masking vector.
44//!
45//! Distinguishing between such A s + e and a uniformly sampled t is the
46//! module learning-with-errors (MLWE) problem. If that is hard, then it
47//! is also hard to recover the private key from the public key as that
48//! would allow you to distinguish between those two.
49//!
50//! To save space in the public key, A is recomputed deterministically
51//! from a seed _rho_.
52//!
53//! A ciphertext for a message m under this public key is a pair (c_1,
54//! c_2) computed roughly as follows:
55//!
56//! c_1 = Compress(A^T r + e_1, d_u)
57//! c_2 = Compress(t^T r + e_2 + Decompress(m, 1), d_v)
58//!
59//! where
60//!
61//! * e_1, e_2 and r are small blinds;
62//!
63//! * Compress(-, d) removes some information, leaving d bits per
64//! coefficient and Decompress is such that Compress after Decompress
65//! does nothing and
66//!
67//! * d_u, d_v are scheme parameters.
68//!
69//! Distinguishing such a ciphertext and uniformly sampled (c_1, c_2) is
70//! an example of the full MLWER problem, see section 4.4 of [KyberV302].
71//!
72//! To decrypt the ciphertext, one computes
73//!
74//! m = Compress(Decompress(c_2, d_v) - s^T Decompress(c_1, d_u), 1).
75//!
76//! It it not straight-forward to see that this formula is correct. In
77//! fact, there is negligible but non-zero probability that a ciphertext
78//! does not decrypt correctly given by the DFP column in Table 4. This
79//! failure probability can be computed by a careful automated analysis
80//! of the probabilities involved, see kyber_failure.py of [SecEst].
81//!
82//! [KyberV302](https://pq-crystals.org/kyber/data/kyber-specification-round3-20210804.pdf)
83//! [I-D](https://github.com/bwesterb/draft-schwabe-cfrg-kyber)
84//! [SecEst](https://github.com/pq-crystals/security-estimates)
85
86// TODO
87//
88// - The bottleneck in Kyber are the various hash/xof calls:
89// - Optimize Zig's keccak implementation.
90// - Use SIMD to compute keccak in parallel.
91// - Can we track bounds of coefficients using comptime types without
92// duplicating code?
93// - Would be neater to have tests closer to the thing under test.
94// - When generating a keypair, we have a copy of the inner public key with
95// its large matrix A in both the public key and the private key. In Go we
96// can just have a pointer in the private key to the public key, but
97// how do we do this elegantly in Zig?
98
99const std = @import("std");
100const builtin = @import("builtin");
101
102const testing = std.testing;
103const assert = std.debug.assert;
104const crypto = std.crypto;
105const errors = std.crypto.errors;
106const math = std.math;
107const mem = std.mem;
108const sha3 = crypto.hash.sha3;
109
110const RndGen = std.Random.DefaultPrng;
111
112// Q is the modulus q ≡ 3329 = 2¹¹ + 2¹⁰ + 2⁸ + 1
113const Q: i16 = 3329;
114
115// Montgomery R = 2^16 mod Q (for Montgomery multiplication)
116const R: i32 = 1 << 16;
117
118// N is the degree of polynomials (polynomial ring dimension)
119const N: usize = 256;
120
121// eta2 is the size of "small" vectors used in encryption blinds
122const eta2: u8 = 2;
123
124const Params = struct {
125 name: []const u8,
126
127 // NIST ML-KEM variant instead of Kyber as originally submitted.
128 ml_kem: bool = false,
129
130 // Width and height of the matrix A.
131 k: u8,
132
133 // Size of "small" vectors used in private key and encryption blinds.
134 eta1: u8,
135
136 // How many bits to retain of u, the private-key independent part
137 // of the ciphertext.
138 du: u8,
139
140 // How many bits to retain of v, the private-key dependent part
141 // of the ciphertext.
142 dv: u8,
143};
144
145pub const d00 = struct {
146 pub const Kyber512 = Kyber(.{
147 .name = "Kyber512",
148 .k = 2,
149 .eta1 = 3,
150 .du = 10,
151 .dv = 4,
152 });
153
154 pub const Kyber768 = Kyber(.{
155 .name = "Kyber768",
156 .k = 3,
157 .eta1 = 2,
158 .du = 10,
159 .dv = 4,
160 });
161
162 pub const Kyber1024 = Kyber(.{
163 .name = "Kyber1024",
164 .k = 4,
165 .eta1 = 2,
166 .du = 11,
167 .dv = 5,
168 });
169};
170
171pub const nist = struct {
172 pub const MLKem512 = Kyber(.{
173 .name = "ML-KEM-512",
174 .ml_kem = true,
175 .k = 2,
176 .eta1 = 3,
177 .du = 10,
178 .dv = 4,
179 });
180
181 pub const MLKem768 = Kyber(.{
182 .name = "ML-KEM-768",
183 .ml_kem = true,
184 .k = 3,
185 .eta1 = 2,
186 .du = 10,
187 .dv = 4,
188 });
189
190 pub const MLKem1024 = Kyber(.{
191 .name = "ML-KEM-1024",
192 .ml_kem = true,
193 .k = 4,
194 .eta1 = 2,
195 .du = 11,
196 .dv = 5,
197 });
198};
199
200const modes = [_]type{
201 d00.Kyber512,
202 d00.Kyber768,
203 d00.Kyber1024,
204 nist.MLKem512,
205 nist.MLKem768,
206 nist.MLKem1024,
207};
208const h_length: usize = 32;
209const inner_seed_length: usize = 32;
210const common_encaps_seed_length: usize = 32;
211const common_shared_key_size: usize = 32;
212
213fn Kyber(comptime p: Params) type {
214 return struct {
215 // Size of a ciphertext, in bytes.
216 pub const ciphertext_length = Poly.compressedSize(p.du) * p.k + Poly.compressedSize(p.dv);
217
218 const Self = @This();
219 const V = PolyVec(p.k);
220 const M = Mat(p.k);
221
222 /// Length (in bytes) of a shared secret.
223 pub const shared_length = common_shared_key_size;
224 /// Length (in bytes) of a seed for deterministic encapsulation.
225 pub const encaps_seed_length = common_encaps_seed_length;
226 /// Length (in bytes) of a seed for key generation.
227 pub const seed_length: usize = inner_seed_length + shared_length;
228 /// Algorithm name.
229 pub const name = p.name;
230
231 /// A shared secret, and an encapsulated (encrypted) representation of it.
232 pub const EncapsulatedSecret = struct {
233 shared_secret: [shared_length]u8,
234 ciphertext: [ciphertext_length]u8,
235 };
236
237 /// A Kyber public key.
238 pub const PublicKey = struct {
239 pk: InnerPk,
240
241 // Cached
242 hpk: [h_length]u8, // H(pk)
243
244 /// Size of a serialized representation of the key, in bytes.
245 pub const encoded_length = InnerPk.encoded_length;
246
247 /// Generates a shared secret, encapsulated for the public key,
248 /// using random bytes.
249 ///
250 /// This is recommended over `encapsDeterministic`.
251 pub fn encaps(pk: PublicKey, io: std.Io) EncapsulatedSecret {
252 var m: [inner_plaintext_length]u8 = undefined;
253 io.random(&m);
254 return encapsInner(pk, &m);
255 }
256
257 /// Generates a shared secret, encapsulated for the public key,
258 /// using the provided seed.
259 ///
260 /// Calling `encaps` instead is recommended.
261 pub fn encapsDeterministic(pk: PublicKey, seed: *const [encaps_seed_length]u8) EncapsulatedSecret {
262 var m: [inner_plaintext_length]u8 = undefined;
263 if (p.ml_kem) {
264 @memcpy(&m, seed);
265 } else {
266 // m = H(seed)
267 sha3.Sha3_256.hash(seed, &m, .{});
268 }
269 return encapsInner(pk, &m);
270 }
271
272 fn encapsInner(pk: PublicKey, m: *[inner_plaintext_length]u8) EncapsulatedSecret {
273 // (K', r) = G(m ‖ H(pk))
274 var kr: [inner_plaintext_length + h_length]u8 = undefined;
275 var g = sha3.Sha3_512.init(.{});
276 g.update(m);
277 g.update(&pk.hpk);
278 g.final(&kr);
279
280 // c = innerEncrypt(pk, m, r)
281 const ct = pk.pk.encrypt(m, kr[32..64]);
282
283 if (p.ml_kem) {
284 return EncapsulatedSecret{
285 .shared_secret = kr[0..shared_length].*, // ML-KEM: K = K'
286 .ciphertext = ct,
287 };
288 } else {
289 // Compute H(c) and put in second slot of kr, which will be (K', H(c)).
290 sha3.Sha3_256.hash(&ct, kr[32..], .{});
291
292 var ss: [shared_length]u8 = undefined;
293 sha3.Shake256.hash(&kr, &ss, .{});
294 return EncapsulatedSecret{
295 .shared_secret = ss, // Kyber: K = KDF(K' ‖ H(c))
296 .ciphertext = ct,
297 };
298 }
299 }
300
301 /// Serializes the key into a byte array.
302 pub fn toBytes(pk: PublicKey) [encoded_length]u8 {
303 return pk.pk.toBytes();
304 }
305
306 /// Deserializes the key from a byte array.
307 pub fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!PublicKey {
308 var ret: PublicKey = undefined;
309 ret.pk = try InnerPk.fromBytes(buf[0..InnerPk.encoded_length]);
310 sha3.Sha3_256.hash(buf, &ret.hpk, .{});
311 return ret;
312 }
313 };
314
315 /// A Kyber secret key.
316 pub const SecretKey = struct {
317 sk: InnerSk,
318 pk: InnerPk,
319 hpk: [h_length]u8, // H(pk)
320 z: [shared_length]u8,
321
322 /// Size of a serialized representation of the key, in bytes.
323 pub const encoded_length: usize =
324 InnerSk.encoded_length + InnerPk.encoded_length + h_length + shared_length;
325
326 /// Decapsulates the shared secret within ct using the private key.
327 pub fn decaps(sk: SecretKey, ct: *const [ciphertext_length]u8) ![shared_length]u8 {
328 // m' = innerDec(ct)
329 const m2 = sk.sk.decrypt(ct);
330
331 // (K'', r') = G(m' ‖ H(pk))
332 var kr2: [64]u8 = undefined;
333 var g = sha3.Sha3_512.init(.{});
334 g.update(&m2);
335 g.update(&sk.hpk);
336 g.final(&kr2);
337
338 // ct' = innerEnc(pk, m', r')
339 const ct2 = sk.pk.encrypt(&m2, kr2[32..64]);
340
341 if (p.ml_kem) {
342 // ML-KEM: K = K'' if ct == ct', else K = J(z || c) per FIPS 203
343 var k_bar: [shared_length]u8 = undefined;
344 var j = sha3.Shake256.init(.{});
345 j.update(&sk.z);
346 j.update(ct);
347 j.squeeze(&k_bar);
348 cmov(shared_length, kr2[0..shared_length], k_bar, ctneq(ciphertext_length, ct.*, ct2));
349 return kr2[0..shared_length].*;
350 } else {
351 // Kyber: K = KDF(K''/z ‖ H(c))
352 sha3.Sha3_256.hash(ct, kr2[32..], .{});
353 cmov(32, kr2[0..32], sk.z, ctneq(ciphertext_length, ct.*, ct2));
354 var ss: [shared_length]u8 = undefined;
355 sha3.Shake256.hash(&kr2, &ss, .{});
356 return ss;
357 }
358 }
359
360 /// Serializes the key into a byte array.
361 pub fn toBytes(sk: SecretKey) [encoded_length]u8 {
362 return sk.sk.toBytes() ++ sk.pk.toBytes() ++ sk.hpk ++ sk.z;
363 }
364
365 /// Deserializes the key from a byte array.
366 pub fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!SecretKey {
367 var ret: SecretKey = undefined;
368 comptime var s: usize = 0;
369 ret.sk = InnerSk.fromBytes(buf[s .. s + InnerSk.encoded_length]);
370 s += InnerSk.encoded_length;
371 ret.pk = try InnerPk.fromBytes(buf[s .. s + InnerPk.encoded_length]);
372 s += InnerPk.encoded_length;
373 ret.hpk = buf[s..][0..h_length].*;
374 s += h_length;
375 ret.z = buf[s..][0..shared_length].*;
376 return ret;
377 }
378 };
379
380 /// A Kyber key pair.
381 pub const KeyPair = struct {
382 secret_key: SecretKey,
383 public_key: PublicKey,
384
385 /// Deterministically derive a key pair from a cryptograpically secure secret seed.
386 ///
387 /// Except in tests, applications should generally call `generate()` instead of this function.
388 pub fn generateDeterministic(seed: [seed_length]u8) !KeyPair {
389 var ret: KeyPair = undefined;
390
391 // Generate inner key
392 innerKeyFromSeed(
393 seed[0..inner_seed_length].*,
394 &ret.public_key.pk,
395 &ret.secret_key.sk,
396 );
397 ret.secret_key.pk = ret.public_key.pk;
398
399 // Copy over z from seed.
400 ret.secret_key.z = seed[inner_seed_length..seed_length].*;
401
402 // Compute H(pk)
403 sha3.Sha3_256.hash(&ret.public_key.pk.toBytes(), &ret.secret_key.hpk, .{});
404 ret.public_key.hpk = ret.secret_key.hpk;
405
406 return ret;
407 }
408
409 /// Generate a new, random key pair.
410 pub fn generate(io: std.Io) KeyPair {
411 var random_seed: [seed_length]u8 = undefined;
412 while (true) {
413 io.random(&random_seed);
414 return generateDeterministic(random_seed) catch {
415 @branchHint(.unlikely);
416 continue;
417 };
418 }
419 }
420 };
421
422 // Size of plaintexts of the in
423 const inner_plaintext_length: usize = Poly.compressedSize(1);
424
425 const InnerPk = struct {
426 rho: [32]u8, // ρ, the seed for the matrix A
427 th: V, // NTT(t), normalized
428
429 // Cached values
430 aT: M,
431
432 const encoded_length = V.encoded_length + 32;
433
434 fn encrypt(
435 pk: InnerPk,
436 pt: *const [inner_plaintext_length]u8,
437 seed: *const [32]u8,
438 ) [ciphertext_length]u8 {
439 // Sample r, e₁ and e₂ appropriately
440 const rh = V.noise(p.eta1, 0, seed).ntt().barrettReduce();
441 const e1 = V.noise(eta2, p.k, seed);
442 const e2 = Poly.noise(eta2, 2 * p.k, seed);
443
444 // Next we compute u = Aᵀ r + e₁. First Aᵀ.
445 var u: V = undefined;
446 for (0..p.k) |i| {
447 // Note that coefficients of r are bounded by q and those of Aᵀ
448 // are bounded by 4.5q and so their product is bounded by 2¹⁵q
449 // as required for multiplication.
450 u.ps[i] = pk.aT.rows[i].dotHat(rh);
451 }
452
453 // Aᵀ and r were not in Montgomery form, so the Montgomery
454 // multiplications in the inner product added a factor R⁻¹ which
455 // the InvNTT cancels out.
456 u = u.barrettReduce().invNTT().add(e1).normalize();
457
458 // Next, compute v = <t, r> + e₂ + Decompress_q(m, 1)
459 const v = pk.th.dotHat(rh).barrettReduce().invNTT()
460 .add(Poly.decompress(1, pt)).add(e2).normalize();
461
462 return u.compress(p.du) ++ v.compress(p.dv);
463 }
464
465 fn toBytes(pk: InnerPk) [encoded_length]u8 {
466 return pk.th.toBytes() ++ pk.rho;
467 }
468
469 fn fromBytes(buf: *const [encoded_length]u8) errors.NonCanonicalError!InnerPk {
470 var ret: InnerPk = undefined;
471
472 const th_bytes = buf[0..V.encoded_length];
473 ret.th = V.fromBytes(th_bytes).normalize();
474
475 if (p.ml_kem) {
476 // Verify that the coefficients used a canonical representation.
477 if (!mem.eql(u8, &ret.th.toBytes(), th_bytes)) {
478 return error.NonCanonical;
479 }
480 }
481
482 ret.rho = buf[V.encoded_length..encoded_length].*;
483 ret.aT = M.uniform(ret.rho, true);
484 return ret;
485 }
486 };
487
488 // Private key of the inner PKE
489 const InnerSk = struct {
490 sh: V, // NTT(s), normalized
491 const encoded_length = V.encoded_length;
492
493 fn decrypt(sk: InnerSk, ct: *const [ciphertext_length]u8) [inner_plaintext_length]u8 {
494 const u = V.decompress(p.du, ct[0..comptime V.compressedSize(p.du)]);
495 const v = Poly.decompress(
496 p.dv,
497 ct[comptime V.compressedSize(p.du)..ciphertext_length],
498 );
499
500 // Compute m = v - <s, u>
501 return v.sub(sk.sh.dotHat(u.ntt()).barrettReduce().invNTT())
502 .normalize().compress(1);
503 }
504
505 fn toBytes(sk: InnerSk) [encoded_length]u8 {
506 return sk.sh.toBytes();
507 }
508
509 fn fromBytes(buf: *const [encoded_length]u8) InnerSk {
510 var ret: InnerSk = undefined;
511 ret.sh = V.fromBytes(buf).normalize();
512 return ret;
513 }
514 };
515
516 // Derives inner PKE keypair from given seed.
517 fn innerKeyFromSeed(seed: [inner_seed_length]u8, pk: *InnerPk, sk: *InnerSk) void {
518 var expanded_seed: [64]u8 = undefined;
519 var h = sha3.Sha3_512.init(.{});
520 h.update(&seed);
521 if (p.ml_kem) h.update(&[1]u8{p.k});
522 h.final(&expanded_seed);
523 pk.rho = expanded_seed[0..32].*;
524 const sigma = expanded_seed[32..64];
525 pk.aT = M.uniform(pk.rho, false); // Expand ρ to A; we'll transpose later on
526
527 // Sample secret vector s.
528 sk.sh = V.noise(p.eta1, 0, sigma).ntt().normalize();
529
530 const eh = PolyVec(p.k).noise(p.eta1, p.k, sigma).ntt(); // sample blind e.
531 var th: V = undefined;
532
533 // Next, we compute t = A s + e.
534 for (0..p.k) |i| {
535 // Note that coefficients of s are bounded by q and those of A
536 // are bounded by 4.5q and so their product is bounded by 2¹⁵q
537 // as required for multiplication.
538 // A and s were not in Montgomery form, so the Montgomery
539 // multiplications in the inner product added a factor R⁻¹ which
540 // we'll cancel out with toMont(). This will also ensure the
541 // coefficients of th are bounded in absolute value by q.
542 th.ps[i] = pk.aT.rows[i].dotHat(sk.sh).toMont();
543 }
544
545 pk.th = th.add(eh).normalize(); // bounded by 8q
546 pk.aT = pk.aT.transpose();
547 }
548 };
549}
550
551// R mod q
552const r_mod_q: i32 = @rem(@as(i32, R), Q);
553
554// R² mod q
555const r2_mod_q: i32 = @rem(r_mod_q * r_mod_q, Q);
556
557// ζ is the degree 256 primitive root of unity used for the NTT.
558const zeta: i16 = 17;
559
560// (128)⁻¹ R². Used in inverse NTT.
561const r2_over_128: i32 = @mod(invertMod(128, Q) * r2_mod_q, Q);
562
563// zetas lists precomputed powers of the primitive root of unity in
564// Montgomery representation used for the NTT:
565//
566// zetas[i] = ζᵇʳᵛ⁽ⁱ⁾ R mod q
567//
568// where ζ = 17, brv(i) is the bitreversal of a 7-bit number and R=2¹⁶ mod q.
569const zetas = computeZetas();
570
571// invNTTReductions keeps track of which coefficients to apply Barrett
572// reduction to in Poly.invNTT().
573//
574// Generated lazily: once a butterfly is computed which is about to
575// overflow the i16, the largest coefficient is reduced. If that is
576// not enough, the other coefficient is reduced as well.
577//
578// This is actually optimal, as proven in https://eprint.iacr.org/2020/1377.pdf
579const inv_ntt_reductions = [_]i16{
580 -1, // after layer 1
581 -1, // after layer 2
582 16,
583 17,
584 48,
585 49,
586 80,
587 81,
588 112,
589 113,
590 144,
591 145,
592 176,
593 177,
594 208,
595 209,
596 240, 241, -1, // after layer 3
597 0, 1, 32,
598 33, 34, 35,
599 64, 65, 96,
600 97, 98, 99,
601 128, 129,
602 160, 161, 162, 163, 192, 193, 224, 225, 226, 227, -1, // after layer 4
603 2, 3, 66, 67, 68, 69, 70, 71, 130, 131, 194,
604 195, 196, 197,
605 198, 199, -1, // after layer 5
606 4, 5, 6,
607 7, 132, 133,
608 134, 135, 136,
609 137, 138, 139,
610 140, 141,
611 142, 143, -1, // after layer 6
612 -1, // after layer 7
613};
614
615test "invNTTReductions bounds" {
616 // Checks whether the reductions proposed by invNTTReductions
617 // don't overflow during invNTT().
618 var xs: [256]i32 = @splat(1); // start at |x| ≤ q
619
620 var r: usize = 0;
621 var layer: math.Log2Int(usize) = 1;
622 while (layer < 8) : (layer += 1) {
623 const w = @as(usize, 1) << layer;
624 var i: usize = 0;
625
626 while (i + w < 256) {
627 xs[i] = xs[i] + xs[i + w];
628 try testing.expect(xs[i] <= 9); // we can't exceed 9q
629 xs[i + w] = 1;
630 i += 1;
631 if (@mod(i, w) == 0) {
632 i += w;
633 }
634 }
635
636 while (true) {
637 const j = inv_ntt_reductions[r];
638 r += 1;
639 if (j < 0) {
640 break;
641 }
642 xs[@as(usize, @intCast(j))] = 1;
643 }
644 }
645}
646
647fn invertMod(a: anytype, p: @TypeOf(a)) @TypeOf(a) {
648 const r = extendedEuclidean(@TypeOf(a), a, p);
649 assert(r.gcd == 1);
650 return r.x;
651}
652
653// Reduce mod q for testing.
654fn modQ32(x: i32) i16 {
655 var y = @as(i16, @intCast(@rem(x, @as(i32, Q))));
656 if (y < 0) {
657 y += Q;
658 }
659 return y;
660}
661
662// Given -2¹⁵ q ≤ x < 2¹⁵ q, returns -q < y < q with x 2⁻¹⁶ = y (mod q).
663fn montReduce(x: i32) i16 {
664 const qInv = comptime invertMod(@as(i32, Q), R);
665 // This is Montgomery reduction with R=2¹⁶.
666 //
667 // Note gcd(2¹⁶, q) = 1 as q is prime. Write q' := 62209 = q⁻¹ mod R.
668 // First we compute
669 //
670 // m := ((x mod R) q') mod R
671 // = x q' mod R
672 // = int16(x q')
673 // = int16(int32(x) * int32(q'))
674 //
675 // Note that x q' might be as big as 2³² and could overflow the int32
676 // multiplication in the last line. However for any int32s a and b,
677 // we have int32(int64(a)*int64(b)) = int32(a*b) and so the result is ok.
678 const m: i16 = @truncate(@as(i32, @truncate(x *% qInv)));
679
680 // Note that x - m q is divisible by R; indeed modulo R we have
681 //
682 // x - m q ≡ x - x q' q ≡ x - x q⁻¹ q ≡ x - x = 0.
683 //
684 // We return y := (x - m q) / R. Note that y is indeed correct as
685 // modulo q we have
686 //
687 // y ≡ x R⁻¹ - m q R⁻¹ = x R⁻¹
688 //
689 // and as both 2¹⁵ q ≤ m q, x < 2¹⁵ q, we have
690 // 2¹⁶ q ≤ x - m q < 2¹⁶ and so q ≤ (x - m q) / R < q as desired.
691 const yR = x - @as(i32, m) * @as(i32, Q);
692 return @bitCast(@as(u16, @truncate(@as(u32, @bitCast(yR)) >> 16)));
693}
694
695test "Test montReduce" {
696 var rnd = RndGen.init(0);
697 for (0..1000) |_| {
698 const bound = comptime @as(i32, Q) * (1 << 15);
699 const x = rnd.random().intRangeLessThan(i32, -bound, bound);
700 const y = montReduce(x);
701 try testing.expect(-Q < y and y < Q);
702 try testing.expectEqual(modQ32(x), modQ32(@as(i32, y) * R));
703 }
704}
705
706// Given any x, return x R mod q where R=2¹⁶.
707fn feToMont(x: i16) i16 {
708 // Note |1353 x| ≤ 1353 2¹⁵ ≤ 13318 q ≤ 2¹⁵ q and so we're within
709 // the bounds of montReduce.
710 return montReduce(@as(i32, x) * r2_mod_q);
711}
712
713test "Test feToMont" {
714 var x: i32 = -(1 << 15);
715 while (x < 1 << 15) : (x += 1) {
716 const y = feToMont(@as(i16, @intCast(x)));
717 try testing.expectEqual(modQ32(@as(i32, y)), modQ32(x * r_mod_q));
718 }
719}
720
721// Given any x, compute 0 ≤ y ≤ q with x = y (mod q).
722//
723// Beware: we might have feBarrettReduce(x) = q ≠ 0 for some x. In fact,
724// this happens if and only if x = -nq for some positive integer n.
725fn feBarrettReduce(x: i16) i16 {
726 // This is standard Barrett reduction.
727 //
728 // For any x we have x mod q = x - ⌊x/q⌋ q. We will use 20159/2²⁶ as
729 // an approximation of 1/q. Note that 0 ≤ 20159/2²⁶ - 1/q ≤ 0.135/2²⁶
730 // and so | x 20156/2²⁶ - x/q | ≤ 2⁻¹⁰ for |x| ≤ 2¹⁶. For all x
731 // not a multiple of q, the number x/q is further than 1/q from any integer
732 // and so ⌊x 20156/2²⁶⌋ = ⌊x/q⌋. If x is a multiple of q and x is positive,
733 // then x 20156/2²⁶ is larger than x/q so ⌊x 20156/2²⁶⌋ = ⌊x/q⌋ as well.
734 // Finally, if x is negative multiple of q, then ⌊x 20156/2²⁶⌋ = ⌊x/q⌋-1.
735 // Thus
736 // [ q if x=-nq for pos. integer n
737 // x - ⌊x 20156/2²⁶⌋ q = [
738 // [ x mod q otherwise
739 //
740 // To actually compute this, note that
741 //
742 // ⌊x 20156/2²⁶⌋ = (20159 x) >> 26.
743 return x -% @as(i16, @intCast((@as(i32, x) * 20159) >> 26)) *% Q;
744}
745
746test "Test Barrett reduction" {
747 var x: i32 = -(1 << 15);
748 while (x < 1 << 15) : (x += 1) {
749 var y1 = feBarrettReduce(@as(i16, @intCast(x)));
750 const y2 = @mod(@as(i16, @intCast(x)), Q);
751 if (x < 0 and @rem(-x, Q) == 0) {
752 y1 -= Q;
753 }
754 try testing.expectEqual(y1, y2);
755 }
756}
757
758// Returns x if x < q and x - q otherwise. Assumes x ≥ -29439.
759fn csubq(x: i16) i16 {
760 var r = x;
761 r -= Q;
762 r += (r >> 15) & Q;
763 return r;
764}
765
766test "Test csubq" {
767 var x: i32 = -29439;
768 while (x < 1 << 15) : (x += 1) {
769 const y1 = csubq(@as(i16, @intCast(x)));
770 var y2 = @as(i16, @intCast(x));
771 if (@as(i16, @intCast(x)) >= Q) {
772 y2 -= Q;
773 }
774 try testing.expectEqual(y1, y2);
775 }
776}
777
778// Computes zetas table used by ntt and invNTT.
779fn computeZetas() [128]i16 {
780 @setEvalBranchQuota(10000);
781 var ret: [128]i16 = undefined;
782 for (&ret, 0..) |*r, i| {
783 const t = @as(i16, @intCast(modularPow(i32, zeta, @bitReverse(@as(u7, @intCast(i))), Q)));
784 r.* = csubq(feBarrettReduce(feToMont(t)));
785 }
786 return ret;
787}
788
789// An element of our base ring R which are polynomials over ℤ_q
790// modulo the equation Xᴺ = -1, where q=3329 and N=256.
791//
792// This type is also used to store NTT-transformed polynomials,
793// see Poly.NTT().
794//
795// Coefficients aren't always reduced. See Normalize().
796const Poly = struct {
797 cs: [N]i16,
798
799 const encoded_length = N / 2 * 3;
800 const zero: Poly = .{ .cs = @splat(0) };
801
802 // Add two polynomials (coefficients not normalized)
803 fn add(a: Poly, b: Poly) Poly {
804 var ret: Poly = undefined;
805 for (0..N) |i| {
806 ret.cs[i] = a.cs[i] + b.cs[i];
807 }
808 return ret;
809 }
810
811 // Subtract two polynomials (coefficients not normalized)
812 fn sub(a: Poly, b: Poly) Poly {
813 var ret: Poly = undefined;
814 for (0..N) |i| {
815 ret.cs[i] = a.cs[i] - b.cs[i];
816 }
817 return ret;
818 }
819
820 // Executes a forward "NTT" on p.
821 //
822 // Assumes the coefficients are in absolute value ≤q. The resulting
823 // coefficients are in absolute value ≤7q. If the input is in Montgomery
824 // form, then the result is in Montgomery form and so (by linearity of the NTT)
825 // if the input is in regular form, then the result is also in regular form.
826 fn ntt(a: Poly) Poly {
827 // Note that ℤ_q does not have a primitive 512ᵗʰ root of unity (as 512
828 // does not divide into q-1) and so we cannot do a regular NTT. ℤ_q
829 // does have a primitive 256ᵗʰ root of unity, the smallest of which
830 // is ζ := 17.
831 //
832 // Recall that our base ring R := ℤ_q[x] / (x²⁵⁶ + 1). The polynomial
833 // x²⁵⁶+1 will not split completely (as its roots would be 512ᵗʰ roots
834 // of unity.) However, it does split almost (using ζ¹²⁸ = -1):
835 //
836 // x²⁵⁶ + 1 = (x²)¹²⁸ - ζ¹²⁸
837 // = ((x²)⁶⁴ - ζ⁶⁴)((x²)⁶⁴ + ζ⁶⁴)
838 // = ((x²)³² - ζ³²)((x²)³² + ζ³²)((x²)³² - ζ⁹⁶)((x²)³² + ζ⁹⁶)
839 // ⋮
840 // = (x² - ζ)(x² + ζ)(x² - ζ⁶⁵)(x² + ζ⁶⁵) … (x² + ζ¹²⁷)
841 //
842 // Note that the powers of ζ that appear (from the second line down) are
843 // in binary
844 //
845 // 0100000 1100000
846 // 0010000 1010000 0110000 1110000
847 // 0001000 1001000 0101000 1101000 0011000 1011000 0111000 1111000
848 // …
849 //
850 // That is: brv(2), brv(3), brv(4), …, where brv(x) denotes the 7-bit
851 // bitreversal of x. These powers of ζ are given by the Zetas array.
852 //
853 // The polynomials x² ± ζⁱ are irreducible and coprime, hence by
854 // the Chinese Remainder Theorem we know
855 //
856 // ℤ_q[x]/(x²⁵⁶+1) → ℤ_q[x]/(x²-ζ) x … x ℤ_q[x]/(x²+ζ¹²⁷)
857 //
858 // given by a ↦ ( a mod x²-ζ, …, a mod x²+ζ¹²⁷ )
859 // is an isomorphism, which is the "NTT". It can be efficiently computed by
860 //
861 //
862 // a ↦ ( a mod (x²)⁶⁴ - ζ⁶⁴, a mod (x²)⁶⁴ + ζ⁶⁴ )
863 // ↦ ( a mod (x²)³² - ζ³², a mod (x²)³² + ζ³²,
864 // a mod (x²)⁹⁶ - ζ⁹⁶, a mod (x²)⁹⁶ + ζ⁹⁶ )
865 //
866 // et cetera
867 // If N was 8 then this can be pictured in the following diagram:
868 //
869 // https://cnx.org/resources/17ee4dfe517a6adda05377b25a00bf6e6c93c334/File0026.png
870 //
871 // Each cross is a Cooley-Tukey butterfly: it's the map
872 //
873 // (a, b) ↦ (a + ζb, a - ζb)
874 //
875 // for the appropriate power ζ for that column and row group.
876 var p = a;
877 var k: usize = 0; // index into zetas
878
879 var l = N >> 1;
880 while (l > 1) : (l >>= 1) {
881 // On the nᵗʰ iteration of the l-loop, the absolute value of the
882 // coefficients are bounded by nq.
883
884 // offset effectively loops over the row groups in this column; it is
885 // the first row in the row group.
886 var offset: usize = 0;
887 while (offset < N - l) : (offset += 2 * l) {
888 k += 1;
889 const z = @as(i32, zetas[k]);
890
891 // j loops over each butterfly in the row group.
892 for (offset..offset + l) |j| {
893 const t = montReduce(z * @as(i32, p.cs[j + l]));
894 p.cs[j + l] = p.cs[j] - t;
895 p.cs[j] += t;
896 }
897 }
898 }
899
900 return p;
901 }
902
903 // Executes an inverse "NTT" on p and multiply by the Montgomery factor R.
904 //
905 // Assumes the coefficients are in absolute value ≤q. The resulting
906 // coefficients are in absolute value ≤q. If the input is in Montgomery
907 // form, then the result is in Montgomery form and so (by linearity)
908 // if the input is in regular form, then the result is also in regular form.
909 fn invNTT(a: Poly) Poly {
910 var k: usize = 127; // index into zetas
911 var r: usize = 0; // index into invNTTReductions
912 var p = a;
913
914 // We basically do the oppposite of NTT, but postpone dividing by 2 in the
915 // inverse of the Cooley-Tukey butterfly and accumulate that into a big
916 // division by 2⁷ at the end. See the comments in the ntt() function.
917
918 var l: usize = 2;
919 while (l < N) : (l <<= 1) {
920 var offset: usize = 0;
921 while (offset < N - l) : (offset += 2 * l) {
922 // As we're inverting, we need powers of ζ⁻¹ (instead of ζ).
923 // To be precise, we need ζᵇʳᵛ⁽ᵏ⁾⁻¹²⁸. However, as ζ⁻¹²⁸ = -1,
924 // we can use the existing zetas table instead of
925 // keeping a separate invZetas table as in Dilithium.
926
927 const minZeta = @as(i32, zetas[k]);
928 k -= 1;
929
930 for (offset..offset + l) |j| {
931 // Gentleman-Sande butterfly: (a, b) ↦ (a + b, ζ(a-b))
932 const t = p.cs[j + l] - p.cs[j];
933 p.cs[j] += p.cs[j + l];
934 p.cs[j + l] = montReduce(minZeta * @as(i32, t));
935
936 // Note that if we had |a| < αq and |b| < βq before the
937 // butterfly, then now we have |a| < (α+β)q and |b| < q.
938 }
939 }
940
941 // We let the invNTTReductions instruct us which coefficients to
942 // Barrett reduce.
943 while (true) {
944 const i = inv_ntt_reductions[r];
945 r += 1;
946 if (i < 0) {
947 break;
948 }
949 p.cs[@as(usize, @intCast(i))] = feBarrettReduce(p.cs[@as(usize, @intCast(i))]);
950 }
951 }
952
953 for (0..N) |j| {
954 // Note 1441 = (128)⁻¹ R². The coefficients are bounded by 9q, so
955 // as 1441 * 9 ≈ 2¹⁴ < 2¹⁵, we're within the required bounds
956 // for montReduce().
957 p.cs[j] = montReduce(r2_over_128 * @as(i32, p.cs[j]));
958 }
959
960 return p;
961 }
962
963 // Normalizes coefficients.
964 //
965 // Ensures each coefficient is in {0, …, q-1}.
966 fn normalize(a: Poly) Poly {
967 var ret: Poly = undefined;
968 for (0..N) |i| {
969 ret.cs[i] = csubq(feBarrettReduce(a.cs[i]));
970 }
971 return ret;
972 }
973
974 // Put p in Montgomery form.
975 fn toMont(a: Poly) Poly {
976 var ret: Poly = undefined;
977 for (0..N) |i| {
978 ret.cs[i] = feToMont(a.cs[i]);
979 }
980 return ret;
981 }
982
983 // Barret reduce coefficients.
984 //
985 // Beware, this does not fully normalize coefficients.
986 fn barrettReduce(a: Poly) Poly {
987 var ret: Poly = undefined;
988 for (0..N) |i| {
989 ret.cs[i] = feBarrettReduce(a.cs[i]);
990 }
991 return ret;
992 }
993
994 fn compressedSize(comptime d: u8) usize {
995 return @divTrunc(N * d, 8);
996 }
997
998 // Returns packed Compress_q(p, d).
999 //
1000 // Assumes p is normalized.
1001 fn compress(p: Poly, comptime d: u8) [compressedSize(d)]u8 {
1002 @setEvalBranchQuota(10000);
1003 const q_over_2: u32 = comptime @divTrunc(Q, 2); // (q-1)/2
1004 const two_d_min_1: u32 = comptime (1 << d) - 1; // 2ᵈ-1
1005 var in_off: usize = 0;
1006 var out_off: usize = 0;
1007
1008 const batch_size: usize = comptime math.lcm(d, 8);
1009 const in_batch_size: usize = comptime batch_size / d;
1010 const out_batch_size: usize = comptime batch_size / 8;
1011
1012 const out_length: usize = comptime @divTrunc(N * d, 8);
1013 comptime assert(out_length * 8 == d * N);
1014 var out: [out_length]u8 = @splat(0);
1015
1016 while (in_off < N) {
1017 // First we compress into in.
1018 var in: [in_batch_size]u16 = undefined;
1019 inline for (0..in_batch_size) |i| {
1020 // Compress_q(x, d) = ⌈(2ᵈ/q)x⌋ mod⁺ 2ᵈ
1021 // = ⌊(2ᵈ/q)x+½⌋ mod⁺ 2ᵈ
1022 // = ⌊((x << d) + q/2) / q⌋ mod⁺ 2ᵈ
1023 // = DIV((x << d) + q/2, q) & ((1<<d) - 1)
1024 const t = @as(u24, @intCast(p.cs[in_off + i])) << d;
1025 // Division by invariant multiplication, equivalent to DIV(t + q/2, q).
1026 // A division may not be a constant-time operation, even with a constant denominator.
1027 // Here, side channels would leak information about the shared secret, see https://kyberslash.cr.yp.to
1028 // Multiplication, on the other hand, is a constant-time operation on the CPUs we currently support.
1029 comptime assert(d <= 11);
1030 comptime assert(((20642679 * @as(u64, Q)) >> 36) == 1);
1031 const u: u32 = @intCast((@as(u64, t + q_over_2) * 20642679) >> 36);
1032 in[i] = @intCast(u & two_d_min_1);
1033 }
1034
1035 // Now we pack the d-bit integers from `in' into out as bytes.
1036 comptime var in_shift: usize = 0;
1037 comptime var j: usize = 0;
1038 comptime var i: usize = 0;
1039 inline while (i < in_batch_size) : (j += 1) {
1040 comptime var todo: usize = 8;
1041 inline while (todo > 0) {
1042 const out_shift = comptime 8 - todo;
1043 out[out_off + j] |= @as(u8, @truncate((in[i] >> in_shift) << out_shift));
1044
1045 const done = comptime @min(@min(d, todo), d - in_shift);
1046 todo -= done;
1047 in_shift += done;
1048
1049 if (in_shift == d) {
1050 in_shift = 0;
1051 i += 1;
1052 }
1053 }
1054 }
1055
1056 in_off += in_batch_size;
1057 out_off += out_batch_size;
1058 }
1059
1060 return out;
1061 }
1062
1063 // Set p to Decompress_q(m, d).
1064 fn decompress(comptime d: u8, in: *const [compressedSize(d)]u8) Poly {
1065 @setEvalBranchQuota(10000);
1066 const in_len = comptime @divTrunc(N * d, 8);
1067 comptime assert(in_len * 8 == d * N);
1068 var ret: Poly = undefined;
1069 var in_off: usize = 0;
1070 var out_off: usize = 0;
1071
1072 const batch_size: usize = comptime math.lcm(d, 8);
1073 const in_batch_size: usize = comptime batch_size / 8;
1074 const out_batch_size: usize = comptime batch_size / d;
1075
1076 while (out_off < N) {
1077 comptime var in_shift: usize = 0;
1078 comptime var j: usize = 0;
1079 comptime var i: usize = 0;
1080 inline while (i < out_batch_size) : (i += 1) {
1081 // First, unpack next coefficient.
1082 comptime var todo = d;
1083 var out: u16 = 0;
1084
1085 inline while (todo > 0) {
1086 const out_shift = comptime d - todo;
1087 const m = comptime (1 << d) - 1;
1088 out |= (@as(u16, in[in_off + j] >> in_shift) << out_shift) & m;
1089
1090 const done = comptime @min(@min(8, todo), 8 - in_shift);
1091 todo -= done;
1092 in_shift += done;
1093
1094 if (in_shift == 8) {
1095 in_shift = 0;
1096 j += 1;
1097 }
1098 }
1099
1100 // Decompress_q(x, d) = ⌈(q/2ᵈ)x⌋
1101 // = ⌊(q/2ᵈ)x+½⌋
1102 // = ⌊(qx + 2ᵈ⁻¹)/2ᵈ⌋
1103 // = (qx + (1<<(d-1))) >> d
1104 const qx = @as(u32, out) * @as(u32, Q);
1105 ret.cs[out_off + i] = @as(i16, @intCast((qx + (1 << (d - 1))) >> d));
1106 }
1107
1108 in_off += in_batch_size;
1109 out_off += out_batch_size;
1110 }
1111
1112 return ret;
1113 }
1114
1115 // Returns the "pointwise" multiplication a o b.
1116 //
1117 // That is: invNTT(a o b) = invNTT(a) * invNTT(b). Assumes a and b are in
1118 // Montgomery form. Products between coefficients of a and b must be strictly
1119 // bounded in absolute value by 2¹⁵q. a o b will be in Montgomery form and
1120 // bounded in absolute value by 2q.
1121 fn mulHat(a: Poly, b: Poly) Poly {
1122 // Recall from the discussion in ntt(), that a transformed polynomial is
1123 // an element of ℤ_q[x]/(x²-ζ) x … x ℤ_q[x]/(x²+ζ¹²⁷);
1124 // that is: 128 degree-one polynomials instead of simply 256 elements
1125 // from ℤ_q as in the regular NTT. So instead of pointwise multiplication,
1126 // we multiply the 128 pairs of degree-one polynomials modulo the
1127 // right equation:
1128 //
1129 // (a₁ + a₂x)(b₁ + b₂x) = a₁b₁ + a₂b₂ζ' + (a₁b₂ + a₂b₁)x,
1130 //
1131 // where ζ' is the appropriate power of ζ.
1132
1133 var p: Poly = undefined;
1134 var k: usize = 64;
1135 var i: usize = 0;
1136 while (i < N) : (i += 4) {
1137 const z = @as(i32, zetas[k]);
1138 k += 1;
1139
1140 const a1b1 = montReduce(@as(i32, a.cs[i + 1]) * @as(i32, b.cs[i + 1]));
1141 const a0b0 = montReduce(@as(i32, a.cs[i]) * @as(i32, b.cs[i]));
1142 const a1b0 = montReduce(@as(i32, a.cs[i + 1]) * @as(i32, b.cs[i]));
1143 const a0b1 = montReduce(@as(i32, a.cs[i]) * @as(i32, b.cs[i + 1]));
1144
1145 p.cs[i] = montReduce(a1b1 * z) + a0b0;
1146 p.cs[i + 1] = a0b1 + a1b0;
1147
1148 const a3b3 = montReduce(@as(i32, a.cs[i + 3]) * @as(i32, b.cs[i + 3]));
1149 const a2b2 = montReduce(@as(i32, a.cs[i + 2]) * @as(i32, b.cs[i + 2]));
1150 const a3b2 = montReduce(@as(i32, a.cs[i + 3]) * @as(i32, b.cs[i + 2]));
1151 const a2b3 = montReduce(@as(i32, a.cs[i + 2]) * @as(i32, b.cs[i + 3]));
1152
1153 p.cs[i + 2] = a2b2 - montReduce(a3b3 * z);
1154 p.cs[i + 3] = a2b3 + a3b2;
1155 }
1156
1157 return p;
1158 }
1159
1160 // Sample p from a centered binomial distribution with n=2η and p=½ - viz:
1161 // coefficients are in {-η, …, η} with probabilities
1162 //
1163 // {ncr(0, 2η)/2^2η, ncr(1, 2η)/2^2η, …, ncr(2η,2η)/2^2η}
1164 fn noise(comptime eta: u8, nonce: u8, seed: *const [32]u8) Poly {
1165 var h = sha3.Shake256.init(.{});
1166 const suffix: [1]u8 = .{nonce};
1167 h.update(seed);
1168 h.update(&suffix);
1169
1170 // The distribution at hand is exactly the same as that
1171 // of (a₁ + a₂ + … + a_η) - (b₁ + … + b_η) where a_i,b_i~U(1).
1172 // Thus we need 2η bits per coefficient.
1173 const buf_len = comptime 2 * eta * N / 8;
1174 var buf: [buf_len]u8 = undefined;
1175 h.squeeze(&buf);
1176
1177 // buf is interpreted as a₁…a_ηb₁…b_ηa₁…a_ηb₁…b_η…. We process
1178 // multiple coefficients in one batch.
1179
1180 const T = switch (builtin.target.cpu.arch) {
1181 .x86_64, .x86 => u32, // Generates better code on Intel CPUs
1182 else => u64, // u128 might be faster on some other CPUs.
1183 };
1184
1185 comptime var batch_count: usize = undefined;
1186 comptime var batch_bytes: usize = undefined;
1187 comptime var mask: T = 0;
1188 comptime {
1189 batch_count = @bitSizeOf(T) / @as(usize, 2 * eta);
1190 while (@rem(N, batch_count) != 0 and batch_count > 0) : (batch_count -= 1) {}
1191 assert(batch_count > 0);
1192 assert(@rem(2 * eta * batch_count, 8) == 0);
1193 batch_bytes = 2 * eta * batch_count / 8;
1194
1195 for (0..2 * eta * batch_count) |_| {
1196 mask <<= eta;
1197 mask |= 1;
1198 }
1199 }
1200
1201 var ret: Poly = undefined;
1202 for (0..comptime N / batch_count) |i| {
1203 // Read coefficients into t. In the case of η=3,
1204 // we have t = a₁ + 2a₂ + 4a₃ + 8b₁ + 16b₂ + …
1205 var t: T = 0;
1206 inline for (0..batch_bytes) |j| {
1207 t |= @as(T, buf[batch_bytes * i + j]) << (8 * j);
1208 }
1209
1210 // Accumulate `a's and `b's together by masking them out, shifting
1211 // and adding. For η=3, we have d = a₁ + a₂ + a₃ + 8(b₁ + b₂ + b₃) + …
1212 var d: T = 0;
1213 inline for (0..eta) |j| {
1214 d += (t >> j) & mask;
1215 }
1216
1217 // Extract each a and b separately and set coefficient in polynomial.
1218 inline for (0..batch_count) |j| {
1219 const mask2 = comptime (1 << eta) - 1;
1220 const a = @as(i16, @intCast((d >> (comptime (2 * j * eta))) & mask2));
1221 const b = @as(i16, @intCast((d >> (comptime ((2 * j + 1) * eta))) & mask2));
1222 ret.cs[batch_count * i + j] = a - b;
1223 }
1224 }
1225
1226 return ret;
1227 }
1228
1229 fn uniform(seed: [32]u8, x: u8, y: u8) Poly {
1230 const domain_sep: [2]u8 = .{ x, y };
1231 return sampleUniformRejection(
1232 Poly,
1233 Q,
1234 12,
1235 N,
1236 &seed,
1237 &domain_sep,
1238 );
1239 }
1240
1241 // Packs p.
1242 //
1243 // Assumes p is normalized (and not just Barrett reduced).
1244 fn toBytes(p: Poly) [encoded_length]u8 {
1245 var ret: [encoded_length]u8 = undefined;
1246 for (0..comptime N / 2) |i| {
1247 const t0 = @as(u16, @intCast(p.cs[2 * i]));
1248 const t1 = @as(u16, @intCast(p.cs[2 * i + 1]));
1249 ret[3 * i] = @as(u8, @truncate(t0));
1250 ret[3 * i + 1] = @as(u8, @truncate((t0 >> 8) | (t1 << 4)));
1251 ret[3 * i + 2] = @as(u8, @truncate(t1 >> 4));
1252 }
1253 return ret;
1254 }
1255
1256 // Unpacks a Poly from buf.
1257 //
1258 // p will not be normalized; instead 0 ≤ p[i] < 4096.
1259 fn fromBytes(buf: *const [encoded_length]u8) Poly {
1260 var ret: Poly = undefined;
1261 for (0..comptime N / 2) |i| {
1262 const b0 = @as(i16, buf[3 * i]);
1263 const b1 = @as(i16, buf[3 * i + 1]);
1264 const b2 = @as(i16, buf[3 * i + 2]);
1265 ret.cs[2 * i] = b0 | ((b1 & 0xf) << 8);
1266 ret.cs[2 * i + 1] = (b1 >> 4) | b2 << 4;
1267 }
1268 return ret;
1269 }
1270};
1271
1272// A vector of k polynomials.
1273fn PolyVec(comptime k: u8) type {
1274 return struct {
1275 ps: [k]Poly,
1276
1277 const Self = @This();
1278 const encoded_length = k * Poly.encoded_length;
1279
1280 fn compressedSize(comptime d: u8) usize {
1281 return Poly.compressedSize(d) * k;
1282 }
1283
1284 /// Apply unary operation to each polynomial
1285 fn map(v: Self, comptime op: fn (Poly) Poly) Self {
1286 var ret: Self = undefined;
1287 inline for (0..k) |i| {
1288 ret.ps[i] = op(v.ps[i]);
1289 }
1290 return ret;
1291 }
1292
1293 /// Apply binary operation pairwise
1294 fn mapBinary(a: Self, b: Self, comptime op: fn (Poly, Poly) Poly) Self {
1295 var ret: Self = undefined;
1296 inline for (0..k) |i| {
1297 ret.ps[i] = op(a.ps[i], b.ps[i]);
1298 }
1299 return ret;
1300 }
1301
1302 fn ntt(v: Self) Self {
1303 return map(v, Poly.ntt);
1304 }
1305
1306 fn invNTT(v: Self) Self {
1307 return map(v, Poly.invNTT);
1308 }
1309
1310 fn normalize(v: Self) Self {
1311 return map(v, Poly.normalize);
1312 }
1313
1314 fn barrettReduce(v: Self) Self {
1315 return map(v, Poly.barrettReduce);
1316 }
1317
1318 fn add(a: Self, b: Self) Self {
1319 return mapBinary(a, b, Poly.add);
1320 }
1321
1322 fn sub(a: Self, b: Self) Self {
1323 return mapBinary(a, b, Poly.sub);
1324 }
1325
1326 // Samples v[i] from centered binomial distribution with the given η,
1327 // seed and nonce+i.
1328 fn noise(comptime eta: u8, nonce: u8, seed: *const [32]u8) Self {
1329 var ret: Self = undefined;
1330 for (0..k) |i| {
1331 ret.ps[i] = Poly.noise(eta, nonce + @as(u8, @intCast(i)), seed);
1332 }
1333 return ret;
1334 }
1335
1336 // Sets p to the inner product of a and b using "pointwise" multiplication.
1337 //
1338 // See MulHat() and NTT() for a description of the multiplication.
1339 // Assumes a and b are in Montgomery form. p will be in Montgomery form,
1340 // and its coefficients will be bounded in absolute value by 2kq.
1341 // If a and b are not in Montgomery form, then the action is the same
1342 // as "pointwise" multiplication followed by multiplying by R⁻¹, the inverse
1343 // of the Montgomery factor.
1344 fn dotHat(a: Self, b: Self) Poly {
1345 var ret: Poly = Poly.zero;
1346 for (0..k) |i| {
1347 ret = ret.add(a.ps[i].mulHat(b.ps[i]));
1348 }
1349 return ret;
1350 }
1351
1352 fn compress(v: Self, comptime d: u8) [compressedSize(d)]u8 {
1353 const cs = comptime Poly.compressedSize(d);
1354 var ret: [compressedSize(d)]u8 = undefined;
1355 inline for (0..k) |i| {
1356 ret[i * cs .. (i + 1) * cs].* = v.ps[i].compress(d);
1357 }
1358 return ret;
1359 }
1360
1361 fn decompress(comptime d: u8, buf: *const [compressedSize(d)]u8) Self {
1362 const cs = comptime Poly.compressedSize(d);
1363 var ret: Self = undefined;
1364 inline for (0..k) |i| {
1365 ret.ps[i] = Poly.decompress(d, buf[i * cs .. (i + 1) * cs]);
1366 }
1367 return ret;
1368 }
1369
1370 /// Serializes the key into a byte array.
1371 fn toBytes(v: Self) [encoded_length]u8 {
1372 var ret: [encoded_length]u8 = undefined;
1373 inline for (0..k) |i| {
1374 ret[i * Poly.encoded_length .. (i + 1) * Poly.encoded_length].* = v.ps[i].toBytes();
1375 }
1376 return ret;
1377 }
1378
1379 /// Deserializes the key from a byte array.
1380 fn fromBytes(buf: *const [encoded_length]u8) Self {
1381 var ret: Self = undefined;
1382 inline for (0..k) |i| {
1383 ret.ps[i] = Poly.fromBytes(
1384 buf[i * Poly.encoded_length .. (i + 1) * Poly.encoded_length],
1385 );
1386 }
1387 return ret;
1388 }
1389 };
1390}
1391
1392// A matrix of k vectors
1393fn Mat(comptime k: u8) type {
1394 return struct {
1395 const Self = @This();
1396 rows: [k]PolyVec(k),
1397
1398 fn uniform(seed: [32]u8, comptime transposed: bool) Self {
1399 var ret: Self = undefined;
1400 var i: u8 = 0;
1401 while (i < k) : (i += 1) {
1402 var j: u8 = 0;
1403 while (j < k) : (j += 1) {
1404 ret.rows[i].ps[j] = Poly.uniform(
1405 seed,
1406 if (transposed) i else j,
1407 if (transposed) j else i,
1408 );
1409 }
1410 }
1411 return ret;
1412 }
1413
1414 // Returns transpose of A
1415 fn transpose(m: Self) Self {
1416 var ret: Self = undefined;
1417 for (0..k) |i| {
1418 for (0..k) |j| {
1419 ret.rows[i].ps[j] = m.rows[j].ps[i];
1420 }
1421 }
1422 return ret;
1423 }
1424 };
1425}
1426
1427// Returns `true` if a ≠ b.
1428fn ctneq(comptime len: usize, a: [len]u8, b: [len]u8) u1 {
1429 return 1 - @intFromBool(crypto.timing_safe.eql([len]u8, a, b));
1430}
1431
1432// Copy src into dst given b = 1.
1433fn cmov(comptime len: usize, dst: *[len]u8, src: [len]u8, b: u1) void {
1434 const mask = @as(u8, 0) -% b;
1435 for (0..len) |i| {
1436 dst[i] ^= mask & (dst[i] ^ src[i]);
1437 }
1438}
1439
1440// Test helper: generates a random polynomial with each coefficient |x| ≤ q
1441fn randPolyAbsLeqQ(rnd: anytype) Poly {
1442 var ret: Poly = undefined;
1443 for (0..N) |i| {
1444 ret.cs[i] = rnd.random().intRangeAtMost(i16, -Q, Q);
1445 }
1446 return ret;
1447}
1448
1449// Test helper: generates a random normalized polynomial
1450fn randPolyNormalized(rnd: anytype) Poly {
1451 var ret: Poly = undefined;
1452 for (0..N) |i| {
1453 ret.cs[i] = rnd.random().intRangeLessThan(i16, 0, Q);
1454 }
1455 return ret;
1456}
1457
1458test "MulHat" {
1459 var rnd = RndGen.init(0);
1460
1461 for (0..100) |_| {
1462 const a = randPolyAbsLeqQ(&rnd);
1463 const b = randPolyAbsLeqQ(&rnd);
1464
1465 const p2 = a.ntt().mulHat(b.ntt()).barrettReduce().invNTT().normalize();
1466 var p: Poly = undefined;
1467
1468 @memset(&p.cs, 0);
1469
1470 for (0..N) |i| {
1471 for (0..N) |j| {
1472 var v = montReduce(@as(i32, a.cs[i]) * @as(i32, b.cs[j]));
1473 var k = i + j;
1474 if (k >= N) {
1475 // Recall Xᴺ = -1.
1476 k -= N;
1477 v = -v;
1478 }
1479 p.cs[k] = feBarrettReduce(v + p.cs[k]);
1480 }
1481 }
1482
1483 p = p.toMont().normalize();
1484
1485 try testing.expectEqual(p, p2);
1486 }
1487}
1488
1489test "NTT" {
1490 var rnd = RndGen.init(0);
1491
1492 for (0..1000) |_| {
1493 var p = randPolyAbsLeqQ(&rnd);
1494 const q = p.toMont().normalize();
1495 p = p.ntt();
1496
1497 for (0..N) |i| {
1498 try testing.expect(p.cs[i] <= 7 * Q and -7 * Q <= p.cs[i]);
1499 }
1500
1501 p = p.normalize().invNTT();
1502 for (0..N) |i| {
1503 try testing.expect(p.cs[i] <= Q and -Q <= p.cs[i]);
1504 }
1505
1506 p = p.normalize();
1507
1508 try testing.expectEqual(p, q);
1509 }
1510}
1511
1512test "Compression" {
1513 var rnd = RndGen.init(0);
1514 inline for (.{ 1, 4, 5, 10, 11 }) |d| {
1515 for (0..1000) |_| {
1516 const p = randPolyNormalized(&rnd);
1517 const pp = p.compress(d);
1518 const pq = Poly.decompress(d, &pp).compress(d);
1519 try testing.expectEqual(pp, pq);
1520 }
1521 }
1522}
1523
1524test "noise" {
1525 var seed: [32]u8 = undefined;
1526 for (&seed, 0..) |*s, i| {
1527 s.* = @as(u8, @intCast(i));
1528 }
1529 try testing.expectEqual(Poly.noise(3, 37, &seed).cs, .{
1530 0, 0, 1, -1, 0, 2, 0, -1, -1, 3, 0, 1, -2, -2, 0, 1, -2,
1531 1, 0, -2, 3, 0, 0, 0, 1, 3, 1, 1, 2, 1, -1, -1, -1, 0,
1532 1, 0, 1, 0, 2, 0, 1, -2, 0, -1, -1, -2, 1, -1, -1, 2, -1,
1533 1, 1, 2, -3, -1, -1, 0, 0, 0, 0, 1, -1, -2, -2, 0, -2, 0,
1534 0, 0, 1, 0, -1, -1, 1, -2, 2, 0, 0, 2, -2, 0, 1, 0, 1,
1535 1, 1, 0, 1, -2, -1, -2, -1, 1, 0, 0, 0, 0, 0, 1, 0, -1,
1536 -1, 0, -1, 1, 0, 1, 0, -1, -1, 0, -2, 2, 0, -2, 1, -1, 0,
1537 1, -1, -1, 2, 1, 0, 0, -2, -1, 2, 0, 0, 0, -1, -1, 3, 1,
1538 0, 1, 0, 1, 0, 2, 1, 0, 0, 1, 0, 1, 0, 0, -1, -1, -1,
1539 0, 1, 3, 1, 0, 1, 0, 1, -1, -1, -1, -1, 0, 0, -2, -1, -1,
1540 2, 0, 1, 0, 1, 0, 2, -2, 0, 1, 1, -3, -1, -2, -1, 0, 1,
1541 0, 1, -2, 2, 2, 1, 1, 0, -1, 0, -1, -1, 1, 0, -1, 2, 1,
1542 -1, 1, 2, -2, 1, 2, 0, 1, 2, 1, 0, 0, 2, 1, 2, 1, 0,
1543 2, 1, 0, 0, -1, -1, 1, -1, 0, 1, -1, 2, 2, 0, 0, -1, 1,
1544 1, 1, 1, 0, 0, -2, 0, -1, 1, 2, 0, 0, 1, 1, -1, 1, 0,
1545 1,
1546 });
1547 try testing.expectEqual(Poly.noise(2, 37, &seed).cs, .{
1548 1, 0, 1, -1, -1, -2, -1, -1, 2, 0, -1, 0, 0, -1,
1549 1, 1, -1, 1, 0, 2, -2, 0, 1, 2, 0, 0, -1, 1,
1550 0, -1, 1, -1, 1, 2, 1, 1, 0, -1, 1, -1, -2, -1,
1551 1, -1, -1, -1, 2, -1, -1, 0, 0, 1, 1, -1, 1, 1,
1552 1, 1, -1, -2, 0, 1, 0, 0, 2, 1, -1, 2, 0, 0,
1553 1, 1, 0, -1, 0, 0, -1, -1, 2, 0, 1, -1, 2, -1,
1554 -1, -1, -1, 0, -2, 0, 2, 1, 0, 0, 0, -1, 0, 0,
1555 0, -1, -1, 0, -1, -1, 0, -1, 0, 0, -2, 1, 1, 0,
1556 1, 0, 1, 0, 1, 1, -1, 2, 0, 1, -1, 1, 2, 0,
1557 0, 0, 0, -1, -1, -1, 0, 1, 0, -1, 2, 0, 0, 1,
1558 1, 1, 0, 1, -1, 1, 2, 1, 0, 2, -1, 1, -1, -2,
1559 -1, -2, -1, 1, 0, -2, -2, -1, 1, 0, 0, 0, 0, 1,
1560 0, 0, 0, 2, 2, 0, 1, 0, -1, -1, 0, 2, 0, 0,
1561 -2, 1, 0, 2, 1, -1, -2, 0, 0, -1, 1, 1, 0, 0,
1562 2, 0, 1, 1, -2, 1, -2, 1, 1, 0, 2, 0, -1, 0,
1563 -1, 0, 1, 2, 0, 1, 0, -2, 1, -2, -2, 1, -1, 0,
1564 -1, 1, 1, 0, 0, 0, 1, 0, -1, 1, 1, 0, 0, 0,
1565 0, 1, 0, 1, -1, 0, 1, -1, -1, 2, 0, 0, 1, -1,
1566 0, 1, -1, 0,
1567 });
1568}
1569
1570test "uniform sampling" {
1571 var seed: [32]u8 = undefined;
1572 for (&seed, 0..) |*s, i| {
1573 s.* = @as(u8, @intCast(i));
1574 }
1575 try testing.expectEqual(Poly.uniform(seed, 1, 0).cs, .{
1576 797, 993, 161, 6, 2608, 2385, 2096, 2661, 1676, 247, 2440,
1577 342, 634, 194, 1570, 2848, 986, 684, 3148, 3208, 2018, 351,
1578 2288, 612, 1394, 170, 1521, 3119, 58, 596, 2093, 1549, 409,
1579 2156, 1934, 1730, 1324, 388, 446, 418, 1719, 2202, 1812, 98,
1580 1019, 2369, 214, 2699, 28, 1523, 2824, 273, 402, 2899, 246,
1581 210, 1288, 863, 2708, 177, 3076, 349, 44, 949, 854, 1371,
1582 957, 292, 2502, 1617, 1501, 254, 7, 1761, 2581, 2206, 2655,
1583 1211, 629, 1274, 2358, 816, 2766, 2115, 2985, 1006, 2433, 856,
1584 2596, 3192, 1, 1378, 2345, 707, 1891, 1669, 536, 1221, 710,
1585 2511, 120, 1176, 322, 1897, 2309, 595, 2950, 1171, 801, 1848,
1586 695, 2912, 1396, 1931, 1775, 2904, 893, 2507, 1810, 2873, 253,
1587 1529, 1047, 2615, 1687, 831, 1414, 965, 3169, 1887, 753, 3246,
1588 1937, 115, 2953, 586, 545, 1621, 1667, 3187, 1654, 1988, 1857,
1589 512, 1239, 1219, 898, 3106, 391, 1331, 2228, 3169, 586, 2412,
1590 845, 768, 156, 662, 478, 1693, 2632, 573, 2434, 1671, 173,
1591 969, 364, 1663, 2701, 2169, 813, 1000, 1471, 720, 2431, 2530,
1592 3161, 733, 1691, 527, 2634, 335, 26, 2377, 1707, 767, 3020,
1593 950, 502, 426, 1138, 3208, 2607, 2389, 44, 1358, 1392, 2334,
1594 875, 2097, 173, 1697, 2578, 942, 1817, 974, 1165, 2853, 1958,
1595 2973, 3282, 271, 1236, 1677, 2230, 673, 1554, 96, 242, 1729,
1596 2518, 1884, 2272, 71, 1382, 924, 1807, 1610, 456, 1148, 2479,
1597 2152, 238, 2208, 2329, 713, 1175, 1196, 757, 1078, 3190, 3169,
1598 708, 3117, 154, 1751, 3225, 1364, 154, 23, 2842, 1105, 1419,
1599 79, 5, 2013,
1600 });
1601}
1602
1603test "Polynomial packing" {
1604 var rnd = RndGen.init(0);
1605
1606 for (0..1000) |_| {
1607 const p = randPolyNormalized(&rnd);
1608 try testing.expectEqual(Poly.fromBytes(&p.toBytes()), p);
1609 }
1610}
1611
1612test "Test inner PKE" {
1613 var seed: [32]u8 = undefined;
1614 var pt: [32]u8 = undefined;
1615 for (&seed, &pt, 0..) |*s, *p, i| {
1616 s.* = @as(u8, @intCast(i));
1617 p.* = @as(u8, @intCast(i + 32));
1618 }
1619 inline for (modes) |mode| {
1620 for (0..10) |i| {
1621 var pk: mode.InnerPk = undefined;
1622 var sk: mode.InnerSk = undefined;
1623 seed[0] = @as(u8, @intCast(i));
1624 mode.innerKeyFromSeed(seed, &pk, &sk);
1625 for (0..10) |j| {
1626 seed[1] = @as(u8, @intCast(j));
1627 try testing.expectEqual(sk.decrypt(&pk.encrypt(&pt, &seed)), pt);
1628 }
1629 }
1630 }
1631}
1632
1633test "Test happy flow" {
1634 var seed: [64]u8 = undefined;
1635 for (&seed, 0..) |*s, i| {
1636 s.* = @as(u8, @intCast(i));
1637 }
1638 inline for (modes) |mode| {
1639 for (0..10) |i| {
1640 seed[0] = @intCast(i);
1641 const kp = try mode.KeyPair.generateDeterministic(seed);
1642 const sk = try mode.SecretKey.fromBytes(&kp.secret_key.toBytes());
1643 try testing.expectEqual(sk, kp.secret_key);
1644 const pk = try mode.PublicKey.fromBytes(&kp.public_key.toBytes());
1645 try testing.expectEqual(pk, kp.public_key);
1646 for (0..10) |j| {
1647 seed[1] = @intCast(j);
1648 const e = pk.encapsDeterministic(seed[0..32]);
1649 try testing.expectEqual(e.shared_secret, try sk.decaps(&e.ciphertext));
1650 }
1651 }
1652 }
1653}
1654
1655// Code to test NIST Known Answer Tests (KAT), see PQCgenKAT.c.
1656
1657test "NIST KAT test d00.Kyber512" {
1658 try testNistKat(d00.Kyber512, "e9c2bd37133fcb40772f81559f14b1f58dccd1c816701be9ba6214d43baf4547");
1659}
1660
1661test "NIST KAT test d00.Kyber1024" {
1662 try testNistKat(d00.Kyber1024, "89248f2f33f7f4f7051729111f3049c409a933ec904aedadf035f30fa5646cd5");
1663}
1664
1665test "NIST KAT test d00.Kyber768" {
1666 try testNistKat(d00.Kyber768, "a1e122cad3c24bc51622e4c242d8b8acbcd3f618fee4220400605ca8f9ea02c2");
1667}
1668
1669fn testNistKat(mode: type, hash: []const u8) !void {
1670 var seed: [48]u8 = undefined;
1671 for (&seed, 0..) |*s, i| {
1672 s.* = @as(u8, @intCast(i));
1673 }
1674 var fw: std.Io.Writer.Hashing(crypto.hash.sha2.Sha256) = .init(&.{});
1675 var g = NistDRBG.init(seed);
1676 try fw.writer.print("# {s}\n\n", .{mode.name});
1677 for (0..100) |i| {
1678 g.fill(&seed);
1679 try fw.writer.print("count = {}\n", .{i});
1680 try fw.writer.print("seed = {X}\n", .{&seed});
1681 var g2 = NistDRBG.init(seed);
1682
1683 // This is not equivalent to g2.fill(kseed[:]). As the reference
1684 // implementation calls randombytes twice generating the keypair,
1685 // we have to do that as well.
1686 var kseed: [64]u8 = undefined;
1687 var eseed: [32]u8 = undefined;
1688 g2.fill(kseed[0..32]);
1689 g2.fill(kseed[32..64]);
1690 g2.fill(&eseed);
1691 const kp = try mode.KeyPair.generateDeterministic(kseed);
1692 const e = kp.public_key.encapsDeterministic(&eseed);
1693 const ss2 = try kp.secret_key.decaps(&e.ciphertext);
1694 try testing.expectEqual(ss2, e.shared_secret);
1695 try fw.writer.print("pk = {X}\n", .{&kp.public_key.toBytes()});
1696 try fw.writer.print("sk = {X}\n", .{&kp.secret_key.toBytes()});
1697 try fw.writer.print("ct = {X}\n", .{&e.ciphertext});
1698 try fw.writer.print("ss = {X}\n\n", .{&e.shared_secret});
1699 }
1700
1701 var out: [32]u8 = undefined;
1702 fw.hasher.final(&out);
1703 var outHex: [64]u8 = undefined;
1704 _ = try std.mem.print(&outHex, "{x}", .{&out});
1705 try testing.expectEqualStrings(&outHex, hash);
1706}
1707
1708const NistDRBG = struct {
1709 key: [32]u8,
1710 v: [16]u8,
1711
1712 fn incV(g: *NistDRBG) void {
1713 const val = std.mem.readInt(u128, &g.v, .big);
1714 std.mem.writeInt(u128, &g.v, val +% 1, .big);
1715 }
1716
1717 // AES256_CTR_DRBG_Update(pd, &g.key, &g.v).
1718 fn update(g: *NistDRBG, pd: ?[48]u8) void {
1719 var buf: [48]u8 = undefined;
1720 const ctx = crypto.core.aes.Aes256.initEnc(g.key);
1721 var i: usize = 0;
1722 while (i < 3) : (i += 1) {
1723 g.incV();
1724 var block: [16]u8 = undefined;
1725 ctx.encrypt(&block, &g.v);
1726 buf[i * 16 ..][0..16].* = block;
1727 }
1728 if (pd) |p| {
1729 for (&buf, p) |*b, x| {
1730 b.* ^= x;
1731 }
1732 }
1733 g.key = buf[0..32].*;
1734 g.v = buf[32..48].*;
1735 }
1736
1737 // randombytes.
1738 fn fill(g: *NistDRBG, out: []u8) void {
1739 var block: [16]u8 = undefined;
1740 var dst = out;
1741
1742 const ctx = crypto.core.aes.Aes256.initEnc(g.key);
1743 while (dst.len > 0) {
1744 g.incV();
1745 ctx.encrypt(&block, &g.v);
1746 if (dst.len < 16) {
1747 @memcpy(dst, block[0..dst.len]);
1748 break;
1749 }
1750 dst[0..block.len].* = block;
1751 dst = dst[16..dst.len];
1752 }
1753 g.update(null);
1754 }
1755
1756 fn init(seed: [48]u8) NistDRBG {
1757 var ret: NistDRBG = .{ .key = @splat(0), .v = @splat(0) };
1758 ret.update(seed);
1759 return ret;
1760 }
1761};
1762
1763/// Extended Euclidian Algorithm
1764/// Only meant to be used on comptime values; correctness matters, performance doesn't.
1765fn extendedEuclidean(comptime T: type, comptime a_: T, comptime b_: T) struct { gcd: T, x: T, y: T } {
1766 var a = a_;
1767 var b = b_;
1768 var x0: T = 1;
1769 var x1: T = 0;
1770 var y0: T = 0;
1771 var y1: T = 1;
1772
1773 while (b != 0) {
1774 const q = @divTrunc(a, b);
1775 const temp_a = a;
1776 a = b;
1777 b = temp_a - q * b;
1778
1779 const temp_x = x0;
1780 x0 = x1;
1781 x1 = temp_x - q * x1;
1782
1783 const temp_y = y0;
1784 y0 = y1;
1785 y1 = temp_y - q * y1;
1786 }
1787
1788 return .{ .gcd = a, .x = x0, .y = y0 };
1789}
1790
1791/// Modular inversion: computes a^(-1) mod p
1792/// Requires gcd(a,p) = 1. The result is normalized to the range [0, p).
1793fn modularInverse(comptime T: type, comptime a: T, comptime p: T) T {
1794 // Use a signed type for EEA computation
1795 const type_info = @typeInfo(T);
1796 const SignedT = if (type_info == .int and type_info.int.signedness == .unsigned)
1797 @Int(.signed, type_info.int.bits)
1798 else
1799 T;
1800
1801 const a_signed = @as(SignedT, @intCast(a));
1802 const p_signed = @as(SignedT, @intCast(p));
1803
1804 const r = extendedEuclidean(SignedT, a_signed, p_signed);
1805 assert(r.gcd == 1);
1806
1807 // Normalize result to [0, p)
1808 var result = r.x;
1809 while (result < 0) {
1810 result += p_signed;
1811 }
1812
1813 return @intCast(result);
1814}
1815
1816/// Modular exponentiation: computes a^s mod p using square-and-multiply algorithm.
1817fn modularPow(comptime T: type, comptime a: T, s: T, comptime p: T) T {
1818 const type_info = @typeInfo(T);
1819 const bits = type_info.int.bits;
1820 const WideT = @Int(.unsigned, bits * 2);
1821
1822 var ret: T = 1;
1823 var base: T = a;
1824 var exp = s;
1825
1826 while (exp > 0) {
1827 if (exp & 1 == 1) {
1828 ret = @intCast((@as(WideT, ret) * @as(WideT, base)) % p);
1829 }
1830 base = @intCast((@as(WideT, base) * @as(WideT, base)) % p);
1831 exp >>= 1;
1832 }
1833
1834 return ret;
1835}
1836
1837/// Creates an all-ones or all-zeros mask from a single bit value.
1838/// Returns all 1s (0xFF...FF) if bit == 1, all 0s if bit == 0.
1839fn bitMask(comptime T: type, bit: T) T {
1840 const type_info = @typeInfo(T);
1841 if (type_info != .int or type_info.int.signedness != .unsigned) {
1842 @compileError("bitMask requires an unsigned integer type");
1843 }
1844 return -%bit;
1845}
1846
1847/// Creates a mask from the sign bit of a signed integer.
1848/// Returns all 1s (0xFF...FF) if x < 0, all 0s if x >= 0.
1849fn signMask(comptime T: type, x: T) @Int(.unsigned, @typeInfo(T).int.bits) {
1850 const type_info = @typeInfo(T);
1851 if (type_info != .int) {
1852 @compileError("signMask requires an integer type");
1853 }
1854
1855 const bits = type_info.int.bits;
1856 const SignedT = @Int(.signed, bits);
1857
1858 // Convert to signed if needed, arithmetic right shift to propagate sign bit
1859 const x_signed: SignedT = if (type_info.int.signedness == .signed) x else @bitCast(x);
1860 const shifted = x_signed >> (bits - 1);
1861 return @bitCast(shifted);
1862}
1863
1864test "bitMask and signMask helpers" {
1865 try testing.expectEqual(@as(u32, 0x00000000), bitMask(u32, 0));
1866 try testing.expectEqual(@as(u32, 0xFFFFFFFF), bitMask(u32, 1));
1867 try testing.expectEqual(@as(u8, 0x00), bitMask(u8, 0));
1868 try testing.expectEqual(@as(u8, 0xFF), bitMask(u8, 1));
1869 try testing.expectEqual(@as(u64, 0x0000000000000000), bitMask(u64, 0));
1870 try testing.expectEqual(@as(u64, 0xFFFFFFFFFFFFFFFF), bitMask(u64, 1));
1871
1872 try testing.expectEqual(@as(u32, 0xFFFFFFFF), signMask(i32, -1));
1873 try testing.expectEqual(@as(u32, 0xFFFFFFFF), signMask(i32, -100));
1874 try testing.expectEqual(@as(u32, 0x00000000), signMask(i32, 0));
1875 try testing.expectEqual(@as(u32, 0x00000000), signMask(i32, 1));
1876 try testing.expectEqual(@as(u32, 0x00000000), signMask(i32, 100));
1877
1878 try testing.expectEqual(@as(u32, 0xFFFFFFFF), signMask(u32, 0x80000000)); // MSB set
1879 try testing.expectEqual(@as(u32, 0x00000000), signMask(u32, 0x7FFFFFFF)); // MSB clear
1880}
1881
1882/// Montgomery reduction: for input x, returns y where y ≡ x*R^(-1) (mod q).
1883/// This is a generic implementation parameterized by the modulus q, its inverse qInv,
1884/// the Montgomery constant R, and the result bound.
1885///
1886/// For ML-DSA: R = 2^32, returns y < 2q
1887/// For ML-KEM: R = 2^16, returns y in range (-q, q)
1888fn montgomeryReduce(
1889 comptime InT: type,
1890 comptime OutT: type,
1891 comptime q: comptime_int,
1892 comptime qInv: comptime_int,
1893 comptime r_bits: comptime_int,
1894 x: InT,
1895) OutT {
1896 const mask = (@as(InT, 1) << r_bits) - 1;
1897 const m_full = (x *% qInv) & mask;
1898 const m: OutT = @truncate(m_full);
1899
1900 const yR = x -% @as(InT, m) * @as(InT, q);
1901 const y_shifted = @as(@Int(.unsigned, @typeInfo(InT).Int.bits), @bitCast(yR)) >> r_bits;
1902 return @bitCast(@as(@Int(.unsigned, @typeInfo(OutT).Int.bits), @truncate(y_shifted)));
1903}
1904
1905/// Uniform sampling using SHAKE-128 with rejection sampling.
1906/// Samples polynomial coefficients uniformly from [0, q) using rejection sampling.
1907///
1908/// Parameters:
1909/// - PolyType: The polynomial type to return
1910/// - q: Modulus
1911/// - bits_per_coef: Number of bits per coefficient (12 or 23)
1912/// - n: Number of coefficients
1913/// - seed: Random seed
1914/// - domain_sep: Domain separation bytes (appended to seed)
1915fn sampleUniformRejection(
1916 comptime PolyType: type,
1917 comptime q: comptime_int,
1918 comptime bits_per_coef: comptime_int,
1919 comptime n: comptime_int,
1920 seed: []const u8,
1921 domain_sep: []const u8,
1922) PolyType {
1923 var h = sha3.Shake128.init(.{});
1924 h.update(seed);
1925 h.update(domain_sep);
1926
1927 const buf_len = sha3.Shake128.block_length; // 168 bytes
1928 var buf: [buf_len]u8 = undefined;
1929
1930 var ret: PolyType = undefined;
1931 var coef_idx: usize = 0;
1932
1933 if (bits_per_coef == 12) {
1934 // ML-KEM path: pack 2 coefficients per 3 bytes (12 bits each)
1935 outer: while (true) {
1936 h.squeeze(&buf);
1937
1938 var j: usize = 0;
1939 while (j < buf_len) : (j += 3) {
1940 const b0 = @as(u16, buf[j]);
1941 const b1 = @as(u16, buf[j + 1]);
1942 const b2 = @as(u16, buf[j + 2]);
1943
1944 const ts: [2]u16 = .{
1945 b0 | ((b1 & 0xf) << 8),
1946 (b1 >> 4) | (b2 << 4),
1947 };
1948
1949 inline for (ts) |t| {
1950 if (t < q) {
1951 ret.cs[coef_idx] = @intCast(t);
1952 coef_idx += 1;
1953 if (coef_idx == n) break :outer;
1954 }
1955 }
1956 }
1957 }
1958 } else if (bits_per_coef == 23) {
1959 // ML-DSA path: 1 coefficient per 3 bytes (23 bits)
1960 while (coef_idx < n) {
1961 h.squeeze(&buf);
1962
1963 var j: usize = 0;
1964 while (j < buf_len and coef_idx < n) : (j += 3) {
1965 const t = (@as(u32, buf[j]) |
1966 (@as(u32, buf[j + 1]) << 8) |
1967 (@as(u32, buf[j + 2]) << 16)) & 0x7fffff;
1968
1969 if (t < q) {
1970 ret.cs[coef_idx] = @intCast(t);
1971 coef_idx += 1;
1972 }
1973 }
1974 }
1975 } else {
1976 @compileError("bits_per_coef must be 12 or 23");
1977 }
1978
1979 return ret;
1980}