nx_kyber_ntt_wasm.nx source
↩ module page · 205 lines · 9674 B
1// nx_kyber_ntt_wasm.nx -- NTT in Z_q (q = 3329) for FIPS 203 ML-KEM.
2//
3// Self-contained for WAT target. Polynomial = 256 INT16 coefficients
4// packed little-endian (512 bytes at poly_ptr). Coefficients live in
5// the range [-q/2, q/2] in Montgomery form during NTT operations and
6// in canonical [0, q) form before/after.
7//
8// API:
9// nx_kyber_ntt(poly_ptr) -> i64
10// nx_kyber_invntt(poly_ptr) -> i64
11// nx_kyber_basemul(out, a, b, zeta_idx) -> i64
12// (one pair of basemul; caller iterates 128 times)
13//
14// KAT: self-inverse round-trip + simple known-poly pointwise mul check
15// (smoke test, not on-page).
16//
17// Bricks for: FIPS 203 ML-KEM-768 (Kyber). Composed with SHA-3-256
18// (hash function H), SHAKE128 (matrix-A expansion), SHAKE256 (noise
19// sampling), nx_hmac_drbg (entropy mixing).
20//
21// Reference: PQClean public-domain Kyber NTT, derived from Cooley-Tukey
22// butterfly over the negacyclic ring Z_q[X]/(X^256 + 1).
23//
24// license_tier: INDEPENDENT_REDERIVE
25// genealogy_id: international-research-sources/nist/fips_203
26// lineage_id: nishi_kyber_ntt_wasm_q1
27// safe_shift_audit: this module does no 64-bit rotations
28
29const KYBER_Q: i64 = 3329
30const KYBER_QINV: i64 = 62209 // -3327 mod 2^16 = 62209
31const KYBER_N: i64 = 256
32const KYBER_MONT: i64 = 2285 // 2^16 mod q
33
34// === Packed i16 polynomial I/O ===
35// Each coefficient is a little-endian signed 16-bit at offset 2*i.
36func _poly_load(p: *u8, i: i64) -> i64 {
37 let lo: i64 = p[i * 2]
38 let hi: i64 = p[i * 2 + 1]
39 let raw: i64 = lo | (hi << 8)
40 if raw >= 32768 { return raw - 65536 }
41 return raw
42}
43
44func _poly_store(p: *u8, i: i64, v: i64) -> i64 {
45 var vv: i64 = v
46 if vv < 0 { vv = vv + 65536 }
47 p[i * 2] = vv & 0xff
48 p[i * 2 + 1] = (vv >> 8) & 0xff
49 return 0
50}
51
52// === Reductions ===
53
54// Montgomery reduce: input a in [-2^15 * q, 2^15 * q),
55// output in (-q, q) representing a * R^{-1} mod q where R = 2^16.
56//
57// Algorithm (PQClean):
58// u = (a * QINV) mod 2^16 treated as signed int16
59// t = (a - u * q) >> 16 arithmetic shift
60// return t
61func _montgomery_reduce(a: i64) -> i64 {
62 var u: i64 = (a * KYBER_QINV) & 0xffff
63 if u >= 32768 { u = u - 65536 }
64 let t: i64 = (a - u * KYBER_Q) >> 16
65 return t
66}
67
68// Barrett reduce: input a, output in [-q/2, q/2] representing a mod q.
69// Algorithm: v = ((1<<26) + q/2) / q = 20159 (Kyber constant)
70// t = (v * a + 2^25) >> 26
71// return a - t * q
72func _barrett_reduce(a: i64) -> i64 {
73 let v: i64 = 20159
74 let t: i64 = (v * a + 33554432) >> 26
75 return a - t * KYBER_Q
76}
77
78// Fast modular multiply a * b mod q with Montgomery reduction.
79// Used inside butterflies where one operand is a Montgomery-form zeta.
80func _fqmul(a: i64, b: i64) -> i64 {
81 return _montgomery_reduce(a * b)
82}
83
84// === Kyber zeta table (FIPS 203 / PQClean reference) ===
85// 128 precomputed Montgomery-form powers of zeta = 17 mod q.
86// Values are signed int16 in (-q/2, q/2).
87func _zetas(i: i64) -> i64 {
88 if i == 0 { return -1044 } if i == 1 { return -758 } if i == 2 { return -359 } if i == 3 { return -1517 }
89 if i == 4 { return 1493 } if i == 5 { return 1422 } if i == 6 { return 287 } if i == 7 { return 202 }
90 if i == 8 { return -171 } if i == 9 { return 622 } if i == 10 { return 1577 } if i == 11 { return 182 }
91 if i == 12 { return 962 } if i == 13 { return -1202 } if i == 14 { return -1474 } if i == 15 { return 1468 }
92 if i == 16 { return 573 } if i == 17 { return -1325 } if i == 18 { return 264 } if i == 19 { return 383 }
93 if i == 20 { return -829 } if i == 21 { return 1458 } if i == 22 { return -1602 } if i == 23 { return -130 }
94 if i == 24 { return -681 } if i == 25 { return 1017 } if i == 26 { return 732 } if i == 27 { return 608 }
95 if i == 28 { return -1542 } if i == 29 { return 411 } if i == 30 { return -205 } if i == 31 { return -1571 }
96 if i == 32 { return 1223 } if i == 33 { return 652 } if i == 34 { return -552 } if i == 35 { return 1015 }
97 if i == 36 { return -1293 } if i == 37 { return 1491 } if i == 38 { return -282 } if i == 39 { return -1544 }
98 if i == 40 { return 516 } if i == 41 { return -8 } if i == 42 { return -320 } if i == 43 { return -666 }
99 if i == 44 { return -1618 } if i == 45 { return -1162 } if i == 46 { return 126 } if i == 47 { return 1469 }
100 if i == 48 { return -853 } if i == 49 { return -90 } if i == 50 { return -271 } if i == 51 { return 830 }
101 if i == 52 { return 107 } if i == 53 { return -1421 } if i == 54 { return -247 } if i == 55 { return -951 }
102 if i == 56 { return -398 } if i == 57 { return 961 } if i == 58 { return -1508 } if i == 59 { return -725 }
103 if i == 60 { return 448 } if i == 61 { return -1065 } if i == 62 { return 677 } if i == 63 { return -1275 }
104 if i == 64 { return -1103 } if i == 65 { return 430 } if i == 66 { return 555 } if i == 67 { return 843 }
105 if i == 68 { return -1251 } if i == 69 { return 871 } if i == 70 { return 1550 } if i == 71 { return 105 }
106 if i == 72 { return 422 } if i == 73 { return 587 } if i == 74 { return 177 } if i == 75 { return -235 }
107 if i == 76 { return -291 } if i == 77 { return -460 } if i == 78 { return 1574 } if i == 79 { return 1653 }
108 if i == 80 { return -246 } if i == 81 { return 778 } if i == 82 { return 1159 } if i == 83 { return -147 }
109 if i == 84 { return -777 } if i == 85 { return 1483 } if i == 86 { return -602 } if i == 87 { return 1119 }
110 if i == 88 { return -1590 } if i == 89 { return 644 } if i == 90 { return -872 } if i == 91 { return 349 }
111 if i == 92 { return 418 } if i == 93 { return 329 } if i == 94 { return -156 } if i == 95 { return -75 }
112 if i == 96 { return 817 } if i == 97 { return 1097 } if i == 98 { return 603 } if i == 99 { return 610 }
113 if i == 100 { return 1322 } if i == 101 { return -1285 } if i == 102 { return -1465 } if i == 103 { return 384 }
114 if i == 104 { return -1215 } if i == 105 { return -136 } if i == 106 { return 1218 } if i == 107 { return -1335 }
115 if i == 108 { return -874 } if i == 109 { return 220 } if i == 110 { return -1187 } if i == 111 { return -1659 }
116 if i == 112 { return -1185 } if i == 113 { return -1530 } if i == 114 { return -1278 } if i == 115 { return 794 }
117 if i == 116 { return -1510 } if i == 117 { return -854 } if i == 118 { return -870 } if i == 119 { return 478 }
118 if i == 120 { return -108 } if i == 121 { return -308 } if i == 122 { return 996 } if i == 123 { return 991 }
119 if i == 124 { return 958 } if i == 125 { return -1460 } if i == 126 { return 1522 }
120 return 1628
121}
122
123// === Forward NTT ===
124// In-place over 256 packed-i16 coefficients.
125// Cooley-Tukey decimation-in-time over the 8-layer butterfly tree.
126func nx_kyber_ntt(poly: *u8) -> i64 {
127 var k: i64 = 1
128 var len: i64 = 128
129 while len >= 2 {
130 var start: i64 = 0
131 while start < KYBER_N {
132 let zeta: i64 = _zetas(k)
133 k = k + 1
134 var j: i64 = start
135 while j < start + len {
136 let aj: i64 = _poly_load(poly, j)
137 let ajl: i64 = _poly_load(poly, j + len)
138 let t: i64 = _fqmul(zeta, ajl)
139 _poly_store(poly, j + len, aj - t)
140 _poly_store(poly, j, aj + t)
141 j = j + 1
142 }
143 start = j + len
144 }
145 len = len >> 1
146 }
147 return 0
148}
149
150// === Inverse NTT ===
151// In-place; mirror of forward with reverse zeta walk + final Montgomery
152// scale by f = 1441 = mont^2 * 128^{-1} mod q (which composes the
153// constant Montgomery factor + the 1/N normalisation).
154func nx_kyber_invntt(poly: *u8) -> i64 {
155 let f: i64 = 1441
156 var k: i64 = 127
157 var len: i64 = 2
158 while len <= 128 {
159 var start: i64 = 0
160 while start < KYBER_N {
161 let zeta: i64 = _zetas(k)
162 k = k - 1
163 var j: i64 = start
164 while j < start + len {
165 let aj: i64 = _poly_load(poly, j)
166 let ajl: i64 = _poly_load(poly, j + len)
167 _poly_store(poly, j, _barrett_reduce(aj + ajl))
168 let diff: i64 = ajl - aj
169 _poly_store(poly, j + len, _fqmul(zeta, diff))
170 j = j + 1
171 }
172 start = j + len
173 }
174 len = len << 1
175 }
176 // Final scale by f for every coefficient.
177 var i: i64 = 0
178 while i < KYBER_N {
179 let v: i64 = _poly_load(poly, i)
180 _poly_store(poly, i, _fqmul(v, f))
181 i = i + 1
182 }
183 return 0
184}
185
186// === Pointwise basemul in NTT domain ===
187// Kyber NTT domain is non-trivial: pairs of coefficients (2i, 2i+1)
188// represent linear polynomials a + b*X over Z_q[X]/(X^2 - zeta_pair).
189// basemul multiplies one such pair-of-pairs by another.
190//
191// r[0] = a[0]*b[0] + zeta_pair * a[1]*b[1]
192// r[1] = a[0]*b[1] + a[1]*b[0]
193//
194// caller passes zeta_pair (one of _zetas(64+i)) and per-pair pointers.
195func nx_kyber_basemul(out: *u8, a: *u8, b: *u8, zeta: i64) -> i64 {
196 let a0: i64 = _poly_load(a, 0)
197 let a1: i64 = _poly_load(a, 1)
198 let b0: i64 = _poly_load(b, 0)
199 let b1: i64 = _poly_load(b, 1)
200 let r0: i64 = _fqmul(_fqmul(a1, b1), zeta) + _fqmul(a0, b0)
201 let r1: i64 = _fqmul(a0, b1) + _fqmul(a1, b0)
202 _poly_store(out, 0, r0)
203 _poly_store(out, 1, r1)
204 return 0
205}