nx_kyber_poly_wasm.nx source
↩ module page · 239 lines · 8584 B
1// nx_kyber_poly_wasm.nx -- K-PKE polynomial helpers for FIPS 203 ML-KEM-768.
2//
3// Composes with the shipped nx_kyber_ntt (L117). Together these are
4// the L2 substrate transforms that nx_ml_kem_768.nx (L118 next) composes
5// the FIPS 203 KEM around.
6//
7// All polynomials are 256-coefficient over Z_{3329}. In-memory format:
8// packed little-endian INT16 at byte ptr (512 bytes per polynomial).
9// CALLER is responsible for canonical-form vs Montgomery-form bookkeeping
10// per FIPS 203 conventions (NTT-domain polynomials are in Mont form per
11// PQClean `_tomont`; canonical-form polynomials are NOT).
12//
13// Public API:
14// nx_kyber_poly_add(out, a, b) -> i64 out = a + b mod q
15// nx_kyber_poly_sub(out, a, b) -> i64 out = a - b mod q
16// nx_kyber_poly_cbd_eta2(out, buf) -> i64 centered binomial
17// sample from 128 bytes
18// nx_kyber_poly_compress10(bytes_out, poly) -> i64 320-byte d_u packing
19// nx_kyber_poly_decompress10(poly, bytes) -> i64 320-byte unpack
20// nx_kyber_poly_compress4(bytes_out, poly) -> i64 128-byte d_v packing
21// nx_kyber_poly_decompress4(poly, bytes) -> i64 128-byte unpack
22// nx_kyber_poly_tobytes12(bytes_out, poly) -> i64 384-byte canonical
23// nx_kyber_poly_frombytes12(poly, bytes) -> i64 384-byte parse
24//
25// Reductions: caller-provided polynomials should be in canonical [0, q).
26// poly_add/sub output coefficients in [0, 2q-2] / [-q+1, q-1]; caller
27// applies barrett_reduce (in NTT module) before passing further if needed.
28//
29// Verified: smoke runs each compress/decompress pair + tobytes/frombytes
30// round-trip; the deterministic encoding round-trips bit-exact.
31//
32// license_tier: INDEPENDENT_REDERIVE
33// genealogy_id: international-research-sources/nist/fips_203
34// lineage_id: nishi_kyber_poly_wasm_q1
35// safe_shift_audit: no 64-bit rotations in this module
36
37const KYBER_Q: i64 = 3329
38const KYBER_N: i64 = 256
39
40// === Packed-i16 polynomial I/O (same layout as nx_kyber_ntt_wasm.nx) ===
41
42func _pload(p: *u8, i: i64) -> i64 {
43 let lo: i64 = p[i * 2]
44 let hi: i64 = p[i * 2 + 1]
45 let raw: i64 = lo | (hi << 8)
46 if raw >= 32768 { return raw - 65536 }
47 return raw
48}
49
50func _pstore(p: *u8, i: i64, v: i64) -> i64 {
51 var vv: i64 = v
52 if vv < 0 { vv = vv + 65536 }
53 p[i * 2] = vv & 0xff
54 p[i * 2 + 1] = (vv >> 8) & 0xff
55 return 0
56}
57
58// Reduce coefficient into canonical [0, q).
59func _canon(v: i64) -> i64 {
60 var x: i64 = v % KYBER_Q
61 if x < 0 { x = x + KYBER_Q }
62 return x
63}
64
65// === poly_add / poly_sub ===
66
67func nx_kyber_poly_add(out: *u8, a: *u8, b: *u8) -> i64 {
68 var i: i64 = 0
69 while i < KYBER_N {
70 let s: i64 = _pload(a, i) + _pload(b, i)
71 _pstore(out, i, s)
72 i = i + 1
73 }
74 return 0
75}
76
77func nx_kyber_poly_sub(out: *u8, a: *u8, b: *u8) -> i64 {
78 var i: i64 = 0
79 while i < KYBER_N {
80 let s: i64 = _pload(a, i) - _pload(b, i)
81 _pstore(out, i, s)
82 i = i + 1
83 }
84 return 0
85}
86
87// === CBD (centered binomial distribution) with eta = 2 ===
88// Per FIPS 203 ยง4.2.2: each 4-bit chunk of input produces one coefficient
89// in {-2, -1, 0, 1, 2}. 64 bytes of random input -> 128 coefficients.
90// For 256 coefficients we consume 128 bytes total -> two halves.
91//
92// CBD_2(buf) for each 4-bit chunk b0..b3:
93// a = popcount(b0 b1) in {0, 1, 2}
94// b = popcount(b2 b3) in {0, 1, 2}
95// coefficient = a - b in {-2, -1, 0, 1, 2}
96func nx_kyber_poly_cbd_eta2(out: *u8, buf: *u8) -> i64 {
97 var i: i64 = 0
98 while i < KYBER_N {
99 let byte_idx: i64 = i >> 1
100 let upper: i64 = i & 1 // 0 -> low nibble, 1 -> high
101 let nibble: i64 = (buf[byte_idx] >> (upper * 4)) & 0xf
102 let b0: i64 = nibble & 1
103 let b1: i64 = (nibble >> 1) & 1
104 let b2: i64 = (nibble >> 2) & 1
105 let b3: i64 = (nibble >> 3) & 1
106 let a: i64 = b0 + b1
107 let bb: i64 = b2 + b3
108 _pstore(out, i, a - bb)
109 i = i + 1
110 }
111 return 0
112}
113
114// === poly_compress / decompress (d=10, used for ML-KEM-768 u-vector) ===
115//
116// FIPS 203 Compress_q(x, d) = round((2^d / q) * x) mod 2^d
117// for d=10: floor((x << 11) / q + 1) >> 1 then mask to 10 bits
118// Round-to-nearest implemented via the standard half-up trick.
119//
120// 256 coefficients * 10 bits = 2560 bits = 320 bytes per polynomial.
121// Bit-packing: 4 coefficients pack into 5 bytes (40 bits).
122func _compress10_one(v: i64) -> i64 {
123 let x: i64 = _canon(v)
124 // round((x << 10) / q) = ((x << 11) / q + 1) >> 1
125 return (((x << 11) / KYBER_Q + 1) >> 1) & 0x3ff
126}
127
128func nx_kyber_poly_compress10(bytes_out: *u8, poly: *u8) -> i64 {
129 var i: i64 = 0
130 while i < KYBER_N {
131 let c0: i64 = _compress10_one(_pload(poly, i + 0))
132 let c1: i64 = _compress10_one(_pload(poly, i + 1))
133 let c2: i64 = _compress10_one(_pload(poly, i + 2))
134 let c3: i64 = _compress10_one(_pload(poly, i + 3))
135 let off: i64 = (i >> 2) * 5
136 bytes_out[off + 0] = c0 & 0xff
137 bytes_out[off + 1] = ((c0 >> 8) | (c1 << 2)) & 0xff
138 bytes_out[off + 2] = ((c1 >> 6) | (c2 << 4)) & 0xff
139 bytes_out[off + 3] = ((c2 >> 4) | (c3 << 6)) & 0xff
140 bytes_out[off + 4] = (c3 >> 2) & 0xff
141 i = i + 4
142 }
143 return 0
144}
145
146// Decompress_q(x, d) = round((q / 2^d) * x) = floor((q * x + 2^(d-1)) / 2^d)
147func _decompress10_one(x: i64) -> i64 {
148 return (KYBER_Q * x + 512) >> 10
149}
150
151func nx_kyber_poly_decompress10(poly: *u8, bytes_in: *u8) -> i64 {
152 var i: i64 = 0
153 while i < KYBER_N {
154 let off: i64 = (i >> 2) * 5
155 let b0: i64 = bytes_in[off + 0]
156 let b1: i64 = bytes_in[off + 1]
157 let b2: i64 = bytes_in[off + 2]
158 let b3: i64 = bytes_in[off + 3]
159 let b4: i64 = bytes_in[off + 4]
160 let c0: i64 = b0 | ((b1 & 0x03) << 8)
161 let c1: i64 = (b1 >> 2) | ((b2 & 0x0f) << 6)
162 let c2: i64 = (b2 >> 4) | ((b3 & 0x3f) << 4)
163 let c3: i64 = (b3 >> 6) | (b4 << 2)
164 _pstore(poly, i + 0, _decompress10_one(c0 & 0x3ff))
165 _pstore(poly, i + 1, _decompress10_one(c1 & 0x3ff))
166 _pstore(poly, i + 2, _decompress10_one(c2 & 0x3ff))
167 _pstore(poly, i + 3, _decompress10_one(c3 & 0x3ff))
168 i = i + 4
169 }
170 return 0
171}
172
173// === poly_compress / decompress (d=4, used for ML-KEM-768 v scalar) ===
174// 256 coefficients * 4 bits = 1024 bits = 128 bytes per polynomial.
175// 2 coefficients per byte.
176func _compress4_one(v: i64) -> i64 {
177 let x: i64 = _canon(v)
178 return (((x << 5) / KYBER_Q + 1) >> 1) & 0xf
179}
180
181func nx_kyber_poly_compress4(bytes_out: *u8, poly: *u8) -> i64 {
182 var i: i64 = 0
183 while i < KYBER_N {
184 let c0: i64 = _compress4_one(_pload(poly, i + 0))
185 let c1: i64 = _compress4_one(_pload(poly, i + 1))
186 bytes_out[i >> 1] = (c0 | (c1 << 4)) & 0xff
187 i = i + 2
188 }
189 return 0
190}
191
192func _decompress4_one(x: i64) -> i64 {
193 return (KYBER_Q * x + 8) >> 4
194}
195
196func nx_kyber_poly_decompress4(poly: *u8, bytes_in: *u8) -> i64 {
197 var i: i64 = 0
198 while i < KYBER_N {
199 let b: i64 = bytes_in[i >> 1]
200 _pstore(poly, i + 0, _decompress4_one(b & 0xf))
201 _pstore(poly, i + 1, _decompress4_one((b >> 4) & 0xf))
202 i = i + 2
203 }
204 return 0
205}
206
207// === poly_tobytes12 / frombytes12 (canonical 12-bit encoding for pk) ===
208//
209// 256 coefficients * 12 bits = 3072 bits = 384 bytes per polynomial.
210// 2 coefficients pack into 3 bytes (24 bits).
211func nx_kyber_poly_tobytes12(bytes_out: *u8, poly: *u8) -> i64 {
212 var i: i64 = 0
213 while i < KYBER_N {
214 let c0: i64 = _canon(_pload(poly, i + 0)) & 0xfff
215 let c1: i64 = _canon(_pload(poly, i + 1)) & 0xfff
216 let off: i64 = (i >> 1) * 3
217 bytes_out[off + 0] = c0 & 0xff
218 bytes_out[off + 1] = ((c0 >> 8) | (c1 << 4)) & 0xff
219 bytes_out[off + 2] = (c1 >> 4) & 0xff
220 i = i + 2
221 }
222 return 0
223}
224
225func nx_kyber_poly_frombytes12(poly: *u8, bytes_in: *u8) -> i64 {
226 var i: i64 = 0
227 while i < KYBER_N {
228 let off: i64 = (i >> 1) * 3
229 let b0: i64 = bytes_in[off + 0]
230 let b1: i64 = bytes_in[off + 1]
231 let b2: i64 = bytes_in[off + 2]
232 let c0: i64 = b0 | ((b1 & 0x0f) << 8)
233 let c1: i64 = (b1 >> 4) | (b2 << 4)
234 _pstore(poly, i + 0, c0 & 0xfff)
235 _pstore(poly, i + 1, c1 & 0xfff)
236 i = i + 2
237 }
238 return 0
239}