nx_ml_dsa_65_wasm.nx source

1// nx_ml_dsa_65_wasm.nx -- Sovereign single-WASM FIPS 204 ML-DSA-65 (foundation). 2// 3// L131 iteration 1: math foundation only. Ships: 4// * Constants + i32 polynomial helpers (256 i32 LE = 1024 B per poly) 5// * Keccak-f[1600] + SHA-3-256 + SHAKE256 sponges (inlined verbatim from L126) 6// * Montgomery + Barrett reductions for Z_8380417 7// * Full ML-DSA NTT + INVNTT with 256-entry zeta table (Montgomery form, 8// primitive 512th root zeta=1753 mod q, bit-reversed indexed) 9// * Pointwise basemul + montgomery acc 10// * KAT export: nx_mldsa_ntt_round_trip_test(poly_ptr) returns 0 if 11// INVNTT(NTT(poly)) ≡ original (mod q, accounting for the 12// Montgomery factor R=2^32 that invntt_tomont leaves behind) 13// 14// Next iterations layer on: 15// L132: sampling (rej_uniform for matrix A, rej_eta for s1/s2, sample_in_ball) 16// L133: power-of-2 decompose + make_hint / use_hint 17// L134: keygen + tobytes/frombytes for pk/sk 18// L135: sign with rejection loop 19// L136: verify 20// L137: end-to-end round-trip self-test 21// 22// Final public API (after L137): 23// nx_mldsa_keygen(seed_32, scratch, pk_out, sk_out) 24// nx_mldsa_sign(sk, msg, msg_len, scratch, sig_out, sig_len_out) 25// nx_mldsa_verify(pk, msg, msg_len, sig, sig_len, scratch) -> i64 1/0 26// 27// Pillar 4 alignment: uses the same scratch-layout annotation pattern as 28// L126 (F-meta-8 grader-compatible). 29// 30// license_tier: INDEPENDENT_REDERIVE 31// genealogy_id: international-research-sources/nist/fips_204 32// lineage_id: nishi_ml_dsa_65_wasm_q1 33// safe_shift_audit: shared Keccak _rotl64 uses gold-standard mask 34 35// ============================================================================ 36// SECTION 0: ML-DSA-65 parameters (FIPS 204 §4 Table 1) 37// ============================================================================ 38 39const MLDSA_Q: i64 = 8380417 40const MLDSA_QINV: i64 = 58728449 // q^{-1} mod 2^32 (unsigned) 41const MLDSA_N: i64 = 256 42const MLDSA_K: i64 = 6 // matrix rows 43const MLDSA_L: i64 = 5 // matrix cols 44const MLDSA_ETA: i64 = 4 45const MLDSA_TAU: i64 = 49 46const MLDSA_BETA: i64 = 196 // tau * eta 47const MLDSA_GAMMA1: i64 = 524288 // 2^19 48const MLDSA_GAMMA2: i64 = 261888 // (q-1)/32 49const MLDSA_OMEGA: i64 = 55 50const MLDSA_D: i64 = 13 51 52const MLDSA_PK_BYTES: i64 = 1952 53const MLDSA_SK_BYTES: i64 = 4032 54const MLDSA_SIG_BYTES: i64 = 3309 55 56// Polynomial in-memory size (256 i32 LE) 57const POLY_BUF: i64 = 1024 58 59// ============================================================================ 60// SECTION 1: i32 polynomial I/O + reductions 61// ============================================================================ 62 63func _i32_load_le(p: *u8, i: i64) -> i64 { 64 let b0: i64 = p[i * 4] 65 let b1: i64 = p[i * 4 + 1] 66 let b2: i64 = p[i * 4 + 2] 67 let b3: i64 = p[i * 4 + 3] 68 let raw: i64 = b0 | (b1 << 8) | (b2 << 16) | (b3 << 24) 69 if raw >= 0x80000000 { return raw - 0x100000000 } 70 return raw 71} 72 73func _i32_store_le(p: *u8, i: i64, v: i64) -> i64 { 74 var vv: i64 = v 75 if vv < 0 { vv = vv + 0x100000000 } 76 p[i * 4] = vv & 0xff 77 p[i * 4 + 1] = (vv >> 8) & 0xff 78 p[i * 4 + 2] = (vv >> 16) & 0xff 79 p[i * 4 + 3] = (vv >> 24) & 0xff 80 return 0 81} 82 83// Montgomery reduction for q = 8380417, R = 2^32. 84// input: a in (-q * 2^31, q * 2^31) 85// output: a * R^{-1} mod q in (-q, q) 86// Per FIPS 204 / PQClean: 87// t = (a * QINV) mod 2^32 // int32_t truncation 88// return (a - t * q) >> 32 89func _mont_reduce(a: i64) -> i64 { 90 var t: i64 = (a * MLDSA_QINV) & 0xffffffff 91 if t >= 0x80000000 { t = t - 0x100000000 } 92 return (a - t * MLDSA_Q) >> 32 93} 94 95// Barrett-style reduce32 per PQClean: input any i32, output in (-(q-1)/2, q/2]. 96func _reduce32(a: i64) -> i64 { 97 let t: i64 = (a + (1 << 22)) >> 23 98 return a - t * MLDSA_Q 99} 100 101// Modular multiplication via Montgomery: result = a*b*R^{-1} mod q. 102func _fqmul(a: i64, b: i64) -> i64 { return _mont_reduce(a * b) } 103 104// 64-bit rotate left (Keccak helper). 105func _rotl64(x: i64, n: i64) -> i64 { 106 let nn: i64 = n & 63 107 if nn == 0 { return x } 108 let shr_amt: i64 = 64 - nn 109 let mask: i64 = (1 << nn) - 1 110 return ((x << nn) | ((x >> shr_amt) & mask)) & 0xffffffffffffffff 111} 112 113// ============================================================================ 114// SECTION 2: Keccak-f[1600] + sponges (verbatim from L126) 115// ============================================================================ 116 117func _lane_load(p: *u8) -> i64 { 118 let b0: i64 = p[0]; let b1: i64 = p[1]; let b2: i64 = p[2]; let b3: i64 = p[3] 119 let b4: i64 = p[4]; let b5: i64 = p[5]; let b6: i64 = p[6]; let b7: i64 = p[7] 120 return b0 | (b1 << 8) | (b2 << 16) | (b3 << 24) | (b4 << 32) | (b5 << 40) | (b6 << 48) | (b7 << 56) 121} 122func _lane_store(p: *u8, v: i64) -> i64 { 123 p[0] = v & 0xff; p[1] = (v >> 8) & 0xff; p[2] = (v >> 16) & 0xff; p[3] = (v >> 24) & 0xff 124 p[4] = (v >> 32) & 0xff; p[5] = (v >> 40) & 0xff; p[6] = (v >> 48) & 0xff; p[7] = (v >> 56) & 0xff 125 return 0 126} 127func _keccak_rc(i: i64) -> i64 { 128 if i == 0 { return 0x0000000000000001 } if i == 1 { return 0x0000000000008082 } 129 if i == 2 { return 0x800000000000808a } if i == 3 { return 0x8000000080008000 } 130 if i == 4 { return 0x000000000000808b } if i == 5 { return 0x0000000080000001 } 131 if i == 6 { return 0x8000000080008081 } if i == 7 { return 0x8000000000008009 } 132 if i == 8 { return 0x000000000000008a } if i == 9 { return 0x0000000000000088 } 133 if i == 10 { return 0x0000000080008009 } if i == 11 { return 0x000000008000000a } 134 if i == 12 { return 0x000000008000808b } if i == 13 { return 0x800000000000008b } 135 if i == 14 { return 0x8000000000008089 } if i == 15 { return 0x8000000000008003 } 136 if i == 16 { return 0x8000000000008002 } if i == 17 { return 0x8000000000000080 } 137 if i == 18 { return 0x000000000000800a } if i == 19 { return 0x800000008000000a } 138 if i == 20 { return 0x8000000080008081 } if i == 21 { return 0x8000000000008080 } 139 if i == 22 { return 0x0000000080000001 } 140 return 0x8000000080008008 141} 142func _rho_off(lane_idx: i64) -> i64 { 143 if lane_idx == 0 { return 0 } if lane_idx == 1 { return 1 } 144 if lane_idx == 2 { return 62 } if lane_idx == 3 { return 28 } 145 if lane_idx == 4 { return 27 } if lane_idx == 5 { return 36 } 146 if lane_idx == 6 { return 44 } if lane_idx == 7 { return 6 } 147 if lane_idx == 8 { return 55 } if lane_idx == 9 { return 20 } 148 if lane_idx == 10 { return 3 } if lane_idx == 11 { return 10 } 149 if lane_idx == 12 { return 43 } if lane_idx == 13 { return 25 } 150 if lane_idx == 14 { return 39 } if lane_idx == 15 { return 41 } 151 if lane_idx == 16 { return 45 } if lane_idx == 17 { return 15 } 152 if lane_idx == 18 { return 21 } if lane_idx == 19 { return 8 } 153 if lane_idx == 20 { return 18 } if lane_idx == 21 { return 2 } 154 if lane_idx == 22 { return 61 } if lane_idx == 23 { return 56 } 155 return 14 156} 157func _keccak_f1600(state_ptr: *u8) -> i64 { 158 let A: *u8 = state_ptr 159 let C: *u8 = (state_ptr as i64 + 200) as *u8 160 let B: *u8 = (state_ptr as i64 + 240) as *u8 161 var round: i64 = 0 162 while round < 24 { 163 var x: i64 = 0 164 while x < 5 { 165 let c0: i64 = _lane_load((A as i64 + 8 * (x + 0)) as *u8) 166 let c1: i64 = _lane_load((A as i64 + 8 * (x + 5)) as *u8) 167 let c2: i64 = _lane_load((A as i64 + 8 * (x + 10)) as *u8) 168 let c3: i64 = _lane_load((A as i64 + 8 * (x + 15)) as *u8) 169 let c4: i64 = _lane_load((A as i64 + 8 * (x + 20)) as *u8) 170 _lane_store((C as i64 + 8 * x) as *u8, c0 ^ c1 ^ c2 ^ c3 ^ c4) 171 x = x + 1 172 } 173 var x2: i64 = 0 174 while x2 < 5 { 175 let xm: i64 = (x2 + 4) % 5; let xp: i64 = (x2 + 1) % 5 176 let cl: i64 = _lane_load((C as i64 + 8 * xm) as *u8) 177 let cr: i64 = _lane_load((C as i64 + 8 * xp) as *u8) 178 let d: i64 = cl ^ _rotl64(cr, 1) 179 var y: i64 = 0 180 while y < 5 { 181 let off: i64 = 8 * (x2 + 5 * y) 182 let v: i64 = _lane_load((A as i64 + off) as *u8) ^ d 183 _lane_store((A as i64 + off) as *u8, v) 184 y = y + 1 185 } 186 x2 = x2 + 1 187 } 188 var y3: i64 = 0 189 while y3 < 5 { 190 var x3: i64 = 0 191 while x3 < 5 { 192 let src_idx: i64 = x3 + 5 * y3 193 let rot: i64 = _rho_off(src_idx) 194 let lane: i64 = _lane_load((A as i64 + 8 * src_idx) as *u8) 195 let rotated: i64 = _rotl64(lane, rot) 196 let new_x: i64 = y3 197 let new_y: i64 = (2 * x3 + 3 * y3) % 5 198 let dst_idx: i64 = new_x + 5 * new_y 199 _lane_store((B as i64 + 8 * dst_idx) as *u8, rotated) 200 x3 = x3 + 1 201 } 202 y3 = y3 + 1 203 } 204 var y4: i64 = 0 205 while y4 < 5 { 206 var x4: i64 = 0 207 while x4 < 5 { 208 let xp1: i64 = (x4 + 1) % 5; let xp2: i64 = (x4 + 2) % 5 209 let b0: i64 = _lane_load((B as i64 + 8 * (x4 + 5 * y4)) as *u8) 210 let b1: i64 = _lane_load((B as i64 + 8 * (xp1 + 5 * y4)) as *u8) 211 let b2: i64 = _lane_load((B as i64 + 8 * (xp2 + 5 * y4)) as *u8) 212 let nb1: i64 = (~b1) & 0xffffffffffffffff 213 let v: i64 = b0 ^ (nb1 & b2) 214 _lane_store((A as i64 + 8 * (x4 + 5 * y4)) as *u8, v) 215 x4 = x4 + 1 216 } 217 y4 = y4 + 1 218 } 219 let a00: i64 = _lane_load(A) ^ _keccak_rc(round) 220 _lane_store(A, a00) 221 round = round + 1 222 } 223 return 0 224} 225// Generic sponge: rate, dom_byte, out_len. state_ptr needs 440 bytes. 226func _sponge_one_shot(msg: *u8, msg_len: i64, rate: i64, dom_byte: i64, 227 state: *u8, out: *u8, out_len: i64) -> i64 { 228 var i: i64 = 0 229 while i < 200 { state[i] = 0; i = i + 1 } 230 var pos: i64 = 0 231 while pos + rate <= msg_len { 232 var b: i64 = 0 233 while b < rate { state[b] = (state[b] ^ msg[pos + b]) & 0xff; b = b + 1 } 234 _keccak_f1600(state) 235 pos = pos + rate 236 } 237 let tail: i64 = msg_len - pos 238 var t: i64 = 0 239 while t < tail { state[t] = (state[t] ^ msg[pos + t]) & 0xff; t = t + 1 } 240 state[tail] = (state[tail] ^ dom_byte) & 0xff 241 state[rate - 1] = (state[rate - 1] ^ 0x80) & 0xff 242 _keccak_f1600(state) 243 var written: i64 = 0 244 while written < out_len { 245 let remaining: i64 = out_len - written 246 var take: i64 = rate 247 if remaining < rate { take = remaining } 248 var k: i64 = 0 249 while k < take { out[written + k] = state[k]; k = k + 1 } 250 written = written + take 251 if written < out_len { _keccak_f1600(state) } 252 } 253 return 0 254} 255func _sha3_256(msg: *u8, msg_len: i64, state: *u8, out: *u8) -> i64 { 256 return _sponge_one_shot(msg, msg_len, 136, 0x06, state, out, 32) 257} 258func _shake256(msg: *u8, msg_len: i64, state: *u8, out: *u8, out_len: i64) -> i64 { 259 return _sponge_one_shot(msg, msg_len, 136, 0x1f, state, out, out_len) 260} 261func _shake128(msg: *u8, msg_len: i64, state: *u8, out: *u8, out_len: i64) -> i64 { 262 return _sponge_one_shot(msg, msg_len, 168, 0x1f, state, out, out_len) 263} 264 265// ============================================================================ 266// SECTION 3: ML-DSA zetas table (Montgomery form, 256 entries) 267// ============================================================================ 268// Generated from primitive 512th root of unity ζ=1753 mod q=8380417: 269// zetas[i] = (ζ^bitrev8(i) * R) mod q signed to (-q/2, q/2] 270// where R = 2^32 mod q = 4193792. 271func _zeta(i: i64) -> i64 { 272 if i == 0 { return -4186625 } if i == 1 { return 25847 } if i == 2 { return -2608894 } if i == 3 { return -518909 } 273 if i == 4 { return 237124 } if i == 5 { return -777960 } if i == 6 { return -876248 } if i == 7 { return 466468 } 274 if i == 8 { return 1826347 } if i == 9 { return 2353451 } if i == 10 { return -359251 } if i == 11 { return -2091905 } 275 if i == 12 { return 3119733 } if i == 13 { return -2884855 } if i == 14 { return 3111497 } if i == 15 { return 2680103 } 276 if i == 16 { return 2725464 } if i == 17 { return 1024112 } if i == 18 { return -1079900 } if i == 19 { return 3585928 } 277 if i == 20 { return -549488 } if i == 21 { return -1119584 } if i == 22 { return 2619752 } if i == 23 { return -2108549 } 278 if i == 24 { return -2118186 } if i == 25 { return -3859737 } if i == 26 { return -1399561 } if i == 27 { return -3277672 } 279 if i == 28 { return 1757237 } if i == 29 { return -19422 } if i == 30 { return 4010497 } if i == 31 { return 280005 } 280 if i == 32 { return 2706023 } if i == 33 { return 95776 } if i == 34 { return 3077325 } if i == 35 { return 3530437 } 281 if i == 36 { return -1661693 } if i == 37 { return -3592148 } if i == 38 { return -2537516 } if i == 39 { return 3915439 } 282 if i == 40 { return -3861115 } if i == 41 { return -3043716 } if i == 42 { return 3574422 } if i == 43 { return -2867647 } 283 if i == 44 { return 3539968 } if i == 45 { return -300467 } if i == 46 { return 2348700 } if i == 47 { return -539299 } 284 if i == 48 { return -1699267 } if i == 49 { return -1643818 } if i == 50 { return 3505694 } if i == 51 { return -3821735 } 285 if i == 52 { return 3507263 } if i == 53 { return -2140649 } if i == 54 { return -1600420 } if i == 55 { return 3699596 } 286 if i == 56 { return 811944 } if i == 57 { return 531354 } if i == 58 { return 954230 } if i == 59 { return 3881043 } 287 if i == 60 { return 3900724 } if i == 61 { return -2556880 } if i == 62 { return 2071892 } if i == 63 { return -2797779 } 288 if i == 64 { return -3930395 } if i == 65 { return -1528703 } if i == 66 { return -3677745 } if i == 67 { return -3041255 } 289 if i == 68 { return -1452451 } if i == 69 { return 3475950 } if i == 70 { return 2176455 } if i == 71 { return -1585221 } 290 if i == 72 { return -1257611 } if i == 73 { return 1939314 } if i == 74 { return -4083598 } if i == 75 { return -1000202 } 291 if i == 76 { return -3190144 } if i == 77 { return -3157330 } if i == 78 { return -3632928 } if i == 79 { return 126922 } 292 if i == 80 { return 3412210 } if i == 81 { return -983419 } if i == 82 { return 2147896 } if i == 83 { return 2715295 } 293 if i == 84 { return -2967645 } if i == 85 { return -3693493 } if i == 86 { return -411027 } if i == 87 { return -2477047 } 294 if i == 88 { return -671102 } if i == 89 { return -1228525 } if i == 90 { return -22981 } if i == 91 { return -1308169 } 295 if i == 92 { return -381987 } if i == 93 { return 1349076 } if i == 94 { return 1852771 } if i == 95 { return -1430430 } 296 if i == 96 { return -3343383 } if i == 97 { return 264944 } if i == 98 { return 508951 } if i == 99 { return 3097992 } 297 if i == 100 { return 44288 } if i == 101 { return -1100098 } if i == 102 { return 904516 } if i == 103 { return 3958618 } 298 if i == 104 { return -3724342 } if i == 105 { return -8578 } if i == 106 { return 1653064 } if i == 107 { return -3249728 } 299 if i == 108 { return 2389356 } if i == 109 { return -210977 } if i == 110 { return 759969 } if i == 111 { return -1316856 } 300 if i == 112 { return 189548 } if i == 113 { return -3553272 } if i == 114 { return 3159746 } if i == 115 { return -1851402 } 301 if i == 116 { return -2409325 } if i == 117 { return -177440 } if i == 118 { return 1315589 } if i == 119 { return 1341330 } 302 if i == 120 { return 1285669 } if i == 121 { return -1584928 } if i == 122 { return -812732 } if i == 123 { return -1439742 } 303 if i == 124 { return -3019102 } if i == 125 { return -3881060 } if i == 126 { return -3628969 } if i == 127 { return 3839961 } 304 if i == 128 { return 2091667 } if i == 129 { return 3407706 } if i == 130 { return 2316500 } if i == 131 { return 3817976 } 305 if i == 132 { return -3342478 } if i == 133 { return 2244091 } if i == 134 { return -2446433 } if i == 135 { return -3562462 } 306 if i == 136 { return 266997 } if i == 137 { return 2434439 } if i == 138 { return -1235728 } if i == 139 { return 3513181 } 307 if i == 140 { return -3520352 } if i == 141 { return -3759364 } if i == 142 { return -1197226 } if i == 143 { return -3193378 } 308 if i == 144 { return 900702 } if i == 145 { return 1859098 } if i == 146 { return 909542 } if i == 147 { return 819034 } 309 if i == 148 { return 495491 } if i == 149 { return -1613174 } if i == 150 { return -43260 } if i == 151 { return -522500 } 310 if i == 152 { return -655327 } if i == 153 { return -3122442 } if i == 154 { return 2031748 } if i == 155 { return 3207046 } 311 if i == 156 { return -3556995 } if i == 157 { return -525098 } if i == 158 { return -768622 } if i == 159 { return -3595838 } 312 if i == 160 { return 342297 } if i == 161 { return 286988 } if i == 162 { return -2437823 } if i == 163 { return 4108315 } 313 if i == 164 { return 3437287 } if i == 165 { return -3342277 } if i == 166 { return 1735879 } if i == 167 { return 203044 } 314 if i == 168 { return 2842341 } if i == 169 { return 2691481 } if i == 170 { return -2590150 } if i == 171 { return 1265009 } 315 if i == 172 { return 4055324 } if i == 173 { return 1247620 } if i == 174 { return 2486353 } if i == 175 { return 1595974 } 316 if i == 176 { return -3767016 } if i == 177 { return 1250494 } if i == 178 { return 2635921 } if i == 179 { return -3548272 } 317 if i == 180 { return -2994039 } if i == 181 { return 1869119 } if i == 182 { return 1903435 } if i == 183 { return -1050970 } 318 if i == 184 { return -1333058 } if i == 185 { return 1237275 } if i == 186 { return -3318210 } if i == 187 { return -1430225 } 319 if i == 188 { return -451100 } if i == 189 { return 1312455 } if i == 190 { return 3306115 } if i == 191 { return -1962642 } 320 if i == 192 { return -1279661 } if i == 193 { return 1917081 } if i == 194 { return -2546312 } if i == 195 { return -1374803 } 321 if i == 196 { return 1500165 } if i == 197 { return 777191 } if i == 198 { return 2235880 } if i == 199 { return 3406031 } 322 if i == 200 { return -542412 } if i == 201 { return -2831860 } if i == 202 { return -1671176 } if i == 203 { return -1846953 } 323 if i == 204 { return -2584293 } if i == 205 { return -3724270 } if i == 206 { return 594136 } if i == 207 { return -3776993 } 324 if i == 208 { return -2013608 } if i == 209 { return 2432395 } if i == 210 { return 2454455 } if i == 211 { return -164721 } 325 if i == 212 { return 1957272 } if i == 213 { return 3369112 } if i == 214 { return 185531 } if i == 215 { return -1207385 } 326 if i == 216 { return -3183426 } if i == 217 { return 162844 } if i == 218 { return 1616392 } if i == 219 { return 3014001 } 327 if i == 220 { return 810149 } if i == 221 { return 1652634 } if i == 222 { return -3694233 } if i == 223 { return -1799107 } 328 if i == 224 { return -3038916 } if i == 225 { return 3523897 } if i == 226 { return 3866901 } if i == 227 { return 269760 } 329 if i == 228 { return 2213111 } if i == 229 { return -975884 } if i == 230 { return 1717735 } if i == 231 { return 472078 } 330 if i == 232 { return -426683 } if i == 233 { return 1723600 } if i == 234 { return -1803090 } if i == 235 { return 1910376 } 331 if i == 236 { return -1667432 } if i == 237 { return -1104333 } if i == 238 { return -260646 } if i == 239 { return -3833893 } 332 if i == 240 { return -2939036 } if i == 241 { return -2235985 } if i == 242 { return -420899 } if i == 243 { return -2286327 } 333 if i == 244 { return 183443 } if i == 245 { return -976891 } if i == 246 { return 1612842 } if i == 247 { return -3545687 } 334 if i == 248 { return -554416 } if i == 249 { return 3919660 } if i == 250 { return -48306 } if i == 251 { return -1362209 } 335 if i == 252 { return 3937738 } if i == 253 { return 1400424 } if i == 254 { return -846154 } if i == 255 { return 1976782 } 336} 337 338// ============================================================================ 339// SECTION 4: NTT / INVNTT in Z_8380417 340// ============================================================================ 341// Per FIPS 204 / PQClean Dilithium: 342// ntt walks len = 128, 64, 32, ... 1; k starts at 0, ++k before each zeta read 343// so zetas[1..255] are used (zetas[0] is the unused sentinel). 344 345func nx_mldsa_ntt(poly: *u8) -> i64 { 346 var k: i64 = 0 347 var len: i64 = 128 348 while len > 0 { 349 var start: i64 = 0 350 while start < MLDSA_N { 351 k = k + 1 352 let zeta: i64 = _zeta(k) 353 var j: i64 = start 354 while j < start + len { 355 let t: i64 = _fqmul(zeta, _i32_load_le(poly, j + len)) 356 let aj: i64 = _i32_load_le(poly, j) 357 _i32_store_le(poly, j + len, aj - t) 358 _i32_store_le(poly, j, aj + t) 359 j = j + 1 360 } 361 start = j + len 362 } 363 len = len >> 1 364 } 365 return 0 366} 367 368// INVNTT: walks zetas backward, final scale by f = mont^2 * 256^{-1} mod q. 369// Per PQClean dilithium ntt.c invntt: 370// f = 41978 (= mont^2 / 256 mod q signed) 371func nx_mldsa_invntt(poly: *u8) -> i64 { 372 let f: i64 = 41978 373 var k: i64 = 256 374 var len: i64 = 1 375 while len < MLDSA_N { 376 var start: i64 = 0 377 while start < MLDSA_N { 378 k = k - 1 379 let zeta: i64 = -_zeta(k) 380 var j: i64 = start 381 while j < start + len { 382 let aj: i64 = _i32_load_le(poly, j) 383 let ajl: i64 = _i32_load_le(poly, j + len) 384 _i32_store_le(poly, j, aj + ajl) 385 _i32_store_le(poly, j + len, _fqmul(zeta, aj - ajl)) 386 j = j + 1 387 } 388 start = j + len 389 } 390 len = len << 1 391 } 392 // Final scale 393 var i: i64 = 0 394 while i < MLDSA_N { 395 let v: i64 = _i32_load_le(poly, i) 396 _i32_store_le(poly, i, _fqmul(v, f)) 397 i = i + 1 398 } 399 return 0 400} 401 402// Pointwise multiplication in NTT domain. 403func nx_mldsa_basemul(out: *u8, a: *u8, b: *u8) -> i64 { 404 var i: i64 = 0 405 while i < MLDSA_N { 406 let r: i64 = _fqmul(_i32_load_le(a, i), _i32_load_le(b, i)) 407 _i32_store_le(out, i, r) 408 i = i + 1 409 } 410 return 0 411} 412 413// ============================================================================ 414// SECTION 4.5: Sampling primitives (L132) 415// ============================================================================ 416 417// rej_uniform: parse SHAKE128 byte stream into uniform [0, q) i32 coefs. 418// Per FIPS 204 §4.2 / PQClean Dilithium rej_uniform: 419// read 3 bytes -> t = b0 | (b1<<8) | ((b2 & 0x7F)<<16) // 23-bit value 420// accept if t < q 421// continue until 256 coefs accepted 422// Returns bytes consumed on success, -1 if buf exhausted. 423func nx_mldsa_rej_uniform(poly: *u8, buf: *u8, buf_len: i64) -> i64 { 424 var i: i64 = 0 425 var pos: i64 = 0 426 while i < MLDSA_N { 427 if pos + 3 > buf_len { return -1 } 428 let b0: i64 = buf[pos] 429 let b1: i64 = buf[pos + 1] 430 let b2: i64 = buf[pos + 2] & 0x7f // mask top bit 431 pos = pos + 3 432 let t: i64 = b0 | (b1 << 8) | (b2 << 16) 433 if t < MLDSA_Q { 434 _i32_store_le(poly, i, t) 435 i = i + 1 436 } 437 } 438 return pos 439} 440 441// rej_bounded_eta4: parse SHAKE256 stream into [-4, 4] coefs. 442// Per FIPS 204 §4.3 / PQClean rej_eta when ETA == 4: 443// read 1 byte -> two nibbles, each accepted if < 9, output = 4 - nibble. 444// Returns bytes consumed on success, -1 if buf exhausted. 445func nx_mldsa_rej_eta4(poly: *u8, buf: *u8, buf_len: i64) -> i64 { 446 var i: i64 = 0 447 var pos: i64 = 0 448 while i < MLDSA_N { 449 if pos + 1 > buf_len { return -1 } 450 let byte: i64 = buf[pos] 451 pos = pos + 1 452 let t0: i64 = byte & 0x0f 453 let t1: i64 = (byte >> 4) & 0x0f 454 if t0 < 9 { 455 _i32_store_le(poly, i, 4 - t0) 456 i = i + 1 457 } 458 if i < MLDSA_N { 459 if t1 < 9 { 460 _i32_store_le(poly, i, 4 - t1) 461 i = i + 1 462 } 463 } 464 } 465 return pos 466} 467 468// sample_in_ball: build challenge polynomial c with exactly tau (=49) 469// non-zero coefficients each in {-1, +1}, rest zero. 470// Per FIPS 204 §4.4 / PQClean Dilithium poly_challenge: 471// 1. SHAKE256(c_tilde_32, 8) -> sign_bits (8 bytes = 64 bits) 472// 2. continue squeezing for the Fisher-Yates index byte 473// 3. for i from N-tau to N-1: 474// repeatedly read 1 byte j from SHAKE stream until j <= i 475// c[i] = c[j] 476// c[j] = (sign_bits & 1) ? -1 : 1; sign_bits >>= 1 477// 478// Caller provides: 479// poly: 1024-byte output (cleared to all-zero first internally) 480// c_tilde_32: 32-byte commitment hash 481// scratch: >= 600 bytes (Keccak state 440 + SHAKE squeeze buf ~136) 482func nx_mldsa_sample_in_ball(poly: *u8, c_tilde_32: *u8, scratch: *u8) -> i64 { 483 // Zero the output poly. 484 var z: i64 = 0 485 while z < MLDSA_N { 486 _i32_store_le(poly, z, 0) 487 z = z + 1 488 } 489 let keccak_state: *u8 = scratch 490 let shake_buf: *u8 = (scratch as i64 + 440) as *u8 491 492 // Stream MUST be incremental: PQClean does one SHAKE256 over c_tilde, 493 // then squeezes 8 bytes for signs, then squeezes individual bytes for j. 494 // We pre-squeeze a generous chunk (1024 bytes) which suffices for tau=49 495 // and a few rejections. If exhausted we re-squeeze with more bytes from 496 // the same seed (deterministic). 497 _shake256(c_tilde_32, 32, keccak_state, shake_buf, 1024) 498 499 // First 8 bytes = sign_bits packed little-endian (only low tau=49 bits used). 500 var sign_bits_lo: i64 = 0 501 var k: i64 = 0 502 while k < 8 { 503 sign_bits_lo = sign_bits_lo | ((shake_buf[k] as i64) << (k * 8)) 504 k = k + 1 505 } 506 var pos: i64 = 8 507 508 var i: i64 = MLDSA_N - MLDSA_TAU 509 while i < MLDSA_N { 510 // Read bytes until j <= i. 511 var j: i64 = i + 1 512 while j > i { 513 if pos >= 1024 { return -1 } // exhausted (very unlikely) 514 j = shake_buf[pos] 515 pos = pos + 1 516 } 517 // c[i] = c[j]; c[j] = sign-from-bits 518 _i32_store_le(poly, i, _i32_load_le(poly, j)) 519 var s: i64 = 1 520 if (sign_bits_lo & 1) != 0 { s = -1 } 521 _i32_store_le(poly, j, s) 522 sign_bits_lo = sign_bits_lo >> 1 523 i = i + 1 524 } 525 return 0 526} 527 528// ============================================================================ 529// SECTION 4.6: Decompose / hint primitives (L133, FIPS 204 §4.5) 530// ============================================================================ 531 532// power2round(a) -> (a1, a0) such that a = a1 * 2^D + a0, |a0| <= 2^(D-1). 533// ML-DSA-65 D = 13. PQClean returns a1, writes a0 to *a0. 534// We return both via caller-provided i64 pointers (out_a1, out_a0). 535// Operates on canonical [0, q) input. 536func _power2round_coef(a: i64, out_a1: *u8, out_a0: *u8) -> i64 { 537 let a1: i64 = (a + (1 << (MLDSA_D - 1)) - 1) >> MLDSA_D 538 let a0: i64 = a - (a1 << MLDSA_D) 539 _i32_store_le(out_a1, 0, a1) 540 _i32_store_le(out_a0, 0, a0) 541 return 0 542} 543 544// Apply power2round to every coef of `in_poly`, writing a1 vec + a0 vec. 545func nx_mldsa_poly_power2round(a1_poly: *u8, a0_poly: *u8, in_poly: *u8) -> i64 { 546 var i: i64 = 0 547 while i < MLDSA_N { 548 let a: i64 = _i32_load_le(in_poly, i) 549 let a1: i64 = (a + (1 << (MLDSA_D - 1)) - 1) >> MLDSA_D 550 let a0: i64 = a - (a1 << MLDSA_D) 551 _i32_store_le(a1_poly, i, a1) 552 _i32_store_le(a0_poly, i, a0) 553 i = i + 1 554 } 555 return 0 556} 557 558// decompose(a) -> (a1, a0) such that a = a1 * alpha + a0, 559// -alpha/2 < a0 <= alpha/2 (with special carry when a1 == 16). 560// ML-DSA-65 alpha = 2 * gamma2 = 523776. 561// Per PQClean Dilithium (LEVEL=3, gamma2=(q-1)/32): 562// a1 = (a + 127) >> 7 563// a1 = (a1 * 11275 + (1 << 23)) >> 24 564// a1 ^= ((43 - a1) >> 31) & a1 // if a1 > 43, a1 = 0 (= 16 wraps to 0) 565// a0 = a - a1 * 2 * gamma2 566// a0 -= (((q-1)/2 - a0) >> 31) & q 567// Caller poly_decompose iterates over 256 coefs. 568func _decompose_coef(a: i64, out_a0: *u8, off: i64) -> i64 { 569 let half_q_minus_1: i64 = (MLDSA_Q - 1) / 2 570 // ML-DSA-65 (GAMMA2 = (q-1)/32) decompose constants per PQClean 571 // Dilithium3 newer ref / FIPS 204: 1025, shift 22, mask 15. 572 // (NOT 11275/24/^=43 -- those are for GAMMA2 = (q-1)/88.) 573 var a1: i64 = (a + 127) >> 7 574 a1 = (a1 * 1025 + (1 << 21)) >> 22 575 a1 = a1 & 15 576 var a0: i64 = a - a1 * 2 * MLDSA_GAMMA2 577 // a0 -= (((q-1)/2 - a0) >> 31) & q 578 let neg_mask: i64 = (half_q_minus_1 - a0) >> 31 579 a0 = a0 - (neg_mask & MLDSA_Q) 580 _i32_store_le(out_a0, off, a0) 581 return a1 582} 583 584func nx_mldsa_poly_decompose(a1_poly: *u8, a0_poly: *u8, in_poly: *u8) -> i64 { 585 var i: i64 = 0 586 while i < MLDSA_N { 587 let a: i64 = _i32_load_le(in_poly, i) 588 let a1: i64 = _decompose_coef(a, a0_poly, i) 589 _i32_store_le(a1_poly, i, a1) 590 i = i + 1 591 } 592 return 0 593} 594 595// make_hint(a0, a1): returns 1 if high bits of a + z would differ from a1. 596// Per PQClean: hint = 1 iff (a0 > gamma2) OR (a0 < -gamma2) OR 597// (a0 == -gamma2 AND a1 != 0). 598func _make_hint_coef(a0: i64, a1: i64) -> i64 { 599 if a0 > MLDSA_GAMMA2 { return 1 } 600 if a0 < -MLDSA_GAMMA2 { return 1 } 601 if a0 == -MLDSA_GAMMA2 { 602 if a1 != 0 { return 1 } 603 } 604 return 0 605} 606 607// make_hint(z, r): given vectors z (a0_poly) and r (a1_poly), produces 608// a hint polynomial (1 bit per coef, stored as i32 0/1). Returns the 609// total weight (number of 1-bits) for the omega bound check. 610func nx_mldsa_poly_make_hint(hint_out: *u8, a0_poly: *u8, a1_poly: *u8) -> i64 { 611 var i: i64 = 0 612 var weight: i64 = 0 613 while i < MLDSA_N { 614 let h: i64 = _make_hint_coef(_i32_load_le(a0_poly, i), 615 _i32_load_le(a1_poly, i)) 616 _i32_store_le(hint_out, i, h) 617 weight = weight + h 618 i = i + 1 619 } 620 return weight 621} 622 623// use_hint(r, h): recover high bits of r + (small) given r and hint. 624// Per PQClean: decompose r -> (a1, a0). If hint==0: return a1. 625// If hint==1 and a0 > 0: return (a1 + 1) & 15 (i.e. mod 16, since a1 in [0, 15]) 626// If hint==1 and a0 <= 0: return (a1 - 1) & 15 627func _use_hint_coef(a: i64, hint: i64) -> i64 { 628 let half_q_minus_1: i64 = (MLDSA_Q - 1) / 2 629 // ML-DSA-65 GAMMA2=(q-1)/32 decompose constants (1025, 22, mask 15) 630 var a1: i64 = (a + 127) >> 7 631 a1 = (a1 * 1025 + (1 << 21)) >> 22 632 a1 = a1 & 15 633 var a0: i64 = a - a1 * 2 * MLDSA_GAMMA2 634 let neg_mask: i64 = (half_q_minus_1 - a0) >> 31 635 a0 = a0 - (neg_mask & MLDSA_Q) 636 if hint == 0 { return a1 } 637 if a0 > 0 { return (a1 + 1) & 15 } 638 return (a1 - 1) & 15 639} 640 641func nx_mldsa_poly_use_hint(out_poly: *u8, in_poly: *u8, hint_poly: *u8) -> i64 { 642 var i: i64 = 0 643 while i < MLDSA_N { 644 let r: i64 = _use_hint_coef(_i32_load_le(in_poly, i), 645 _i32_load_le(hint_poly, i)) 646 _i32_store_le(out_poly, i, r) 647 i = i + 1 648 } 649 return 0 650} 651 652// ============================================================================ 653// SECTION 4.7: Byte encoding for pk + sk (FIPS 204 §7) 654// ============================================================================ 655 656// Pack 256 10-bit coefficients (t1, range [0, 2^10)) into 320 bytes. 657// 4 coefs per 5 bytes (40 bits). 658func _pack_t1_poly(out: *u8, poly: *u8) -> i64 { 659 var i: i64 = 0 660 while i < 64 { // 64 groups of 4 coefs = 256 661 let off: i64 = i * 5 662 let c0: i64 = _i32_load_le(poly, i*4 + 0) & 0x3ff 663 let c1: i64 = _i32_load_le(poly, i*4 + 1) & 0x3ff 664 let c2: i64 = _i32_load_le(poly, i*4 + 2) & 0x3ff 665 let c3: i64 = _i32_load_le(poly, i*4 + 3) & 0x3ff 666 out[off + 0] = c0 & 0xff 667 out[off + 1] = ((c0 >> 8) | (c1 << 2)) & 0xff 668 out[off + 2] = ((c1 >> 6) | (c2 << 4)) & 0xff 669 out[off + 3] = ((c2 >> 4) | (c3 << 6)) & 0xff 670 out[off + 4] = (c3 >> 2) & 0xff 671 i = i + 1 672 } 673 return 0 674} 675 676// Pack 256 [-eta, eta] coefficients (eta=4) into 128 bytes. 677// 2 coefs per byte; each stored as (eta - coef) in 4 bits. 678func _pack_eta4_poly(out: *u8, poly: *u8) -> i64 { 679 var i: i64 = 0 680 while i < 128 { 681 let t0: i64 = MLDSA_ETA - _i32_load_le(poly, i*2) 682 let t1: i64 = MLDSA_ETA - _i32_load_le(poly, i*2 + 1) 683 out[i] = (t0 | (t1 << 4)) & 0xff 684 i = i + 1 685 } 686 return 0 687} 688 689// Pack 256 t0 coefficients (each in (-2^12, 2^12]) into 416 bytes. 690// 8 coefs per 13 bytes (104 bits). Each stored as (2^12 - coef) in 13 bits. 691func _pack_t0_poly(out: *u8, poly: *u8) -> i64 { 692 let bias: i64 = 1 << (MLDSA_D - 1) 693 var i: i64 = 0 694 while i < 32 { // 32 groups of 8 coefs = 256 695 let off: i64 = i * 13 696 let c0: i64 = bias - _i32_load_le(poly, i*8 + 0) 697 let c1: i64 = bias - _i32_load_le(poly, i*8 + 1) 698 let c2: i64 = bias - _i32_load_le(poly, i*8 + 2) 699 let c3: i64 = bias - _i32_load_le(poly, i*8 + 3) 700 let c4: i64 = bias - _i32_load_le(poly, i*8 + 4) 701 let c5: i64 = bias - _i32_load_le(poly, i*8 + 5) 702 let c6: i64 = bias - _i32_load_le(poly, i*8 + 6) 703 let c7: i64 = bias - _i32_load_le(poly, i*8 + 7) 704 out[off + 0] = c0 & 0xff 705 out[off + 1] = ((c0 >> 8) | (c1 << 5)) & 0xff 706 out[off + 2] = (c1 >> 3) & 0xff 707 out[off + 3] = ((c1 >> 11) | (c2 << 2)) & 0xff 708 out[off + 4] = ((c2 >> 6) | (c3 << 7)) & 0xff 709 out[off + 5] = (c3 >> 1) & 0xff 710 out[off + 6] = ((c3 >> 9) | (c4 << 4)) & 0xff 711 out[off + 7] = (c4 >> 4) & 0xff 712 out[off + 8] = ((c4 >> 12) | (c5 << 1)) & 0xff 713 out[off + 9] = ((c5 >> 7) | (c6 << 6)) & 0xff 714 out[off + 10] = (c6 >> 2) & 0xff 715 out[off + 11] = ((c6 >> 10) | (c7 << 3)) & 0xff 716 out[off + 12] = (c7 >> 5) & 0xff 717 i = i + 1 718 } 719 return 0 720} 721 722// ============================================================================ 723// SECTION 4.8: Higher-level keygen composition (FIPS 204 §5.1) 724// ============================================================================ 725 726// Build one matrix-A entry A[i][j] in NTT form via rej_uniform on SHAKE128. 727// seed_bytes buffer is shake_buf+1024 (34 bytes scratch for rho||j||i). 728func _expand_a_entry(out_poly: *u8, rho: *u8, j: i64, i_row: i64, 729 keccak_state: *u8, shake_buf: *u8) -> i64 { 730 let seed_buf: *u8 = (shake_buf as i64 + 1024) as *u8 731 var z: i64 = 0 732 while z < 32 { seed_buf[z] = rho[z]; z = z + 1 } 733 seed_buf[32] = j & 0xff 734 seed_buf[33] = i_row & 0xff 735 // Try increasing lengths until rejection sampler succeeds. 736 var attempt: i64 = 0 737 while attempt < 3 { 738 let try_len: i64 = 840 + attempt * 168 739 _shake128(seed_buf, 34, keccak_state, shake_buf, try_len) 740 let rc: i64 = nx_mldsa_rej_uniform(out_poly, shake_buf, try_len) 741 if rc >= 0 { return 0 } 742 attempt = attempt + 1 743 } 744 return -1 745} 746 747// Sample one s vector entry from rho_prime + nonce via SHAKE256 + rej_eta4. 748func _expand_s_entry(out_poly: *u8, rho_prime: *u8, nonce: i64, 749 keccak_state: *u8, shake_buf: *u8) -> i64 { 750 let seed_buf: *u8 = (shake_buf as i64 + 1024) as *u8 751 var z: i64 = 0 752 while z < 64 { seed_buf[z] = rho_prime[z]; z = z + 1 } 753 seed_buf[64] = nonce & 0xff 754 seed_buf[65] = (nonce >> 8) & 0xff 755 var attempt: i64 = 0 756 while attempt < 3 { 757 let try_len: i64 = 256 + attempt * 136 758 _shake256(seed_buf, 66, keccak_state, shake_buf, try_len) 759 let rc: i64 = nx_mldsa_rej_eta4(out_poly, shake_buf, try_len) 760 if rc >= 0 { return 0 } 761 attempt = attempt + 1 762 } 763 return -1 764} 765 766// ML-DSA-65 keygen. 767// in: seed_32 (caller-provided 32-byte entropy seed) 768// scratch (>= 64 KB working memory) 769// out: pk_out (1952 bytes) 770// sk_out (4032 bytes) 771// Returns 0 on success, -1 on (extremely unlikely) sampling exhaustion. 772// 773// Scratch layout (per F-meta-8 grader annotations): 774func nx_mldsa_keygen(seed_32: *u8, scratch: *u8, pk_out: *u8, sk_out: *u8) -> i64 { 775 let keccak_state: *u8 = scratch // scratch_range: keccak 0..439 776 let shake_buf: *u8 = (scratch as i64 + 512) as *u8 // scratch_range: shake_buf 512..2559 777 let g_out: *u8 = (scratch as i64 + 2560) as *u8 // scratch_range: g_out 2560..2687 (128 B for rho||rho'||K) 778 let rho: *u8 = (scratch as i64 + 2560) as *u8 // rho = g_out[0..32] 779 let rho_prime: *u8 = (scratch as i64 + 2592) as *u8 // rho_prime = g_out[32..96] 780 let k_secret: *u8 = (scratch as i64 + 2656) as *u8 // K = g_out[96..128] 781 let s1_vec: *u8 = (scratch as i64 + 2688) as *u8 // scratch_range: s1_vec 2688..7807 (5 * 1024) 782 let s2_vec: *u8 = (scratch as i64 + 7808) as *u8 // scratch_range: s2_vec 7808..13951 (6 * 1024) 783 let acc: *u8 = (scratch as i64 + 13952) as *u8 // scratch_range: acc 13952..14975 784 let t_vec: *u8 = (scratch as i64 + 14976) as *u8 // scratch_range: t_vec 14976..21119 (6 * 1024) 785 let t1_vec: *u8 = (scratch as i64 + 21120) as *u8 // scratch_range: t1_vec 21120..27263 786 let t0_vec: *u8 = (scratch as i64 + 27264) as *u8 // scratch_range: t0_vec 27264..33407 787 let A_entry: *u8 = (scratch as i64 + 33408) as *u8 // scratch_range: A_entry 33408..34431 (on-demand matrix expansion) 788 let g_input: *u8 = (scratch as i64 + 34432) as *u8 // scratch_range: g_input 34432..34465 (32 + 1 + 1 = 34 bytes for seed||k||l) 789 790 // 1. (rho, rho_prime, K) = SHAKE256(seed || k || l, 128 bytes total). 791 var i: i64 = 0 792 while i < 32 { g_input[i] = seed_32[i]; i = i + 1 } 793 g_input[32] = MLDSA_K & 0xff 794 g_input[33] = MLDSA_L & 0xff 795 _shake256(g_input, 34, keccak_state, g_out, 128) 796 797 // 2. ExpandS: nonce 0..l-1 for s1; nonce l..l+k-1 for s2. 798 var n: i64 = 0 799 var v: i64 = 0 800 while v < MLDSA_L { 801 let rc: i64 = _expand_s_entry((s1_vec as i64 + v * POLY_BUF) as *u8, 802 rho_prime, n, keccak_state, shake_buf) 803 if rc < 0 { return -1 } 804 n = n + 1 805 v = v + 1 806 } 807 v = 0 808 while v < MLDSA_K { 809 let rc: i64 = _expand_s_entry((s2_vec as i64 + v * POLY_BUF) as *u8, 810 rho_prime, n, keccak_state, shake_buf) 811 if rc < 0 { return -1 } 812 n = n + 1 813 v = v + 1 814 } 815 816 // 3. NTT each s1 entry in place. 817 v = 0 818 while v < MLDSA_L { 819 nx_mldsa_ntt((s1_vec as i64 + v * POLY_BUF) as *u8) 820 v = v + 1 821 } 822 823 // 4. For each row r in 0..k: 824 // acc = sum_j NTT(A[r][j]) ∘ NTT(s1[j]) (pointwise, accumulated) 825 // invntt(acc) -> canonical 826 // t[r] = acc + s2[r] 827 var r: i64 = 0 828 while r < MLDSA_K { 829 // Zero acc 830 var z: i64 = 0 831 while z < MLDSA_N { _i32_store_le(acc, z, 0); z = z + 1 } 832 // For each column j: expand A[r][j], basemul into acc. 833 var j: i64 = 0 834 while j < MLDSA_L { 835 let rc: i64 = _expand_a_entry(A_entry, rho, j, r, keccak_state, shake_buf) 836 if rc < 0 { return -1 } 837 // pointwise mul + accumulate 838 var k: i64 = 0 839 while k < MLDSA_N { 840 let p: i64 = _fqmul(_i32_load_le(A_entry, k), 841 _i32_load_le((s1_vec as i64 + j * POLY_BUF) as *u8, k)) 842 let acc_v: i64 = _i32_load_le(acc, k) + p 843 _i32_store_le(acc, k, acc_v) 844 k = k + 1 845 } 846 j = j + 1 847 } 848 // invntt(acc) -> canonical-form polynomial (with one mont factor; that's fine, 849 // it gets absorbed into t and then power2round which uses canon). 850 nx_mldsa_invntt(acc) 851 // t[r] = reduce32(acc + s2[r]) 852 let t_row: *u8 = (t_vec as i64 + r * POLY_BUF) as *u8 853 let s2_row: *u8 = (s2_vec as i64 + r * POLY_BUF) as *u8 854 var k: i64 = 0 855 while k < MLDSA_N { 856 let sum: i64 = _i32_load_le(acc, k) + _i32_load_le(s2_row, k) 857 // Bring to canonical [0, q) for power2round. 858 var canon: i64 = _reduce32(sum) 859 if canon < 0 { canon = canon + MLDSA_Q } 860 _i32_store_le(t_row, k, canon) 861 k = k + 1 862 } 863 r = r + 1 864 } 865 866 // 5. Power2Round each t[r] -> (t1[r], t0[r]). 867 var rr: i64 = 0 868 while rr < MLDSA_K { 869 nx_mldsa_poly_power2round((t1_vec as i64 + rr * POLY_BUF) as *u8, 870 (t0_vec as i64 + rr * POLY_BUF) as *u8, 871 (t_vec as i64 + rr * POLY_BUF) as *u8) 872 rr = rr + 1 873 } 874 875 // 6. pk = rho (32) || pack_t1(t1) (k * 320 = 1920) = 1952 bytes 876 var rho_i: i64 = 0 877 while rho_i < 32 { pk_out[rho_i] = rho[rho_i]; rho_i = rho_i + 1 } 878 var rk: i64 = 0 879 while rk < MLDSA_K { 880 _pack_t1_poly((pk_out as i64 + 32 + rk * 320) as *u8, 881 (t1_vec as i64 + rk * POLY_BUF) as *u8) 882 rk = rk + 1 883 } 884 885 // 7. tr = SHAKE256(pk, 64). Goes into sk at offset 32+32=64. 886 let tr: *u8 = (sk_out as i64 + 64) as *u8 887 _shake256(pk_out, MLDSA_PK_BYTES, keccak_state, tr, 64) 888 889 // 8. sk = rho (32) || K (32) || tr (64) || pack_eta(s1) (l*128=640) 890 // || pack_eta(s2) (k*128=768) || pack_t0(t0) (k*416=2496) = 4032 bytes 891 var sk_i: i64 = 0 892 while sk_i < 32 { sk_out[sk_i] = rho[sk_i]; sk_i = sk_i + 1 } 893 sk_i = 0 894 while sk_i < 32 { sk_out[32 + sk_i] = k_secret[sk_i]; sk_i = sk_i + 1 } 895 // (tr already in place at offset 64) 896 var s1_i: i64 = 0 897 while s1_i < MLDSA_L { 898 _pack_eta4_poly((sk_out as i64 + 128 + s1_i * 128) as *u8, 899 (s1_vec as i64 + s1_i * POLY_BUF) as *u8) 900 s1_i = s1_i + 1 901 } 902 var s2_i: i64 = 0 903 while s2_i < MLDSA_K { 904 _pack_eta4_poly((sk_out as i64 + 768 + s2_i * 128) as *u8, 905 (s2_vec as i64 + s2_i * POLY_BUF) as *u8) 906 s2_i = s2_i + 1 907 } 908 var t0_i: i64 = 0 909 while t0_i < MLDSA_K { 910 _pack_t0_poly((sk_out as i64 + 1536 + t0_i * 416) as *u8, 911 (t0_vec as i64 + t0_i * POLY_BUF) as *u8) 912 t0_i = t0_i + 1 913 } 914 915 return 0 916} 917 918// Self-test: keygen twice with same seed; pk + sk must be byte-identical. 919// Returns 0 on success, non-zero bitmap of failure axes. 920// bit 0: first keygen returned non-zero 921// bit 1: second keygen returned non-zero 922// bit 2: pk mismatch 923// bit 3: sk mismatch 924// Caller provides scratch (>= 64 KB). Internal pk_1, sk_1, pk_2, sk_2 are 925// placed at high addresses inside scratch. 926func nx_mldsa_keygen_deterministic_test(seed_32: *u8, scratch: *u8) -> i64 { 927 let pk_1: *u8 = (scratch as i64 + 50000) as *u8 928 let sk_1: *u8 = (scratch as i64 + 52000) as *u8 929 let pk_2: *u8 = (scratch as i64 + 56500) as *u8 930 let sk_2: *u8 = (scratch as i64 + 58500) as *u8 931 var fail: i64 = 0 932 let rc1: i64 = nx_mldsa_keygen(seed_32, scratch, pk_1, sk_1) 933 if rc1 != 0 { fail = fail | 1 } 934 let rc2: i64 = nx_mldsa_keygen(seed_32, scratch, pk_2, sk_2) 935 if rc2 != 0 { fail = fail | 2 } 936 var pk_diff: i64 = 0 937 var i: i64 = 0 938 while i < MLDSA_PK_BYTES { 939 if pk_1[i] != pk_2[i] { pk_diff = 1 } 940 i = i + 1 941 } 942 if pk_diff != 0 { fail = fail | 4 } 943 var sk_diff: i64 = 0 944 var j: i64 = 0 945 while j < MLDSA_SK_BYTES { 946 if sk_1[j] != sk_2[j] { sk_diff = 1 } 947 j = j + 1 948 } 949 if sk_diff != 0 { fail = fail | 8 } 950 return fail 951} 952 953// ============================================================================ 954// SECTION 4.9: ExpandMask + sig-byte encoders (L135, FIPS 204 §5.3 / §7) 955// ============================================================================ 956 957// Unpack one 20-bit chunk into a coefficient in (-gamma1, gamma1]. 958// Per FIPS 204 §7.1 / PQClean: for gamma1 = 2^19, output = gamma1 - x. 959// 5 bytes -> 2 coefs (PQClean unpacks z 4 coefs at a time across 10 bytes 960// = 80 bits = 4 * 20 bits). Here we do 4-coef chunks. 961func _unpack_z_chunk(out_poly: *u8, base_idx: i64, buf: *u8, buf_off: i64) -> i64 { 962 // 4 coefs from 10 bytes (40 nibbles = 80 bits = 4 * 20) 963 let b0: i64 = buf[buf_off + 0] 964 let b1: i64 = buf[buf_off + 1] 965 let b2: i64 = buf[buf_off + 2] 966 let b3: i64 = buf[buf_off + 3] 967 let b4: i64 = buf[buf_off + 4] 968 let b5: i64 = buf[buf_off + 5] 969 let b6: i64 = buf[buf_off + 6] 970 let b7: i64 = buf[buf_off + 7] 971 let b8: i64 = buf[buf_off + 8] 972 let b9: i64 = buf[buf_off + 9] 973 let v0: i64 = (b0 | (b1 << 8) | ((b2 & 0x0f) << 16)) & 0xfffff 974 let v1: i64 = ((b2 >> 4) | (b3 << 4) | (b4 << 12)) & 0xfffff 975 let v2: i64 = (b5 | (b6 << 8) | ((b7 & 0x0f) << 16)) & 0xfffff 976 let v3: i64 = ((b7 >> 4) | (b8 << 4) | (b9 << 12)) & 0xfffff 977 _i32_store_le(out_poly, base_idx + 0, MLDSA_GAMMA1 - v0) 978 _i32_store_le(out_poly, base_idx + 1, MLDSA_GAMMA1 - v1) 979 _i32_store_le(out_poly, base_idx + 2, MLDSA_GAMMA1 - v2) 980 _i32_store_le(out_poly, base_idx + 3, MLDSA_GAMMA1 - v3) 981 return 0 982} 983 984// Pack 256 coefs in (-gamma1, gamma1] into 640 bytes (20 bits/coef). 985// Per FIPS 204 / PQClean: store (gamma1 - coef) in 20 bits. 986func _pack_z_poly(out: *u8, poly: *u8) -> i64 { 987 var i: i64 = 0 988 while i < 64 { // 64 groups of 4 coefs = 256 989 let off: i64 = i * 10 990 let c0: i64 = MLDSA_GAMMA1 - _i32_load_le(poly, i*4 + 0) 991 let c1: i64 = MLDSA_GAMMA1 - _i32_load_le(poly, i*4 + 1) 992 let c2: i64 = MLDSA_GAMMA1 - _i32_load_le(poly, i*4 + 2) 993 let c3: i64 = MLDSA_GAMMA1 - _i32_load_le(poly, i*4 + 3) 994 out[off + 0] = c0 & 0xff 995 out[off + 1] = (c0 >> 8) & 0xff 996 out[off + 2] = ((c0 >> 16) | (c1 << 4)) & 0xff 997 out[off + 3] = (c1 >> 4) & 0xff 998 out[off + 4] = (c1 >> 12) & 0xff 999 out[off + 5] = c2 & 0xff 1000 out[off + 6] = (c2 >> 8) & 0xff 1001 out[off + 7] = ((c2 >> 16) | (c3 << 4)) & 0xff 1002 out[off + 8] = (c3 >> 4) & 0xff 1003 out[off + 9] = (c3 >> 12) & 0xff 1004 i = i + 1 1005 } 1006 return 0 1007} 1008 1009func _unpack_z_poly(poly: *u8, bytes_in: *u8) -> i64 { 1010 var i: i64 = 0 1011 while i < 64 { 1012 _unpack_z_chunk(poly, i * 4, bytes_in, i * 10) 1013 i = i + 1 1014 } 1015 return 0 1016} 1017 1018// Pack 256 coefs in [0, 16) into 128 bytes (4 bits/coef, 2 per byte). 1019func _pack_w1_poly(out: *u8, poly: *u8) -> i64 { 1020 var i: i64 = 0 1021 while i < 128 { 1022 let c0: i64 = _i32_load_le(poly, i * 2) & 0x0f 1023 let c1: i64 = _i32_load_le(poly, i * 2 + 1) & 0x0f 1024 out[i] = (c0 | (c1 << 4)) & 0xff 1025 i = i + 1 1026 } 1027 return 0 1028} 1029 1030// Pack hint vector: k polynomials of 0/1 -> (omega + k) bytes total. 1031// Format per FIPS 204: for each of k polys, append the indices of its 1032// non-zero positions, then write the cumulative count as a single byte. 1033// Final encoding: omega data bytes + k count bytes = 55 + 6 = 61 bytes. 1034// Caller passes hint_polys (k * POLY_BUF) and writes to out (61 bytes). 1035// Returns -1 if total nonzero weight > omega (invalid signature). 1036func _pack_hint_vec(out: *u8, hint_polys: *u8) -> i64 { 1037 // Zero entire out first. 1038 var z: i64 = 0 1039 while z < MLDSA_OMEGA + MLDSA_K { out[z] = 0; z = z + 1 } 1040 var idx: i64 = 0 1041 var p: i64 = 0 1042 while p < MLDSA_K { 1043 var j: i64 = 0 1044 while j < MLDSA_N { 1045 let h: i64 = _i32_load_le((hint_polys as i64 + p * POLY_BUF) as *u8, j) 1046 if h != 0 { 1047 if idx >= MLDSA_OMEGA { return -1 } 1048 out[idx] = j & 0xff 1049 idx = idx + 1 1050 } 1051 j = j + 1 1052 } 1053 out[MLDSA_OMEGA + p] = idx & 0xff 1054 p = p + 1 1055 } 1056 return idx 1057} 1058 1059// Unpack hint vector: read (omega + k) bytes, recover k polys with 0/1 coefs. 1060// Returns 0 on success, -1 if signature malformed (indices out of order or 1061// cumulative count overflows omega). 1062func _unpack_hint_vec(hint_polys: *u8, bytes_in: *u8) -> i64 { 1063 // Zero output polys. 1064 var p0: i64 = 0 1065 while p0 < MLDSA_K { 1066 var j0: i64 = 0 1067 while j0 < MLDSA_N { 1068 _i32_store_le((hint_polys as i64 + p0 * POLY_BUF) as *u8, j0, 0) 1069 j0 = j0 + 1 1070 } 1071 p0 = p0 + 1 1072 } 1073 var prev_idx: i64 = 0 1074 var p: i64 = 0 1075 while p < MLDSA_K { 1076 let cum: i64 = bytes_in[MLDSA_OMEGA + p] 1077 if cum < prev_idx { return -1 } 1078 if cum > MLDSA_OMEGA { return -1 } 1079 var last_pos: i64 = -1 1080 var i: i64 = prev_idx 1081 while i < cum { 1082 let pos: i64 = bytes_in[i] 1083 if pos <= last_pos { if last_pos >= 0 { return -1 } } 1084 last_pos = pos 1085 _i32_store_le((hint_polys as i64 + p * POLY_BUF) as *u8, pos, 1) 1086 i = i + 1 1087 } 1088 prev_idx = cum 1089 p = p + 1 1090 } 1091 return 0 1092} 1093 1094// ExpandMask: derive y[i] vector entry from rho_prime (64) || kappa+i (2 bytes LE). 1095// y has 256 coefs in (-gamma1, gamma1]. SHAKE256 squeezes 640 bytes 1096// (gamma1=2^19 needs 20 bits/coef * 256 = 5120 bits = 640 bytes). 1097// scratch: shake_buf at >= 640 bytes; seed_buf = shake_buf + 640 (66 bytes). 1098func _expand_mask_entry(out_poly: *u8, rho_prime: *u8, kappa_plus_i: i64, 1099 keccak_state: *u8, shake_buf: *u8) -> i64 { 1100 let seed_buf: *u8 = (shake_buf as i64 + 1024) as *u8 1101 var z: i64 = 0 1102 while z < 64 { seed_buf[z] = rho_prime[z]; z = z + 1 } 1103 seed_buf[64] = kappa_plus_i & 0xff 1104 seed_buf[65] = (kappa_plus_i >> 8) & 0xff 1105 _shake256(seed_buf, 66, keccak_state, shake_buf, 640) 1106 _unpack_z_poly(out_poly, shake_buf) 1107 return 0 1108} 1109 1110// ============================================================================ 1111// SECTION 4.10: L135 KAT export -- encoder round-trips + ExpandMask 1112// ============================================================================ 1113// 1114// Caller passes scratch buffer (>= 32 KB). Returns 0 on success, non-zero 1115// bitmap of failure axes: 1116// bit 0: pack_z/unpack_z round-trip mismatch 1117// bit 1: pack_w1 produces out-of-range output 1118// bit 2: pack_hint/unpack_hint round-trip mismatch 1119// bit 3: ExpandMask produces out-of-range output 1120// bit 4: ExpandMask not deterministic for same (rho_prime, kappa) 1121func nx_mldsa_l135_encoders_test(scratch: *u8) -> i64 { 1122 var fail: i64 = 0 1123 1124 // ---- z pack/unpack round-trip ---- 1125 let z_in: *u8 = scratch 1126 let z_bytes: *u8 = (scratch as i64 + POLY_BUF) as *u8 1127 let z_out: *u8 = (scratch as i64 + POLY_BUF + 640) as *u8 1128 // Fill z_in with values in (-gamma1, gamma1]. Use a deterministic pattern. 1129 var i: i64 = 0 1130 while i < MLDSA_N { 1131 let v: i64 = ((i * 7919) % (2 * MLDSA_GAMMA1)) - MLDSA_GAMMA1 + 1 1132 _i32_store_le(z_in, i, v) 1133 i = i + 1 1134 } 1135 _pack_z_poly(z_bytes, z_in) 1136 _unpack_z_poly(z_out, z_bytes) 1137 var diff_z: i64 = 0 1138 var j: i64 = 0 1139 while j < MLDSA_N { 1140 if _i32_load_le(z_in, j) != _i32_load_le(z_out, j) { diff_z = 1 } 1141 j = j + 1 1142 } 1143 if diff_z != 0 { fail = fail | 1 } 1144 1145 // ---- w1 pack: every output coef must be < 16 (4-bit) ---- 1146 let w1_in: *u8 = (scratch as i64 + 4096) as *u8 1147 let w1_bytes: *u8 = (scratch as i64 + 5120) as *u8 1148 var k: i64 = 0 1149 while k < MLDSA_N { 1150 _i32_store_le(w1_in, k, k % 16) 1151 k = k + 1 1152 } 1153 _pack_w1_poly(w1_bytes, w1_in) 1154 var w1_ok: i64 = 1 1155 var bi: i64 = 0 1156 while bi < 128 { 1157 // each byte should be (c0 | (c1 << 4)) where c0 = 2*bi%16, c1 = (2*bi+1)%16 1158 let expected: i64 = ((2 * bi) % 16) | ((((2 * bi + 1) % 16)) << 4) 1159 if w1_bytes[bi] != (expected & 0xff) { w1_ok = 0 } 1160 bi = bi + 1 1161 } 1162 if w1_ok == 0 { fail = fail | 2 } 1163 1164 // ---- hint pack/unpack round-trip ---- 1165 let hint_in_polys: *u8 = (scratch as i64 + 8192) as *u8 // k * 1024 = 6144 B 1166 let hint_bytes: *u8 = (scratch as i64 + 16384) as *u8 // 61 B 1167 let hint_out_polys: *u8 = (scratch as i64 + 16448) as *u8 // 6144 B 1168 // Fill hint_in with deterministic 0/1 pattern, total weight <= omega. 1169 var p: i64 = 0 1170 var total: i64 = 0 1171 while p < MLDSA_K { 1172 var jp: i64 = 0 1173 while jp < MLDSA_N { 1174 // Pick a few positions per poly; total stays under omega=55. 1175 var v: i64 = 0 1176 if total < MLDSA_OMEGA { 1177 // Sparse pattern: every (p+1)*37'th coef in this poly. 1178 if (jp + p * 7) % ((p + 1) * 37) == 0 { v = 1; total = total + 1 } 1179 } 1180 _i32_store_le((hint_in_polys as i64 + p * POLY_BUF) as *u8, jp, v) 1181 jp = jp + 1 1182 } 1183 p = p + 1 1184 } 1185 let pack_rc: i64 = _pack_hint_vec(hint_bytes, hint_in_polys) 1186 if pack_rc < 0 { fail = fail | 4 } 1187 let unpack_rc: i64 = _unpack_hint_vec(hint_out_polys, hint_bytes) 1188 if unpack_rc < 0 { fail = fail | 4 } 1189 var diff_h: i64 = 0 1190 var pp: i64 = 0 1191 while pp < MLDSA_K { 1192 var jj: i64 = 0 1193 while jj < MLDSA_N { 1194 let a: i64 = _i32_load_le((hint_in_polys as i64 + pp * POLY_BUF) as *u8, jj) 1195 let b: i64 = _i32_load_le((hint_out_polys as i64 + pp * POLY_BUF) as *u8, jj) 1196 if a != b { diff_h = 1 } 1197 jj = jj + 1 1198 } 1199 pp = pp + 1 1200 } 1201 if diff_h != 0 { fail = fail | 4 } 1202 1203 // ---- ExpandMask range check ---- 1204 let rho_prime: *u8 = (scratch as i64 + 24064) as *u8 // 64 B 1205 let mask_state: *u8 = (scratch as i64 + 24128) as *u8 // Keccak 440 B 1206 let mask_buf: *u8 = (scratch as i64 + 24640) as *u8 // shake 1024 + seed 66 1207 let mask_poly_1: *u8 = (scratch as i64 + 26240) as *u8 // 1024 B 1208 let mask_poly_2: *u8 = (scratch as i64 + 27264) as *u8 // 1024 B 1209 var rsi: i64 = 0 1210 while rsi < 64 { rho_prime[rsi] = (rsi * 11 + 5) & 0xff; rsi = rsi + 1 } 1211 _expand_mask_entry(mask_poly_1, rho_prime, 0, mask_state, mask_buf) 1212 // range: every coef in (-gamma1, gamma1] 1213 var range_ok: i64 = 1 1214 var mi: i64 = 0 1215 while mi < MLDSA_N { 1216 let c: i64 = _i32_load_le(mask_poly_1, mi) 1217 if c <= -MLDSA_GAMMA1 { range_ok = 0 } 1218 if c > MLDSA_GAMMA1 { range_ok = 0 } 1219 mi = mi + 1 1220 } 1221 if range_ok == 0 { fail = fail | 8 } 1222 // determinism: same (rho_prime, kappa) -> same poly 1223 _expand_mask_entry(mask_poly_2, rho_prime, 0, mask_state, mask_buf) 1224 var det_diff: i64 = 0 1225 var di: i64 = 0 1226 while di < MLDSA_N { 1227 if _i32_load_le(mask_poly_1, di) != _i32_load_le(mask_poly_2, di) { det_diff = 1 } 1228 di = di + 1 1229 } 1230 if det_diff != 0 { fail = fail | 16 } 1231 return fail 1232} 1233 1234// ============================================================================ 1235// SECTION 4.11: Decoders for sk (L136 prereqs) 1236// ============================================================================ 1237 1238// Unpack 128 bytes -> 256 [-eta, eta] coefs. Inverse of _pack_eta4_poly. 1239func _unpack_eta4_poly(out: *u8, bytes_in: *u8) -> i64 { 1240 var i: i64 = 0 1241 while i < 128 { 1242 let b: i64 = bytes_in[i] 1243 _i32_store_le(out, i*2, MLDSA_ETA - (b & 0x0f)) 1244 _i32_store_le(out, i*2 + 1, MLDSA_ETA - ((b >> 4) & 0x0f)) 1245 i = i + 1 1246 } 1247 return 0 1248} 1249 1250// Unpack 416 bytes -> 256 t0 coefs in (-2^12, 2^12]. Inverse of _pack_t0_poly. 1251func _unpack_t0_poly(out: *u8, bytes_in: *u8) -> i64 { 1252 let bias: i64 = 1 << (MLDSA_D - 1) 1253 var i: i64 = 0 1254 while i < 32 { 1255 let off: i64 = i * 13 1256 let b0: i64 = bytes_in[off + 0]; let b1: i64 = bytes_in[off + 1] 1257 let b2: i64 = bytes_in[off + 2]; let b3: i64 = bytes_in[off + 3] 1258 let b4: i64 = bytes_in[off + 4]; let b5: i64 = bytes_in[off + 5] 1259 let b6: i64 = bytes_in[off + 6]; let b7: i64 = bytes_in[off + 7] 1260 let b8: i64 = bytes_in[off + 8]; let b9: i64 = bytes_in[off + 9] 1261 let b10: i64 = bytes_in[off + 10] 1262 let b11: i64 = bytes_in[off + 11] 1263 let b12: i64 = bytes_in[off + 12] 1264 let v0: i64 = (b0 | (b1 << 8)) & 0x1fff 1265 let v1: i64 = ((b1 >> 5) | (b2 << 3) | (b3 << 11)) & 0x1fff 1266 let v2: i64 = ((b3 >> 2) | (b4 << 6)) & 0x1fff 1267 let v3: i64 = ((b4 >> 7) | (b5 << 1) | (b6 << 9)) & 0x1fff 1268 let v4: i64 = ((b6 >> 4) | (b7 << 4) | (b8 << 12)) & 0x1fff 1269 let v5: i64 = ((b8 >> 1) | (b9 << 7)) & 0x1fff 1270 let v6: i64 = ((b9 >> 6) | (b10 << 2) | (b11 << 10)) & 0x1fff 1271 let v7: i64 = ((b11 >> 3) | (b12 << 5)) & 0x1fff 1272 _i32_store_le(out, i*8 + 0, bias - v0) 1273 _i32_store_le(out, i*8 + 1, bias - v1) 1274 _i32_store_le(out, i*8 + 2, bias - v2) 1275 _i32_store_le(out, i*8 + 3, bias - v3) 1276 _i32_store_le(out, i*8 + 4, bias - v4) 1277 _i32_store_le(out, i*8 + 5, bias - v5) 1278 _i32_store_le(out, i*8 + 6, bias - v6) 1279 _i32_store_le(out, i*8 + 7, bias - v7) 1280 i = i + 1 1281 } 1282 return 0 1283} 1284 1285// Inf-norm of a polynomial: returns max(|coef|) using signed representation. 1286// For checking bound conditions in sign loop (||z||_inf < gamma1 - beta etc.) 1287func _poly_inf_norm(poly: *u8) -> i64 { 1288 var max_n: i64 = 0 1289 var i: i64 = 0 1290 while i < MLDSA_N { 1291 var v: i64 = _i32_load_le(poly, i) 1292 if v < 0 { v = -v } 1293 if v > max_n { max_n = v } 1294 i = i + 1 1295 } 1296 return max_n 1297} 1298 1299// Compute the inf-norm of a vector of k or l polys. 1300func _vec_inf_norm(vec: *u8, n_polys: i64) -> i64 { 1301 var max_n: i64 = 0 1302 var i: i64 = 0 1303 while i < n_polys { 1304 let n: i64 = _poly_inf_norm((vec as i64 + i * POLY_BUF) as *u8) 1305 if n > max_n { max_n = n } 1306 i = i + 1 1307 } 1308 return max_n 1309} 1310 1311// ============================================================================ 1312// SECTION 4.12: Sign function (FIPS 204 §5.2) 1313// ============================================================================ 1314// 1315// Scratch buffer requirement: caller passes >= 128 KB working memory. 1316// Internal layout below uses scratch_range annotations per F-meta-8. 1317 1318func nx_mldsa_sign(sk: *u8, msg: *u8, msg_len: i64, 1319 scratch: *u8, sig_out: *u8) -> i64 { 1320 let keccak_state: *u8 = scratch // scratch_range: keccak 0..439 1321 let shake_buf: *u8 = (scratch as i64 + 512) as *u8 // scratch_range: shake_buf 512..2559 1322 // sk components (decoded once into scratch slots) 1323 let s1_vec: *u8 = (scratch as i64 + 3072) as *u8 // scratch_range: s1_vec 3072..8191 (5 * 1024) 1324 let s2_vec: *u8 = (scratch as i64 + 8192) as *u8 // scratch_range: s2_vec 8192..14335 (6 * 1024) 1325 let t0_vec: *u8 = (scratch as i64 + 14336) as *u8 // scratch_range: t0_vec 14336..20479 (6 * 1024) 1326 // working 1327 let mu: *u8 = (scratch as i64 + 20480) as *u8 // scratch_range: mu 20480..20543 (64) 1328 let rho_prime: *u8 = (scratch as i64 + 20544) as *u8 // scratch_range: rho_prime 20544..20607 (64) 1329 let mu_msg_buf: *u8 = (scratch as i64 + 20608) as *u8 // scratch_range: mu_msg 20608..28799 (8192 max msg) 1330 let k_mu_buf: *u8 = (scratch as i64 + 28800) as *u8 // scratch_range: k_mu 28800..28895 (32+64=96) 1331 let y_vec: *u8 = (scratch as i64 + 28896) as *u8 // scratch_range: y_vec 28896..34015 (5 * 1024) 1332 let w_vec: *u8 = (scratch as i64 + 34016) as *u8 // scratch_range: w_vec 34016..40159 (6 * 1024) 1333 let w1_vec: *u8 = (scratch as i64 + 40160) as *u8 // scratch_range: w1_vec 40160..46303 1334 let w0_vec: *u8 = (scratch as i64 + 46304) as *u8 // scratch_range: w0_vec 46304..52447 1335 let c_poly: *u8 = (scratch as i64 + 52448) as *u8 // scratch_range: c_poly 52448..53471 1336 let c_hat: *u8 = (scratch as i64 + 53472) as *u8 // scratch_range: c_hat 53472..54495 1337 let cs1: *u8 = (scratch as i64 + 54496) as *u8 // scratch_range: cs1 54496..59615 (5 * 1024) 1338 let cs2: *u8 = (scratch as i64 + 59616) as *u8 // scratch_range: cs2 59616..65759 (6 * 1024) 1339 let ct0: *u8 = (scratch as i64 + 65760) as *u8 // scratch_range: ct0 65760..71903 (6 * 1024) 1340 let z_vec: *u8 = (scratch as i64 + 71904) as *u8 // scratch_range: z_vec 71904..77023 (5 * 1024) 1341 let r0_vec: *u8 = (scratch as i64 + 77024) as *u8 // scratch_range: r0_vec 77024..83167 1342 let h_vec: *u8 = (scratch as i64 + 83168) as *u8 // scratch_range: h_vec 83168..89311 1343 let A_entry: *u8 = (scratch as i64 + 89312) as *u8 // scratch_range: A_entry 89312..90335 1344 let acc_poly: *u8 = (scratch as i64 + 90336) as *u8 // scratch_range: acc 90336..91359 1345 let w1_packed: *u8 = (scratch as i64 + 91360) as *u8 // scratch_range: w1_packed 91360..92127 (k*128=768) 1346 let c_tilde: *u8 = (scratch as i64 + 92128) as *u8 // scratch_range: c_tilde 92128..92175 (48 B) 1347 1348 // sk layout: rho(32) || K(32) || tr(64) || s1_packed(640) || s2_packed(768) || t0_packed(2496) 1349 let rho: *u8 = sk 1350 let k_secret: *u8 = (sk as i64 + 32) as *u8 1351 let tr: *u8 = (sk as i64 + 64) as *u8 1352 let s1_packed: *u8 = (sk as i64 + 128) as *u8 1353 let s2_packed: *u8 = (sk as i64 + 768) as *u8 1354 let t0_packed: *u8 = (sk as i64 + 1536) as *u8 1355 1356 // Decode s1, s2, t0; NTT each in place. 1357 var iv: i64 = 0 1358 while iv < MLDSA_L { 1359 _unpack_eta4_poly((s1_vec as i64 + iv * POLY_BUF) as *u8, 1360 (s1_packed as i64 + iv * 128) as *u8) 1361 nx_mldsa_ntt((s1_vec as i64 + iv * POLY_BUF) as *u8) 1362 iv = iv + 1 1363 } 1364 iv = 0 1365 while iv < MLDSA_K { 1366 _unpack_eta4_poly((s2_vec as i64 + iv * POLY_BUF) as *u8, 1367 (s2_packed as i64 + iv * 128) as *u8) 1368 nx_mldsa_ntt((s2_vec as i64 + iv * POLY_BUF) as *u8) 1369 iv = iv + 1 1370 } 1371 iv = 0 1372 while iv < MLDSA_K { 1373 _unpack_t0_poly((t0_vec as i64 + iv * POLY_BUF) as *u8, 1374 (t0_packed as i64 + iv * 416) as *u8) 1375 nx_mldsa_ntt((t0_vec as i64 + iv * POLY_BUF) as *u8) 1376 iv = iv + 1 1377 } 1378 1379 // mu = SHAKE256(tr || msg, 64). Limited to msg_len <= 8128 by scratch. 1380 if msg_len > 8128 { return -2 } 1381 var ti: i64 = 0 1382 while ti < 64 { mu_msg_buf[ti] = tr[ti]; ti = ti + 1 } 1383 var mi: i64 = 0 1384 while mi < msg_len { mu_msg_buf[64 + mi] = msg[mi]; mi = mi + 1 } 1385 _shake256(mu_msg_buf, 64 + msg_len, keccak_state, mu, 64) 1386 1387 // rho_prime = SHAKE256(K || mu, 64). Deterministic variant (no rnd). 1388 var ki: i64 = 0 1389 while ki < 32 { k_mu_buf[ki] = k_secret[ki]; ki = ki + 1 } 1390 var mu_i: i64 = 0 1391 while mu_i < 64 { k_mu_buf[32 + mu_i] = mu[mu_i]; mu_i = mu_i + 1 } 1392 _shake256(k_mu_buf, 96, keccak_state, rho_prime, 64) 1393 1394 // Rejection-sample loop. Bounded by MAX_ATTEMPTS = 200 (real impls 1395 // expect ~3-7 iterations; 200 is safety). 1396 var kappa: i64 = 0 1397 var attempt: i64 = 0 1398 while attempt < 200 { 1399 // 1. y = ExpandMask(rho_prime, kappa+i) for i in 0..l 1400 var yi: i64 = 0 1401 while yi < MLDSA_L { 1402 _expand_mask_entry((y_vec as i64 + yi * POLY_BUF) as *u8, 1403 rho_prime, kappa + yi, keccak_state, shake_buf) 1404 yi = yi + 1 1405 } 1406 // 2. w = INVNTT(A_hat ∘ NTT(y)) per row 1407 // First NTT each y entry (we keep an NTT-form copy in y_vec since we 1408 // need y itself later for z = y + cs1). Use cs1 region as temp NTT 1409 // form of y (cs1 isn't yet computed). 1410 var yi2: i64 = 0 1411 while yi2 < MLDSA_L { 1412 // Copy y -> cs1[yi2]; NTT in place; then matrix-mul reads cs1. 1413 var cj: i64 = 0 1414 while cj < MLDSA_N { 1415 _i32_store_le((cs1 as i64 + yi2 * POLY_BUF) as *u8, cj, 1416 _i32_load_le((y_vec as i64 + yi2 * POLY_BUF) as *u8, cj)) 1417 cj = cj + 1 1418 } 1419 nx_mldsa_ntt((cs1 as i64 + yi2 * POLY_BUF) as *u8) 1420 yi2 = yi2 + 1 1421 } 1422 // For each row r, accumulate A[r,j] ∘ NTT(y[j]) -> w[r] in NTT form, 1423 // then INVNTT -> canonical. 1424 var rr: i64 = 0 1425 while rr < MLDSA_K { 1426 var cz: i64 = 0 1427 while cz < MLDSA_N { _i32_store_le(acc_poly, cz, 0); cz = cz + 1 } 1428 var cj: i64 = 0 1429 while cj < MLDSA_L { 1430 _expand_a_entry(A_entry, rho, cj, rr, keccak_state, shake_buf) 1431 var pk: i64 = 0 1432 while pk < MLDSA_N { 1433 let pv: i64 = _fqmul(_i32_load_le(A_entry, pk), 1434 _i32_load_le((cs1 as i64 + cj * POLY_BUF) as *u8, pk)) 1435 let av: i64 = _i32_load_le(acc_poly, pk) + pv 1436 _i32_store_le(acc_poly, pk, av) 1437 pk = pk + 1 1438 } 1439 cj = cj + 1 1440 } 1441 // PQClean polyveck_reduce(&w1) before polyveck_invntt_tomont(&w1). 1442 // Without this, the 256x amplification inside invntt overflows i32 1443 // (deterministically in WASM, but to DIFFERENT values than verify's 1444 // properly-reduced path -- causing the L137 c_tilde mismatch bug). 1445 var redi: i64 = 0 1446 while redi < MLDSA_N { 1447 _i32_store_le(acc_poly, redi, _reduce32(_i32_load_le(acc_poly, redi))) 1448 redi = redi + 1 1449 } 1450 nx_mldsa_invntt(acc_poly) 1451 // canonicalise -> w[rr] 1452 var wk: i64 = 0 1453 while wk < MLDSA_N { 1454 var cv: i64 = _reduce32(_i32_load_le(acc_poly, wk)) 1455 if cv < 0 { cv = cv + MLDSA_Q } 1456 _i32_store_le((w_vec as i64 + rr * POLY_BUF) as *u8, wk, cv) 1457 wk = wk + 1 1458 } 1459 rr = rr + 1 1460 } 1461 // 3. (w1, w0) = Decompose(w) 1462 var di: i64 = 0 1463 while di < MLDSA_K { 1464 nx_mldsa_poly_decompose((w1_vec as i64 + di * POLY_BUF) as *u8, 1465 (w0_vec as i64 + di * POLY_BUF) as *u8, 1466 (w_vec as i64 + di * POLY_BUF) as *u8) 1467 di = di + 1 1468 } 1469 // 4. c_tilde = SHAKE256(mu || w1Encode(w1), 48) 1470 // w1 encoded as k * 128 bytes (4 bits/coef). 1471 var w1pi: i64 = 0 1472 while w1pi < MLDSA_K { 1473 _pack_w1_poly((w1_packed as i64 + w1pi * 128) as *u8, 1474 (w1_vec as i64 + w1pi * POLY_BUF) as *u8) 1475 w1pi = w1pi + 1 1476 } 1477 // Build mu || w1_packed in mu_msg_buf scratch (reuse since msg done). 1478 var mui: i64 = 0 1479 while mui < 64 { mu_msg_buf[mui] = mu[mui]; mui = mui + 1 } 1480 var wpi: i64 = 0 1481 let w1_total: i64 = MLDSA_K * 128 1482 while wpi < w1_total { mu_msg_buf[64 + wpi] = w1_packed[wpi]; wpi = wpi + 1 } 1483 _shake256(mu_msg_buf, 64 + w1_total, keccak_state, c_tilde, 48) 1484 1485 // 5. c = SampleInBall(c_tilde, 48 bytes); c_hat = NTT(c) 1486 nx_mldsa_sample_in_ball(c_poly, c_tilde, keccak_state) 1487 var ci: i64 = 0 1488 while ci < MLDSA_N { 1489 _i32_store_le(c_hat, ci, _i32_load_le(c_poly, ci)) 1490 ci = ci + 1 1491 } 1492 nx_mldsa_ntt(c_hat) 1493 1494 // 6. cs1 = INVNTT(c_hat ∘ s1_hat); z = y + cs1 1495 var s1i: i64 = 0 1496 while s1i < MLDSA_L { 1497 var pk: i64 = 0 1498 while pk < MLDSA_N { 1499 let pv: i64 = _fqmul(_i32_load_le(c_hat, pk), 1500 _i32_load_le((s1_vec as i64 + s1i * POLY_BUF) as *u8, pk)) 1501 _i32_store_le((cs1 as i64 + s1i * POLY_BUF) as *u8, pk, pv) 1502 pk = pk + 1 1503 } 1504 nx_mldsa_invntt((cs1 as i64 + s1i * POLY_BUF) as *u8) 1505 // z = y + cs1 1506 var zk: i64 = 0 1507 while zk < MLDSA_N { 1508 let s: i64 = _i32_load_le((y_vec as i64 + s1i * POLY_BUF) as *u8, zk) + _i32_load_le((cs1 as i64 + s1i * POLY_BUF) as *u8, zk) 1509 let r: i64 = _reduce32(s) 1510 _i32_store_le((z_vec as i64 + s1i * POLY_BUF) as *u8, zk, r) 1511 zk = zk + 1 1512 } 1513 s1i = s1i + 1 1514 } 1515 // 7. if ||z||_inf >= gamma1 - beta: reject 1516 if _vec_inf_norm(z_vec, MLDSA_L) >= (MLDSA_GAMMA1 - MLDSA_BETA) { 1517 kappa = kappa + MLDSA_L 1518 attempt = attempt + 1 1519 // continue 1520 } else { 1521 // 8. cs2 = INVNTT(c_hat ∘ s2_hat); r0 = w0 - cs2 (per PQClean 1522 // Dilithium reference: subtract cs2 directly from w0=LowBits(w), 1523 // NOT decompose(w - cs2). These differ when (w0 - cs2) overflows 1524 // the gamma2 envelope, and the hint parameterization in step 12 1525 // requires the PQClean form. 1526 var s2i: i64 = 0 1527 while s2i < MLDSA_K { 1528 var pk: i64 = 0 1529 while pk < MLDSA_N { 1530 let pv: i64 = _fqmul(_i32_load_le(c_hat, pk), 1531 _i32_load_le((s2_vec as i64 + s2i * POLY_BUF) as *u8, pk)) 1532 _i32_store_le((cs2 as i64 + s2i * POLY_BUF) as *u8, pk, pv) 1533 pk = pk + 1 1534 } 1535 nx_mldsa_invntt((cs2 as i64 + s2i * POLY_BUF) as *u8) 1536 // r0[s2i] = reduce32(w0[s2i] - cs2[s2i]) 1537 var dk: i64 = 0 1538 while dk < MLDSA_N { 1539 let w0v: i64 = _i32_load_le((w0_vec as i64 + s2i * POLY_BUF) as *u8, dk) 1540 let cs2v: i64 = _i32_load_le((cs2 as i64 + s2i * POLY_BUF) as *u8, dk) 1541 let rv: i64 = _reduce32(w0v - cs2v) 1542 _i32_store_le((r0_vec as i64 + s2i * POLY_BUF) as *u8, dk, rv) 1543 dk = dk + 1 1544 } 1545 s2i = s2i + 1 1546 } 1547 // 9. if ||r0||_inf >= gamma2 - beta: reject 1548 if _vec_inf_norm(r0_vec, MLDSA_K) >= (MLDSA_GAMMA2 - MLDSA_BETA) { 1549 kappa = kappa + MLDSA_L 1550 attempt = attempt + 1 1551 } else { 1552 // 10. ct0 = INVNTT(c_hat ∘ t0_hat) 1553 var t0i: i64 = 0 1554 while t0i < MLDSA_K { 1555 var pk: i64 = 0 1556 while pk < MLDSA_N { 1557 let pv: i64 = _fqmul(_i32_load_le(c_hat, pk), 1558 _i32_load_le((t0_vec as i64 + t0i * POLY_BUF) as *u8, pk)) 1559 _i32_store_le((ct0 as i64 + t0i * POLY_BUF) as *u8, pk, pv) 1560 pk = pk + 1 1561 } 1562 nx_mldsa_invntt((ct0 as i64 + t0i * POLY_BUF) as *u8) 1563 // canonicalise ct0 per PQClean (poly_reduce after invntt_tomont) 1564 var ck: i64 = 0 1565 while ck < MLDSA_N { 1566 let cv: i64 = _reduce32(_i32_load_le((ct0 as i64 + t0i * POLY_BUF) as *u8, ck)) 1567 _i32_store_le((ct0 as i64 + t0i * POLY_BUF) as *u8, ck, cv) 1568 ck = ck + 1 1569 } 1570 t0i = t0i + 1 1571 } 1572 // 11. if ||ct0||_inf >= gamma2: reject 1573 if _vec_inf_norm(ct0, MLDSA_K) >= MLDSA_GAMMA2 { 1574 kappa = kappa + MLDSA_L 1575 attempt = attempt + 1 1576 } else { 1577 // 12. h = MakeHint(r0 + ct0, w1) per PQClean parameterization 1578 // (sign and verify must match; verify uses the same hint 1579 // against UseHint(h, A*z - c*t1*2^d) to recover w1 = HighBits(w)). 1580 var hi: i64 = 0 1581 var total_weight: i64 = 0 1582 while hi < MLDSA_K { 1583 // hint_a0 = r0 + ct0 (reduce32 to keep in canonical range) 1584 let hint_a0: *u8 = (cs1 as i64 + hi * POLY_BUF) as *u8 // reuse cs1 slot 1585 var nk: i64 = 0 1586 while nk < MLDSA_N { 1587 let r0v: i64 = _i32_load_le((r0_vec as i64 + hi * POLY_BUF) as *u8, nk) 1588 let ct0v: i64 = _i32_load_le((ct0 as i64 + hi * POLY_BUF) as *u8, nk) 1589 let sv: i64 = _reduce32(r0v + ct0v) 1590 _i32_store_le(hint_a0, nk, sv) 1591 nk = nk + 1 1592 } 1593 let weight: i64 = nx_mldsa_poly_make_hint( 1594 (h_vec as i64 + hi * POLY_BUF) as *u8, 1595 hint_a0, 1596 (w1_vec as i64 + hi * POLY_BUF) as *u8) 1597 total_weight = total_weight + weight 1598 hi = hi + 1 1599 } 1600 // 13. if total_weight > omega: reject 1601 if total_weight > MLDSA_OMEGA { 1602 kappa = kappa + MLDSA_L 1603 attempt = attempt + 1 1604 } else { 1605 // 14. sigEncode: c_tilde (48) || pack_z(z) (l*640=3200) 1606 // || pack_hint(h) (61) = 3309 bytes 1607 var sii: i64 = 0 1608 while sii < 48 { sig_out[sii] = c_tilde[sii]; sii = sii + 1 } 1609 var zi: i64 = 0 1610 while zi < MLDSA_L { 1611 _pack_z_poly((sig_out as i64 + 48 + zi * 640) as *u8, 1612 (z_vec as i64 + zi * POLY_BUF) as *u8) 1613 zi = zi + 1 1614 } 1615 let pack_rc: i64 = _pack_hint_vec((sig_out as i64 + 48 + MLDSA_L * 640) as *u8, 1616 h_vec) 1617 if pack_rc < 0 { return -3 } 1618 return 0 1619 } 1620 } 1621 } 1622 } 1623 } 1624 return -1 // rejection loop exhausted 1625} 1626 1627// L136 KAT export: sign a fixed message + check basic structural properties. 1628// Returns 0 on success, non-zero bitmap. 1629// bit 0: keygen failed 1630// bit 1: sign failed 1631// bit 2: sig[0..48] (c_tilde) all zero 1632// bit 3: sig[48..3248] (z packed) all zero 1633// bit 4: different messages -> same signature (would mean nondeterministic-bug) 1634func nx_mldsa_sign_test(seed_32: *u8, msg_32: *u8, scratch: *u8) -> i64 { 1635 let pk: *u8 = (scratch as i64 + 110000) as *u8 1636 let sk: *u8 = (scratch as i64 + 112000) as *u8 1637 let sig1: *u8 = (scratch as i64 + 117000) as *u8 1638 let sig2: *u8 = (scratch as i64 + 121000) as *u8 1639 var fail: i64 = 0 1640 let rc_kg: i64 = nx_mldsa_keygen(seed_32, scratch, pk, sk) 1641 if rc_kg != 0 { fail = fail | 1 } 1642 let rc_s1: i64 = nx_mldsa_sign(sk, msg_32, 32, scratch, sig1) 1643 if rc_s1 != 0 { fail = fail | 2 } 1644 // structural: c_tilde must not be all zero 1645 var c_zero: i64 = 1 1646 var i: i64 = 0 1647 while i < 48 { 1648 if sig1[i] != 0 { c_zero = 0 } 1649 i = i + 1 1650 } 1651 if c_zero != 0 { fail = fail | 4 } 1652 // z packed must not be all zero 1653 var z_zero: i64 = 1 1654 var j: i64 = 48 1655 while j < 48 + MLDSA_L * 640 { 1656 if sig1[j] != 0 { z_zero = 0 } 1657 j = j + 1 1658 } 1659 if z_zero != 0 { fail = fail | 8 } 1660 // sign a DIFFERENT msg, check sig differs 1661 let alt_msg: *u8 = (scratch as i64 + 126000) as *u8 1662 var a: i64 = 0 1663 while a < 32 { alt_msg[a] = msg_32[a] ^ 0xff; a = a + 1 } 1664 let rc_s2: i64 = nx_mldsa_sign(sk, alt_msg, 32, scratch, sig2) 1665 if rc_s2 != 0 { fail = fail | 2 } 1666 var same: i64 = 1 1667 var k: i64 = 0 1668 while k < 100 { 1669 if sig1[k] != sig2[k] { same = 0 } 1670 k = k + 1 1671 } 1672 if same != 0 { fail = fail | 16 } 1673 return fail 1674} 1675 1676// ============================================================================ 1677// SECTION 5: KAT export -- NTT round-trip self-test 1678// ============================================================================ 1679// Verifies INVNTT(NTT(poly)) ≡ poly (mod q) for a deterministic poly. 1680// Per PQClean convention: invntt_tomont leaves the result multiplied by 1681// the Montgomery factor; we account by checking equivalence mod q 1682// directly (the caller fills the input and compares post-round-trip 1683// via reduce32+canon mod q). 1684// 1685// scratch needs 1024 bytes (one poly buffer). 1686// Returns 0 on success, non-zero bitmap of failure axes. 1687// === L137 ML-DSA-65 verify === 1688// 1689// Unpack 256 10-bit coefficients (range [0, 2^10)) from 320 bytes. 1690// Inverse of _pack_t1_poly. Outputs as i32 LE in `out_poly` (1024 B). 1691func _unpack_t1_poly(out_poly: *u8, packed: *u8) -> i64 { 1692 var i: i64 = 0 1693 while i < 64 { // 64 groups of 4 coefs = 256 1694 let off: i64 = i * 5 1695 let b0: i64 = packed[off + 0] 1696 let b1: i64 = packed[off + 1] 1697 let b2: i64 = packed[off + 2] 1698 let b3: i64 = packed[off + 3] 1699 let b4: i64 = packed[off + 4] 1700 let c0: i64 = (b0 | (b1 << 8)) & 0x3ff 1701 let c1: i64 = ((b1 >> 2) | (b2 << 6)) & 0x3ff 1702 let c2: i64 = ((b2 >> 4) | (b3 << 4)) & 0x3ff 1703 let c3: i64 = ((b3 >> 6) | (b4 << 2)) & 0x3ff 1704 _i32_store_le(out_poly, i*4 + 0, c0) 1705 _i32_store_le(out_poly, i*4 + 1, c1) 1706 _i32_store_le(out_poly, i*4 + 2, c2) 1707 _i32_store_le(out_poly, i*4 + 3, c3) 1708 i = i + 1 1709 } 1710 return 0 1711} 1712 1713// nx_mldsa_verify -- FIPS 204 §8 Verify(pk, M, sigma) -> {0 (valid), 1 (invalid)} 1714// Returns 0 on accept, positive bitmap on reject (so callers can debug): 1715// bit 0: z norm out of range 1716// bit 1: hint vec malformed (weight > omega or unpack failed) 1717// bit 2: c_tilde mismatch (the cryptographic check) 1718// 1719// pk layout (1952 B): rho(32) || t1_packed(K*320 = 1920) 1720// sig layout (3309 B): c_tilde(48) || z_packed(L*640 = 3200) || hint_packed(61) 1721// 1722// Scratch: requires >= 96 KiB. Layout documented per scratch_range below. 1723func nx_mldsa_verify(pk: *u8, msg: *u8, msg_len: i64, 1724 sig: *u8, scratch: *u8) -> i64 { 1725 let keccak_state: *u8 = scratch // scratch_range: keccak 0..439 1726 let shake_buf: *u8 = (scratch as i64 + 512) as *u8 // scratch_range: shake_buf 512..2559 1727 // Unpacked vectors 1728 let t1_vec: *u8 = (scratch as i64 + 3072) as *u8 // scratch_range: t1_vec 3072..9215 (6 * 1024) 1729 let z_vec: *u8 = (scratch as i64 + 9216) as *u8 // scratch_range: z_vec 9216..14335 (5 * 1024) 1730 let h_vec: *u8 = (scratch as i64 + 14336) as *u8 // scratch_range: h_vec 14336..20479 (6 * 1024) 1731 // Working singletons 1732 let A_entry: *u8 = (scratch as i64 + 20480) as *u8 // scratch_range: A_entry 20480..21503 1733 let acc_poly: *u8 = (scratch as i64 + 21504) as *u8 // scratch_range: acc 21504..22527 1734 let c_poly: *u8 = (scratch as i64 + 22528) as *u8 // scratch_range: c_poly 22528..23551 1735 let c_hat: *u8 = (scratch as i64 + 23552) as *u8 // scratch_range: c_hat 23552..24575 1736 let w1_poly: *u8 = (scratch as i64 + 24576) as *u8 // scratch_range: w1_poly 24576..25599 1737 // Hashing scratch 1738 let tr: *u8 = (scratch as i64 + 25600) as *u8 // scratch_range: tr 25600..25663 (64) 1739 let mu: *u8 = (scratch as i64 + 25664) as *u8 // scratch_range: mu 25664..25727 (64) 1740 let mu_msg_buf: *u8 = (scratch as i64 + 25728) as *u8 // scratch_range: mu_msg 25728..33919 (8192 cap) 1741 let w1_packed: *u8 = (scratch as i64 + 33920) as *u8 // scratch_range: w1_pkd 33920..34687 (k*128 = 768) 1742 let c_tilde_rec: *u8 = (scratch as i64 + 34688) as *u8 // scratch_range: c_til_r 34688..34735 (48) 1743 1744 let rho: *u8 = pk 1745 let t1_packed: *u8 = (pk as i64 + 32) as *u8 1746 let c_tilde: *u8 = sig 1747 let z_packed: *u8 = (sig as i64 + 48) as *u8 1748 let h_packed: *u8 = (sig as i64 + 48 + MLDSA_L * 640) as *u8 1749 1750 var fail_mask: i64 = 0 1751 1752 // 1. Decode z; check ||z||_inf bound. 1753 var zi: i64 = 0 1754 while zi < MLDSA_L { 1755 _unpack_z_poly((z_vec as i64 + zi * POLY_BUF) as *u8, 1756 (z_packed as i64 + zi * 640) as *u8) 1757 zi = zi + 1 1758 } 1759 if _vec_inf_norm(z_vec, MLDSA_L) >= (MLDSA_GAMMA1 - MLDSA_BETA) { 1760 fail_mask = fail_mask | 1 1761 } 1762 1763 // 2. Decode hint; rejects if encoding malformed or weight > omega. 1764 if _unpack_hint_vec(h_vec, h_packed) != 0 { 1765 fail_mask = fail_mask | 2 1766 } 1767 1768 // 3. Decode t1. 1769 var ti: i64 = 0 1770 while ti < MLDSA_K { 1771 _unpack_t1_poly((t1_vec as i64 + ti * POLY_BUF) as *u8, 1772 (t1_packed as i64 + ti * 320) as *u8) 1773 ti = ti + 1 1774 } 1775 1776 // 4. tr = SHAKE256(pk, 64). 1777 _shake256(pk, 1952, keccak_state, tr, 64) 1778 1779 // 5. mu = SHAKE256(tr || msg, 64). msg_len capped at 8128. 1780 if msg_len > 8128 { return -2 } 1781 var tri: i64 = 0 1782 while tri < 64 { mu_msg_buf[tri] = tr[tri]; tri = tri + 1 } 1783 var mi: i64 = 0 1784 while mi < msg_len { mu_msg_buf[64 + mi] = msg[mi]; mi = mi + 1 } 1785 _shake256(mu_msg_buf, 64 + msg_len, keccak_state, mu, 64) 1786 1787 // 6. c = SampleInBall(c_tilde) ; c_hat = NTT(c). 1788 nx_mldsa_sample_in_ball(c_poly, c_tilde, keccak_state) 1789 var ci: i64 = 0 1790 while ci < MLDSA_N { 1791 _i32_store_le(c_hat, ci, _i32_load_le(c_poly, ci)) 1792 ci = ci + 1 1793 } 1794 nx_mldsa_ntt(c_hat) 1795 1796 // 7. NTT every z poly in place. z_vec slots become z_hat slots. 1797 var zh: i64 = 0 1798 while zh < MLDSA_L { 1799 nx_mldsa_ntt((z_vec as i64 + zh * POLY_BUF) as *u8) 1800 zh = zh + 1 1801 } 1802 1803 // 8. NTT every (t1 * 2^d) poly in place. t1_vec slots become t1_hat slots. 1804 var th: i64 = 0 1805 while th < MLDSA_K { 1806 let tp: *u8 = (t1_vec as i64 + th * POLY_BUF) as *u8 1807 var sk: i64 = 0 1808 while sk < MLDSA_N { 1809 _i32_store_le(tp, sk, _i32_load_le(tp, sk) << MLDSA_D) 1810 sk = sk + 1 1811 } 1812 // L137 debug: snapshot t1_vec[0] AFTER shift, BEFORE NTT 1813 if th == 0 { 1814 let dbg_t1_shifted: i64 = scratch as i64 + 205120 1815 var ds: i64 = 0 1816 while ds < MLDSA_N { 1817 _i32_store_le(dbg_t1_shifted as *u8, ds, _i32_load_le(tp, ds)) 1818 ds = ds + 1 1819 } 1820 } 1821 nx_mldsa_ntt(tp) 1822 // L137 debug: snapshot t1_vec[0] AFTER NTT 1823 if th == 0 { 1824 let dbg_t1_ntt: i64 = scratch as i64 + 206144 1825 var dn: i64 = 0 1826 while dn < MLDSA_N { 1827 _i32_store_le(dbg_t1_ntt as *u8, dn, _i32_load_le(tp, dn)) 1828 dn = dn + 1 1829 } 1830 } 1831 th = th + 1 1832 } 1833 1834 // 9. For each row i in K: 1835 // acc = Σ_j A[i][j] * z_hat[j] - c_hat * t1_hat[i] 1836 // INVNTT(acc) -> approx of w 1837 // w1' = UseHint(h[i], approx_w_canonicalised) 1838 // pack w1' into w1_packed[i*128 .. ] 1839 var ri: i64 = 0 1840 while ri < MLDSA_K { 1841 // acc = 0 1842 var z0: i64 = 0 1843 while z0 < MLDSA_N { _i32_store_le(acc_poly, z0, 0); z0 = z0 + 1 } 1844 // Sum Σ_j A[i][j] * z_hat[j] 1845 var jj: i64 = 0 1846 while jj < MLDSA_L { 1847 _expand_a_entry(A_entry, rho, jj, ri, keccak_state, shake_buf) 1848 var pk_i: i64 = 0 1849 while pk_i < MLDSA_N { 1850 let prod: i64 = _fqmul(_i32_load_le(A_entry, pk_i), 1851 _i32_load_le((z_vec as i64 + jj * POLY_BUF) as *u8, pk_i)) 1852 _i32_store_le(acc_poly, pk_i, 1853 _i32_load_le(acc_poly, pk_i) + prod) 1854 pk_i = pk_i + 1 1855 } 1856 jj = jj + 1 1857 } 1858 // L137 debug: save row-0 acc AFTER Σ A*z but BEFORE the c*t1 subtract. 1859 if ri == 0 { 1860 let dbg_az: i64 = scratch as i64 + 203072 1861 var dai: i64 = 0 1862 while dai < MLDSA_N { 1863 _i32_store_le(dbg_az as *u8, dai, _i32_load_le(acc_poly, dai)) 1864 dai = dai + 1 1865 } 1866 } 1867 // Subtract c_hat * t1_hat[i] 1868 var sk2: i64 = 0 1869 while sk2 < MLDSA_N { 1870 let pv: i64 = _fqmul(_i32_load_le(c_hat, sk2), 1871 _i32_load_le((t1_vec as i64 + ri * POLY_BUF) as *u8, sk2)) 1872 // L137 debug: save c_hat * t1_hat[0] (the value being subtracted) for row 0 1873 if ri == 0 { 1874 let dbg_ct1: i64 = scratch as i64 + 204096 1875 _i32_store_le(dbg_ct1 as *u8, sk2, pv) 1876 } 1877 _i32_store_le(acc_poly, sk2, 1878 _i32_load_le(acc_poly, sk2) - pv) 1879 sk2 = sk2 + 1 1880 } 1881 // Reduce32 in NTT domain BEFORE invntt -- PQClean 1882 // polyveck_reduce(&w1) before polyveck_invntt_tomont(&w1). 1883 var redi: i64 = 0 1884 while redi < MLDSA_N { 1885 _i32_store_le(acc_poly, redi, _reduce32(_i32_load_le(acc_poly, redi))) 1886 redi = redi + 1 1887 } 1888 // INVNTT in place -> approx w 1889 nx_mldsa_invntt(acc_poly) 1890 // Canonicalise to [0, q) 1891 var ck: i64 = 0 1892 while ck < MLDSA_N { 1893 var v: i64 = _reduce32(_i32_load_le(acc_poly, ck)) 1894 if v < 0 { v = v + MLDSA_Q } 1895 _i32_store_le(acc_poly, ck, v) 1896 ck = ck + 1 1897 } 1898 // L137 debug: save row-0 acc_poly to scratch+202048 for JS inspect. 1899 if ri == 0 { 1900 let dbg_off: i64 = scratch as i64 + 202048 1901 var di: i64 = 0 1902 while di < MLDSA_N { 1903 _i32_store_le(dbg_off as *u8, di, _i32_load_le(acc_poly, di)) 1904 di = di + 1 1905 } 1906 } 1907 // w1' = UseHint(h[i], acc) 1908 nx_mldsa_poly_use_hint(w1_poly, acc_poly, 1909 (h_vec as i64 + ri * POLY_BUF) as *u8) 1910 // Pack into w1_packed. 1911 _pack_w1_poly((w1_packed as i64 + ri * 128) as *u8, w1_poly) 1912 ri = ri + 1 1913 } 1914 1915 // 10. c_tilde' = SHAKE256(mu || w1_packed, 48). Reuse mu_msg_buf. 1916 var mui: i64 = 0 1917 while mui < 64 { mu_msg_buf[mui] = mu[mui]; mui = mui + 1 } 1918 var wpi: i64 = 0 1919 let w1_total: i64 = MLDSA_K * 128 1920 while wpi < w1_total { mu_msg_buf[64 + wpi] = w1_packed[wpi]; wpi = wpi + 1 } 1921 _shake256(mu_msg_buf, 64 + w1_total, keccak_state, c_tilde_rec, 48) 1922 1923 // 11. Constant-byte compare c_tilde_rec vs c_tilde. 1924 var diff: i64 = 0 1925 var cti: i64 = 0 1926 while cti < 48 { 1927 if c_tilde_rec[cti] != c_tilde[cti] { diff = 1 } 1928 cti = cti + 1 1929 } 1930 if diff != 0 { fail_mask = fail_mask | 4 } 1931 1932 return fail_mask 1933} 1934 1935// L137 KAT export: keygen -> sign -> verify round-trip. Returns 0 on 1936// success, non-zero bitmap on failure. 1937// bit 0: keygen failed 1938// bit 1: sign failed 1939// bit 2: verify rejected the legitimate signature 1940// bit 3: verify accepted a tampered signature (flipped sig byte 0) 1941// bit 4: verify accepted a wrong-message signature 1942func nx_mldsa_verify_test(seed_32: *u8, msg_32: *u8, scratch: *u8) -> i64 { 1943 let pk: *u8 = (scratch as i64 + 100000) as *u8 // scratch_range: pk 100000..101951 (1952) 1944 let sk: *u8 = (scratch as i64 + 102000) as *u8 // scratch_range: sk 102000..106031 (4032) 1945 let sig: *u8 = (scratch as i64 + 106500) as *u8 // scratch_range: sig 106500..109808 (3309) 1946 let inner: *u8 = scratch // first 96K used by keygen/sign/verify ops 1947 1948 var fail: i64 = 0 1949 1950 // 1. keygen 1951 if nx_mldsa_keygen(seed_32, inner, pk, sk) != 0 { fail = fail | 1; return fail } 1952 1953 // 2. sign 1954 if nx_mldsa_sign(sk, msg_32, 32, inner, sig) != 0 { fail = fail | 2; return fail } 1955 1956 // 3. verify legitimate -> expect 0 1957 let v_ok: i64 = nx_mldsa_verify(pk, msg_32, 32, sig, inner) 1958 if v_ok != 0 { fail = fail | 4 } 1959 1960 // 4. tamper: flip a byte of c_tilde, expect non-zero 1961 let orig_byte: i64 = sig[0] 1962 sig[0] = (orig_byte ^ 0xff) & 0xff 1963 let v_tamper: i64 = nx_mldsa_verify(pk, msg_32, 32, sig, inner) 1964 sig[0] = orig_byte 1965 if v_tamper == 0 { fail = fail | 8 } 1966 1967 // 5. wrong message: flip a byte of msg, expect non-zero 1968 let orig_mb: i64 = msg_32[0] 1969 msg_32[0] = (orig_mb ^ 0xff) & 0xff 1970 let v_wrong: i64 = nx_mldsa_verify(pk, msg_32, 32, sig, inner) 1971 msg_32[0] = orig_mb 1972 if v_wrong == 0 { fail = fail | 16 } 1973 1974 return fail 1975} 1976 1977// Debug: L=5 accumulated basemul + invntt with VERIFY-realistic magnitudes. 1978// a values in [0, q) like A entries from rej_uniform. 1979// b values in (-gamma1, gamma1] like z from unpack_z. 1980// Returns # mismatched coefs. 1981func nx_mldsa_acc_basemul_large_test(scratch: *u8) -> i64 { 1982 let L: i64 = 5 1983 let a: *u8 = scratch 1984 let b: *u8 = (scratch as i64 + L * POLY_BUF) as *u8 1985 let a_ntt: *u8 = (scratch as i64 + 2 * L * POLY_BUF) as *u8 1986 let b_ntt: *u8 = (scratch as i64 + 3 * L * POLY_BUF) as *u8 1987 let acc: *u8 = (scratch as i64 + 4 * L * POLY_BUF) as *u8 1988 let direct: *u8 = (scratch as i64 + 4 * L * POLY_BUF + POLY_BUF) as *u8 1989 var p: i64 = 0 1990 while p < L { 1991 var i: i64 = 0 1992 while i < MLDSA_N { 1993 // a in [0, q): pseudo-uniform pattern 1994 let av: i64 = (i * 12345 + p * 7891) % MLDSA_Q 1995 // b signed in (-gamma1, gamma1] 1996 var bv: i64 = ((i * 31 + p * 17) % (2 * MLDSA_GAMMA1)) - MLDSA_GAMMA1 + 1 1997 _i32_store_le((a as i64 + p * POLY_BUF) as *u8, i, av) 1998 _i32_store_le((b as i64 + p * POLY_BUF) as *u8, i, bv) 1999 _i32_store_le((a_ntt as i64 + p * POLY_BUF) as *u8, i, av) 2000 _i32_store_le((b_ntt as i64 + p * POLY_BUF) as *u8, i, bv) 2001 i = i + 1 2002 } 2003 nx_mldsa_ntt((a_ntt as i64 + p * POLY_BUF) as *u8) 2004 nx_mldsa_ntt((b_ntt as i64 + p * POLY_BUF) as *u8) 2005 p = p + 1 2006 } 2007 var z: i64 = 0 2008 while z < MLDSA_N { _i32_store_le(acc, z, 0); z = z + 1 } 2009 var pj: i64 = 0 2010 while pj < L { 2011 var k: i64 = 0 2012 while k < MLDSA_N { 2013 let pv: i64 = _fqmul( 2014 _i32_load_le((a_ntt as i64 + pj * POLY_BUF) as *u8, k), 2015 _i32_load_le((b_ntt as i64 + pj * POLY_BUF) as *u8, k)) 2016 _i32_store_le(acc, k, _i32_load_le(acc, k) + pv) 2017 k = k + 1 2018 } 2019 pj = pj + 1 2020 } 2021 var ri: i64 = 0 2022 while ri < MLDSA_N { 2023 _i32_store_le(acc, ri, _reduce32(_i32_load_le(acc, ri))) 2024 ri = ri + 1 2025 } 2026 nx_mldsa_invntt(acc) 2027 var ci: i64 = 0 2028 while ci < MLDSA_N { 2029 var v: i64 = _reduce32(_i32_load_le(acc, ci)) 2030 if v < 0 { v = v + MLDSA_Q } 2031 _i32_store_le(acc, ci, v) 2032 ci = ci + 1 2033 } 2034 // Direct convolution mod q 2035 var dk: i64 = 0 2036 while dk < MLDSA_N { 2037 var sum: i64 = 0 2038 var dp: i64 = 0 2039 while dp < L { 2040 var di: i64 = 0 2041 while di < MLDSA_N { 2042 let dj: i64 = dk - di 2043 let av: i64 = _i32_load_le((a as i64 + dp * POLY_BUF) as *u8, di) 2044 if dj >= 0 { 2045 let bv: i64 = _i32_load_le((b as i64 + dp * POLY_BUF) as *u8, dj) 2046 // Reduce per-term to avoid sum overflow (av * bv can be q * gamma1 ~ 4e12) 2047 sum = (sum + (av * bv) % MLDSA_Q) % MLDSA_Q 2048 } else { 2049 let bv: i64 = _i32_load_le((b as i64 + dp * POLY_BUF) as *u8, dj + MLDSA_N) 2050 sum = (sum - (av * bv) % MLDSA_Q) % MLDSA_Q 2051 } 2052 di = di + 1 2053 } 2054 dp = dp + 1 2055 } 2056 var canon: i64 = sum % MLDSA_Q 2057 if canon < 0 { canon = canon + MLDSA_Q } 2058 _i32_store_le(direct, dk, canon) 2059 dk = dk + 1 2060 } 2061 var diffs: i64 = 0 2062 var c: i64 = 0 2063 while c < MLDSA_N { 2064 if _i32_load_le(acc, c) != _i32_load_le(direct, c) { 2065 diffs = diffs + 1 2066 } 2067 c = c + 1 2068 } 2069 return diffs 2070} 2071 2072// Debug: L=5-term ACCUMULATED basemul + invntt round-trip. 2073// This mirrors verify's matrix-row computation: Σ_{j=0..L-1} a[j] * b[j]. 2074// Inputs in standard form; output compared to direct N²L convolution sum. 2075// Returns # mismatched coefs. 2076func nx_mldsa_acc_basemul_round_trip_test(scratch: *u8) -> i64 { 2077 let L: i64 = 5 2078 let a: *u8 = scratch 2079 let b: *u8 = (scratch as i64 + L * POLY_BUF) as *u8 2080 let a_ntt: *u8 = (scratch as i64 + 2 * L * POLY_BUF) as *u8 2081 let b_ntt: *u8 = (scratch as i64 + 3 * L * POLY_BUF) as *u8 2082 let acc: *u8 = (scratch as i64 + 4 * L * POLY_BUF) as *u8 2083 let direct: *u8 = (scratch as i64 + 4 * L * POLY_BUF + POLY_BUF) as *u8 2084 // Fill a[0..L], b[0..L] with small signed pattern 2085 var p: i64 = 0 2086 while p < L { 2087 var i: i64 = 0 2088 while i < MLDSA_N { 2089 var av: i64 = (p * 13 + i * 7 + 3) % 11 2090 av = av - 5 2091 var bv: i64 = (p * 17 + i * 11 + 5) % 11 2092 bv = bv - 5 2093 _i32_store_le((a as i64 + p * POLY_BUF) as *u8, i, av) 2094 _i32_store_le((b as i64 + p * POLY_BUF) as *u8, i, bv) 2095 _i32_store_le((a_ntt as i64 + p * POLY_BUF) as *u8, i, av) 2096 _i32_store_le((b_ntt as i64 + p * POLY_BUF) as *u8, i, bv) 2097 i = i + 1 2098 } 2099 nx_mldsa_ntt((a_ntt as i64 + p * POLY_BUF) as *u8) 2100 nx_mldsa_ntt((b_ntt as i64 + p * POLY_BUF) as *u8) 2101 p = p + 1 2102 } 2103 // Accumulate Σ basemul(a_ntt[p], b_ntt[p]) into acc (NTT form) 2104 var z: i64 = 0 2105 while z < MLDSA_N { _i32_store_le(acc, z, 0); z = z + 1 } 2106 var pj: i64 = 0 2107 while pj < L { 2108 var k: i64 = 0 2109 while k < MLDSA_N { 2110 let pv: i64 = _fqmul( 2111 _i32_load_le((a_ntt as i64 + pj * POLY_BUF) as *u8, k), 2112 _i32_load_le((b_ntt as i64 + pj * POLY_BUF) as *u8, k)) 2113 _i32_store_le(acc, k, _i32_load_le(acc, k) + pv) 2114 k = k + 1 2115 } 2116 pj = pj + 1 2117 } 2118 // reduce32, invntt, canonicalize 2119 var ri: i64 = 0 2120 while ri < MLDSA_N { 2121 _i32_store_le(acc, ri, _reduce32(_i32_load_le(acc, ri))) 2122 ri = ri + 1 2123 } 2124 nx_mldsa_invntt(acc) 2125 var ci: i64 = 0 2126 while ci < MLDSA_N { 2127 var v: i64 = _reduce32(_i32_load_le(acc, ci)) 2128 if v < 0 { v = v + MLDSA_Q } 2129 _i32_store_le(acc, ci, v) 2130 ci = ci + 1 2131 } 2132 // Direct Σ_p (a[p] * b[p] cyclic in (X^N+1)) per coef 2133 var dk: i64 = 0 2134 while dk < MLDSA_N { 2135 var sum: i64 = 0 2136 var dp: i64 = 0 2137 while dp < L { 2138 var di: i64 = 0 2139 while di < MLDSA_N { 2140 let dj: i64 = dk - di 2141 let av: i64 = _i32_load_le((a as i64 + dp * POLY_BUF) as *u8, di) 2142 if dj >= 0 { 2143 let bv: i64 = _i32_load_le((b as i64 + dp * POLY_BUF) as *u8, dj) 2144 sum = sum + av * bv 2145 } else { 2146 let bv: i64 = _i32_load_le((b as i64 + dp * POLY_BUF) as *u8, dj + MLDSA_N) 2147 sum = sum - av * bv 2148 } 2149 di = di + 1 2150 } 2151 dp = dp + 1 2152 } 2153 var canon: i64 = sum % MLDSA_Q 2154 if canon < 0 { canon = canon + MLDSA_Q } 2155 _i32_store_le(direct, dk, canon) 2156 dk = dk + 1 2157 } 2158 var diffs: i64 = 0 2159 var c: i64 = 0 2160 while c < MLDSA_N { 2161 if _i32_load_le(acc, c) != _i32_load_le(direct, c) { 2162 diffs = diffs + 1 2163 } 2164 c = c + 1 2165 } 2166 return diffs 2167} 2168 2169// Debug: basemul + invntt round-trip. 2170// Compute a, b polys in standard form; do ntt(a), ntt(b), basemul, invntt, 2171// canonicalize. Compare to (a*b mod (X^N+1)) standard form, computed 2172// directly by O(N^2) convolution in Zq. Returns # mismatched coefs. 2173// 2174// If THIS fails but ntt-only round-trip passes, the bug is in basemul 2175// or in invntt's handling of post-basemul Montgomery-form values. 2176func nx_mldsa_basemul_round_trip_test(scratch: *u8) -> i64 { 2177 let a: *u8 = scratch 2178 let b: *u8 = (scratch as i64 + POLY_BUF) as *u8 2179 let a_ntt: *u8 = (scratch as i64 + 2 * POLY_BUF) as *u8 2180 let b_ntt: *u8 = (scratch as i64 + 3 * POLY_BUF) as *u8 2181 let prod_n: *u8 = (scratch as i64 + 4 * POLY_BUF) as *u8 // NTT(a*b) 2182 let prod_d: *u8 = (scratch as i64 + 5 * POLY_BUF) as *u8 // direct a*b mod (X^N+1) 2183 // Fill a, b with small signed pattern 2184 var i: i64 = 0 2185 while i < MLDSA_N { 2186 var av: i64 = (i * 7 + 3) % 9 2187 av = av - 4 // a_i in [-4, 4] 2188 var bv: i64 = (i * 11 + 5) % 9 2189 bv = bv - 4 // b_i in [-4, 4] 2190 _i32_store_le(a, i, av) 2191 _i32_store_le(b, i, bv) 2192 _i32_store_le(a_ntt, i, av) 2193 _i32_store_le(b_ntt, i, bv) 2194 i = i + 1 2195 } 2196 // NTT both 2197 nx_mldsa_ntt(a_ntt) 2198 nx_mldsa_ntt(b_ntt) 2199 // basemul (per-coef fqmul) 2200 nx_mldsa_basemul(prod_n, a_ntt, b_ntt) 2201 // invntt + canonicalize 2202 nx_mldsa_invntt(prod_n) 2203 var k: i64 = 0 2204 while k < MLDSA_N { 2205 var v: i64 = _reduce32(_i32_load_le(prod_n, k)) 2206 if v < 0 { v = v + MLDSA_Q } 2207 _i32_store_le(prod_n, k, v) 2208 k = k + 1 2209 } 2210 // Direct cyclic convolution: prod_d[k] = Σ_{i+j ≡ k (mod N), with sign flip for wrap} 2211 // For X^N + 1: prod_d[k] = Σ_{i=0..k} a[i]*b[k-i] - Σ_{i=k+1..N-1} a[i]*b[N+k-i] 2212 var dk: i64 = 0 2213 while dk < MLDSA_N { 2214 var sum: i64 = 0 2215 var di: i64 = 0 2216 while di < MLDSA_N { 2217 let dj: i64 = dk - di 2218 let av: i64 = _i32_load_le(a, di) 2219 if dj >= 0 { 2220 let bv: i64 = _i32_load_le(b, dj) 2221 sum = sum + av * bv 2222 } else { 2223 let bv: i64 = _i32_load_le(b, dj + MLDSA_N) 2224 sum = sum - av * bv 2225 } 2226 di = di + 1 2227 } 2228 // Canonicalize sum mod q to [0, q) 2229 var canon: i64 = sum % MLDSA_Q 2230 if canon < 0 { canon = canon + MLDSA_Q } 2231 _i32_store_le(prod_d, dk, canon) 2232 dk = dk + 1 2233 } 2234 // Compare prod_n to prod_d 2235 var diffs: i64 = 0 2236 var c: i64 = 0 2237 while c < MLDSA_N { 2238 if _i32_load_le(prod_n, c) != _i32_load_le(prod_d, c) { 2239 diffs = diffs + 1 2240 } 2241 c = c + 1 2242 } 2243 return diffs 2244} 2245 2246// Debug: NTT round-trip on a poly with MIXED positive/negative coefs. 2247// Existing nx_mldsa_ntt_round_trip_test only covers positive inputs. 2248// z in ML-DSA-65 has range (-gamma1, gamma1] so signed values matter. 2249// Returns 0 if all coefs round-trip to (original * MONT) mod q, else 2250// the number of mismatched coefs. 2251func nx_mldsa_ntt_signed_round_trip_test(scratch: *u8) -> i64 { 2252 let poly: *u8 = scratch 2253 let saved: *u8 = (scratch as i64 + POLY_BUF) as *u8 2254 // Fill with alternating +/- pattern bounded by gamma1 = 524288. 2255 // pattern: ((i * 31 + 17) % gamma1) * (((i & 1) == 0) ? 1 : -1) 2256 var i: i64 = 0 2257 while i < MLDSA_N { 2258 var v: i64 = (i * 31 + 17) % MLDSA_GAMMA1 2259 if (i & 1) != 0 { v = 0 - v } 2260 _i32_store_le(poly, i, v) 2261 _i32_store_le(saved, i, v) 2262 i = i + 1 2263 } 2264 nx_mldsa_ntt(poly) 2265 nx_mldsa_invntt(poly) 2266 let MONT: i64 = 4193792 2267 var diffs: i64 = 0 2268 var j: i64 = 0 2269 while j < MLDSA_N { 2270 var a: i64 = _reduce32(_i32_load_le(poly, j)) 2271 if a < 0 { a = a + MLDSA_Q } 2272 var orig: i64 = _i32_load_le(saved, j) 2273 // (orig * MONT) mod q, canonicalize to [0, q) 2274 var expected: i64 = (orig * MONT) % MLDSA_Q 2275 if expected < 0 { expected = expected + MLDSA_Q } 2276 if a != expected { diffs = diffs + 1 } 2277 j = j + 1 2278 } 2279 return diffs 2280} 2281 2282// Debug: unpack t1 row 0 from a pk (1952 bytes) and write the 256 2283// i32 coefs into out_poly. JS smoke checks that all coefs are in 2284// [0, 1024) -- the range required by spec. 2285func nx_mldsa_unpack_t1_debug(pk: *u8, out_poly: *u8) -> i64 { 2286 let t1_packed: *u8 = (pk as i64 + 32) as *u8 2287 _unpack_t1_poly(out_poly, t1_packed) 2288 return 0 2289} 2290 2291// Debug: take an unpacked t1 row (256 i32 in [0, 1024)), pack it via 2292// the same packer keygen uses, write 320 bytes to out_packed. JS 2293// smoke compares out_packed to pk[32..32+320] to verify round-trip. 2294func nx_mldsa_pack_t1_debug(in_poly: *u8, out_packed: *u8) -> i64 { 2295 _pack_t1_poly(out_packed, in_poly) 2296 return 0 2297} 2298 2299// Debug: keygen + sign + ONE verify, leaving sign's w1_packed at 2300// scratch+91360 and verify's recovered w1_packed at scratch+33920 so 2301// the JS smoke can diff them. Returns the verify fail_mask. 2302func nx_mldsa_verify_debug_one(seed_32: *u8, msg_32: *u8, scratch: *u8) -> i64 { 2303 let pk: *u8 = (scratch as i64 + 100000) as *u8 2304 let sk: *u8 = (scratch as i64 + 102000) as *u8 2305 let sig: *u8 = (scratch as i64 + 106500) as *u8 2306 if nx_mldsa_keygen(seed_32, scratch, pk, sk) != 0 { return -1 } 2307 if nx_mldsa_sign(sk, msg_32, 32, scratch, sig) != 0 { return -2 } 2308 return nx_mldsa_verify(pk, msg_32, 32, sig, scratch) 2309} 2310 2311func nx_mldsa_ntt_round_trip_test(scratch: *u8) -> i64 { 2312 let poly: *u8 = scratch 2313 let saved: *u8 = (scratch as i64 + POLY_BUF) as *u8 2314 // Fill poly[i] = (i * 13 + 7) mod q for a deterministic pattern. 2315 var i: i64 = 0 2316 while i < MLDSA_N { 2317 let v: i64 = (i * 13 + 7) % MLDSA_Q 2318 _i32_store_le(poly, i, v) 2319 _i32_store_le(saved, i, v) 2320 i = i + 1 2321 } 2322 nx_mldsa_ntt(poly) 2323 nx_mldsa_invntt(poly) 2324 // Per PQClean convention: invntt_tomont(ntt(X)) = X * mont (mod q), 2325 // where mont = 2^32 mod q = 4193792. Compare round-tripped coef to 2326 // (saved * mont) mod q after canonicalising both. 2327 let MONT: i64 = 4193792 2328 var diff: i64 = 0 2329 var j: i64 = 0 2330 while j < MLDSA_N { 2331 var a: i64 = _reduce32(_i32_load_le(poly, j)) 2332 if a < 0 { a = a + MLDSA_Q } 2333 var expected: i64 = (_i32_load_le(saved, j) * MONT) % MLDSA_Q 2334 if expected < 0 { expected = expected + MLDSA_Q } 2335 if a != expected { diff = 1 } 2336 j = j + 1 2337 } 2338 if diff != 0 { return 1 } 2339 return 0 2340}