| ... | @@ -105,19 +105,20 @@ const crypto = std.crypto; | ... | @@ -105,19 +105,20 @@ const crypto = std.crypto; |
| 105 | const errors = std.crypto.errors; | 105 | const errors = std.crypto.errors; |
| 106 | const math = std.math; | 106 | const math = std.math; |
| 107 | const mem = std.mem; | 107 | const mem = std.mem; |
| 108 | const RndGen = std.Random.DefaultPrng; | | |
| 109 | const sha3 = crypto.hash.sha3; | 108 | const sha3 = crypto.hash.sha3; |
| 110 | | 109 | |
| 111 | // Q is the parameter q ≡ 3329 = 2¹¹ + 2¹⁰ + 2⁸ + 1. | 110 | const RndGen = std.Random.DefaultPrng; |
| | 111 | |
| | 112 | // Q is the modulus q ≡ 3329 = 2¹¹ + 2¹⁰ + 2⁸ + 1 |
| 112 | const Q: i16 = 3329; | 113 | const Q: i16 = 3329; |
| 113 | | 114 | |
| 114 | // Montgomery R | 115 | // Montgomery R = 2^16 mod Q (for Montgomery multiplication) |
| 115 | const R: i32 = 1 << 16; | 116 | const R: i32 = 1 << 16; |
| 116 | | 117 | |
| 117 | // Parameter n, degree of polynomials. | 118 | // N is the degree of polynomials (polynomial ring dimension) |
| 118 | const N: usize = 256; | 119 | const N: usize = 256; |
| 119 | | 120 | |
| 120 | // Size of "small" vectors used in encryption blinds. | 121 | // eta2 is the size of "small" vectors used in encryption blinds |
| 121 | const eta2: u8 = 2; | 122 | const eta2: u8 = 2; |
| 122 | | 123 | |
| 123 | const Params = struct { | 124 | const 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); |
| 216 | | 217 | |
| 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); |
| 220 | | 221 | |
| 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) |
| 242 | | 243 | |
| 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; |
| 245 | | 246 | |
| 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 | } |
| 290 | | 291 | |
| 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 | } |
| 295 | | 296 | |
| 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, |
| 311 | | 312 | |
| 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; |
| 315 | | 316 | |
| 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 | } |
| 347 | | 348 | |
| 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 | } |
| 352 | | 353 | |
| 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 values | 419 | // Cached values |
| 419 | aT: M, | 420 | aT: M, |
| 420 | | 421 | |
| 421 | const bytes_length = V.bytes_length + 32; | 422 | const encoded_length = V.encoded_length + 32; |
| 422 | | 423 | |
| 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¹⁵q | 438 | // 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 | } |
| 441 | | 442 | |
| 442 | // Aᵀ and r were not in Montgomery form, so the Montgomery | 443 | // 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 | } |
| 453 | | 454 | |
| 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 | } |
| 457 | | 458 | |
| 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; |
| 460 | | 461 | |
| 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(); |
| 463 | | 464 | |
| 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 | } |
| 470 | | 471 | |
| 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 PKE | 478 | // Private key of the inner PKE |
| 478 | const InnerSk = struct { | 479 | const InnerSk = struct { |
| 479 | sh: V, // NTT(s), normalized | 480 | sh: V, // NTT(s), normalized |
| 480 | const bytes_length = V.bytes_length; | 481 | const encoded_length = V.encoded_length; |
| 481 | | 482 | |
| 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 | } |
| 493 | | 494 | |
| 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 | } |
| 497 | | 498 | |
| 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(); |
| 518 | | 519 | |
| 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; |
| 521 | | 522 | |
| 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⁻¹ which | 529 | // multiplications in the inner product added a factor R⁻¹ which |
| 529 | // we'll cancel out with toMont(). This will also ensure the | 530 | // 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 | } |
| 533 | | 534 | |
| 534 | pk.th = th.add(eh).normalize(); // bounded by 8q | 535 | 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.pdf | 568 | // This is actually optimal, as proven in https://eprint.iacr.org/2020/1377.pdf |
| 568 | // TODO generate comptime? | | |
| 569 | const inv_ntt_reductions = [_]i16{ | 569 | const inv_ntt_reductions = [_]i16{ |
| 570 | -1, // after layer 1 | 570 | -1, // after layer 1 |
| 571 | -1, // after layer 2 | 571 | -1, // after layer 2 |
| ... | @@ -634,31 +634,8 @@ test "invNTTReductions bounds" { | ... | @@ -634,31 +634,8 @@ test "invNTTReductions bounds" { |
| 634 | } | 634 | } |
| 635 | } | 635 | } |
| 636 | | 636 | |
| 637 | // Extended euclidean algorithm. | | |
| 638 | // | | |
| 639 | // For a, b finds x, y such that x a + y b = gcd(a, b). Used to compute | | |
| 640 | // modular inverse. | | |
| 641 | fn eea(a: anytype, b: @TypeOf(a)) EeaResult(@TypeOf(a)) { | | |
| 642 | if (a == 0) { | | |
| 643 | return .{ .gcd = b, .x = 0, .y = 1 }; | | |
| 644 | } | | |
| 645 | const r = eea(@rem(b, a), a); | | |
| 646 | return .{ .gcd = r.gcd, .x = r.y - @divTrunc(b, a) * r.x, .y = r.x }; | | |
| 647 | } | | |
| 648 | | | |
| 649 | fn EeaResult(comptime T: type) type { | | |
| 650 | return struct { gcd: T, x: T, y: T }; | | |
| 651 | } | | |
| 652 | | | |
| 653 | // Returns least common multiple of a and b. | | |
| 654 | fn lcm(a: anytype, b: @TypeOf(a)) @TypeOf(a) { | | |
| 655 | const r = eea(a, b); | | |
| 656 | return a * b / r.gcd; | | |
| 657 | } | | |
| 658 | | | |
| 659 | // Invert modulo p. | | |
| 660 | fn invertMod(a: anytype, p: @TypeOf(a)) @TypeOf(a) { | 637 | fn 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 | } |
| 790 | | 767 | |
| 791 | // Compute a^s mod p. | | |
| 792 | fn 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. |
| 811 | fn computeZetas() [128]i16 { | 769 | fn 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 { |
| 828 | const Poly = struct { | 786 | const Poly = struct { |
| 829 | cs: [N]i16, | 787 | cs: [N]i16, |
| 830 | | 788 | |
| 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 }; |
| 833 | | 791 | |
| | 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 | } |
| 841 | | 800 | |
| | 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 | } |
| 849 | | 809 | |
| 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 resulting | 812 | // 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; |
| 1056 | | 997 | |
| 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; |
| 1060 | | 1001 | |
| ... | @@ -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; |
| 1120 | | 1061 | |
| 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; |
| 1124 | | 1065 | |
| ... | @@ -1275,53 +1216,23 @@ const Poly = struct { | ... | @@ -1275,53 +1216,23 @@ const Poly = struct { |
| 1275 | return ret; | 1216 | return ret; |
| 1276 | } | 1217 | } |
| 1277 | | 1218 | |
| 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, |
| 1284 | | 1224 | 12, |
| 1285 | const buf_len = sha3.Shake128.block_length; // rate SHAKE-128 | 1225 | N, |
| 1286 | var buf: [buf_len]u8 = undefined; | 1226 | &seed, |
| 1287 | | 1227 | &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 | } |
| 1319 | | 1230 | |
| 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 | }; |
| 1350 | | 1261 | |
| 1351 | // A vector of K polynomials. | 1262 | // A vector of k polynomials. |
| 1352 | fn Vec(comptime K: u8) type { | 1263 | fn PolyVec(comptime k: u8) type { |
| 1353 | return struct { | 1264 | return struct { |
| 1354 | ps: [K]Poly, | 1265 | ps: [k]Poly, |
| 1355 | | 1266 | |
| 1356 | const Self = @This(); | 1267 | const Self = @This(); |
| 1357 | const bytes_length = K * Poly.bytes_length; | 1268 | const encoded_length = k * Poly.encoded_length; |
| 1358 | | 1269 | |
| 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 | } |
| 1362 | | 1273 | |
| 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 | } |
| 1370 | | 1282 | |
| 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 | } |
| 1378 | | 1291 | |
| 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 | } |
| 1386 | | 1295 | |
| 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 | } |
| 1394 | | 1307 | |
| 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 | } |
| 1402 | | 1311 | |
| 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 | } |
| 1410 | | 1315 | |
| 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 | } |
| 1454 | | 1359 | |
| 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 | } |
| 1463 | | 1368 | |
| 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 | } |
| 1476 | | 1381 | |
| 1477 | // A matrix of K vectors | 1382 | // A matrix of k vectors |
| 1478 | fn Mat(comptime K: u8) type { | 1383 | fn 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), |
| 1482 | | 1387 | |
| 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 A | 1404 | // 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 | } |
| 1524 | | 1429 | |
| | 1430 | // Test helper: generates a random polynomial with each coefficient |x| ≤ q |
| | 1431 | fn 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 |
| | 1440 | fn 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 | |
| 1525 | test "MulHat" { | 1448 | test "MulHat" { |
| 1526 | var rnd = RndGen.init(0); | 1449 | var rnd = RndGen.init(0); |
| 1527 | | 1450 | |
| 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); |
| 1531 | | 1454 | |
| 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); |
| 1558 | | 1481 | |
| 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(); |
| 1563 | | 1486 | |
| ... | @@ -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); |
| 1672 | | 1595 | |
| 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. |
| | 1768 | fn 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). |
| | 1796 | fn 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. |
| | 1820 | fn 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. |
| | 1842 | fn 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. |
| | 1852 | fn 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 | |
| | 1867 | test "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) |
| | 1891 | fn 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) |
| | 1918 | fn 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 | } |