luna - wolfssl/wolfcrypt/src/wc_mlkem

   1/* wc_mlkem_poly.c
   2 *
   3 * Copyright (C) 2006-2026 wolfSSL Inc.
   4 *
   5 * This file is part of wolfSSL.
   6 *
   7 * wolfSSL is free software; you can redistribute it and/or modify
   8 * it under the terms of the GNU General Public License as published by
   9 * the Free Software Foundation; either version 3 of the License, or
  10 * (at your option) any later version.
  11 *
  12 * wolfSSL is distributed in the hope that it will be useful,
  13 * but WITHOUT ANY WARRANTY; without even the implied warranty of
  14 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  15 * GNU General Public License for more details.
  16 *
  17 * You should have received a copy of the GNU General Public License
  18 * along with this program; if not, write to the Free Software
  19 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1335, USA
  20 */
  21
  22/* Implementation based on FIPS 203:
  23 *   https://nvlpubs.nist.gov/nistpubs/FIPS/NIST.FIPS.203.pdf
  24 *
  25 * Original implementation based on NIST 3rd Round submission package.
  26 * See link at:
  27 *   https://csrc.nist.gov/Projects/post-quantum-cryptography/
  28 *   post-quantum-cryptography-standardization/round-3-submissions
  29 */
  30
  31/* Implementation of the functions that operate on polynomials or vectors of
  32 * polynomials.
  33 */
  34
  35/* Possible ML-KEM options:
  36 *
  37 * WOLFSSL_HAVE_MLKEM                                         Default: OFF
  38 *   Enables this code, wolfSSL implementation, to be built.
  39 *
  40 * WOLFSSL_WC_ML_KEM_512                                      Default: OFF
  41 *   Enables the ML-KEM 512 parameter implementations.
  42 * WOLFSSL_WC_ML_KEM_768                                      Default: OFF
  43 *   Enables the ML-KEM 768 parameter implementations.
  44 * WOLFSSL_WC_ML_KEM_1024                                     Default: OFF
  45 *   Enables the ML-KEM 1024 parameter implementations.
  46 * WOLFSSL_KYBER512                                           Default: OFF
  47 *   Enables the KYBER512 parameter implementations.
  48 * WOLFSSL_KYBER768                                           Default: OFF
  49 *   Enables the KYBER768 parameter implementations.
  50 * WOLFSSL_KYBER1024                                          Default: OFF
  51 *   Enables the KYBER1024 parameter implementations.
  52 *
  53 * USE_INTEL_SPEEDUP                                          Default: OFF
  54 *   Compiles in Intel x64 specific implementations that are faster.
  55 * WOLFSSL_MLKEM_NO_LARGE_CODE                                Default: OFF
  56 *   Compiles smaller, fast code size with a speed trade-off.
  57 * WOLFSSL_MLKEM_SMALL                                        Default: OFF
  58 *   Compiles to small code size with a speed trade-off.
  59 * WOLFSSL_SMALL_STACK                                        Default: OFF
  60 *   Use less stack by dynamically allocating local variables.
  61 *
  62 * WOLFSSL_MLKEM_NTT_UNROLL                                   Default: OFF
  63 *   Enable an alternative NTT implementation that may be faster on some
  64 *   platforms and is smaller in code size.
  65 * WOLFSSL_MLKEM_INVNTT_UNROLL                                Default: OFF
  66 *   Enables an alternative inverse NTT implementation that may be faster on
  67 *   some platforms and is smaller in code size.
  68 */
  69
  70#include <wolfssl/wolfcrypt/libwolfssl_sources.h>
  71
  72#ifdef WC_MLKEM_NO_ASM
  73    #undef USE_INTEL_SPEEDUP
  74    #undef WOLFSSL_ARMASM
  75    #undef WOLFSSL_RISCV_ASM
  76#endif
  77
  78#include <wolfssl/wolfcrypt/wc_mlkem.h>
  79#include <wolfssl/wolfcrypt/sha3.h>
  80#include <wolfssl/wolfcrypt/cpuid.h>
  81
  82#ifdef WOLFSSL_HAVE_MLKEM
  83
  84#ifdef NO_INLINE
  85    #include <wolfssl/wolfcrypt/misc.h>
  86#else
  87    #define WOLFSSL_MISC_INCLUDED
  88    #include <wolfcrypt/src/misc.c>
  89#endif
  90
  91#if defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) || \
  92    defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM)
  93static int mlkem_gen_matrix_i(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed,
  94    int i, int transposed);
  95static int mlkem_get_noise_i(MLKEM_PRF_T* prf, int k, sword16* vec2,
  96    byte* seed, int i, int make);
  97static int mlkem_get_noise_eta2_c(MLKEM_PRF_T* prf, sword16* p,
  98    const byte* seed);
  99#endif
 100
 101/* Declared in wc_mlkem.c to stop compiler optimizer from simplifying. */
 102extern sword16 wc_mlkem_opt_blocker(void);
 103
 104#if defined(USE_INTEL_SPEEDUP) || (defined(__aarch64__) && \
 105    defined(WOLFSSL_ARMASM))
 106static cpuid_flags_t cpuid_flags = WC_CPUID_INITIALIZER;
 107#endif
 108
 109/* Half of Q plus one. Converted message bit value of 1. */
 110#define MLKEM_Q_1_HALF      ((MLKEM_Q + 1) / 2)
 111/* Half of Q */
 112#define MLKEM_Q_HALF        (MLKEM_Q / 2)
 113
 114
 115/* q^-1 mod 2^16 (inverse of 3329 mod 65536) */
 116#define MLKEM_QINV       62209
 117
 118/* Used in Barrett Reduction:
 119 *    r = a mod q
 120 * => r = a - ((V * a) >> 26) * q), as V based on 2^26
 121 * V is the multiplier that gets the quotient after shifting.
 122 */
 123#define MLKEM_V          (((1U << 26) + (MLKEM_Q / 2)) / MLKEM_Q)
 124
 125/* Used in converting to Montgomery form.
 126 * f is the normalizer = 2^k % m.
 127 * 16-bit value cast to sword32 in use.
 128 */
 129#define MLKEM_F          (((word64)1 << 32) % MLKEM_Q)
 130
 131/* Number of bytes in an output block of SHA-3-128 */
 132#define SHA3_128_BYTES   (WC_SHA3_128_COUNT * 8)
 133/* Number of bytes in an output block of SHA-3-256 */
 134#define SHA3_256_BYTES   (WC_SHA3_256_COUNT * 8)
 135
 136/* Number of blocks to generate for matrix. */
 137#define GEN_MATRIX_NBLOCKS \
 138    ((12 * MLKEM_N / 8 * (1 << 12) / MLKEM_Q + XOF_BLOCK_SIZE) / XOF_BLOCK_SIZE)
 139/* Number of bytes to generate for matrix. */
 140#define GEN_MATRIX_SIZE     GEN_MATRIX_NBLOCKS * XOF_BLOCK_SIZE
 141
 142
 143/* Number of random bytes to generate for ETA3. */
 144#define ETA3_RAND_SIZE     ((3 * MLKEM_N) / 4)
 145/* Number of random bytes to generate for ETA2. */
 146#define ETA2_RAND_SIZE     ((2 * MLKEM_N) / 4)
 147
 148
 149/* Montgomery reduce a.
 150 *
 151 * @param  [in]  a  32-bit value to be reduced.
 152 * @return  Montgomery reduction result.
 153 */
 154#define MLKEM_MONT_RED(a) \
 155    (sword16)(((a) - (sword32)(((sword16)((sword16)(a) * \
 156                                (sword16)MLKEM_QINV)) * \
 157                               (sword32)MLKEM_Q)) >> 16)
 158
 159/* Barrett reduce a. r = a mod q.
 160 *
 161 * Converted division to multiplication.
 162 *
 163 * @param  [in]  a  16-bit value to be reduced to range of q.
 164 * @return  Modulo result.
 165 */
 166#define MLKEM_BARRETT_RED(a) \
 167    (sword16)((sword16)(a) - (sword16)((sword16)( \
 168        ((sword32)((sword32)MLKEM_V * (sword16)(a))) >> 26) * (word16)MLKEM_Q))
 169
 170
 171/* Zetas for NTT. */
 172const sword16 zetas[MLKEM_N / 2] = {
 173    2285, 2571, 2970, 1812, 1493, 1422,  287,  202,
 174    3158,  622, 1577,  182,  962, 2127, 1855, 1468,
 175     573, 2004,  264,  383, 2500, 1458, 1727, 3199,
 176    2648, 1017,  732,  608, 1787,  411, 3124, 1758,
 177    1223,  652, 2777, 1015, 2036, 1491, 3047, 1785,
 178     516, 3321, 3009, 2663, 1711, 2167,  126, 1469,
 179    2476, 3239, 3058,  830,  107, 1908, 3082, 2378,
 180    2931,  961, 1821, 2604,  448, 2264,  677, 2054,
 181    2226,  430,  555,  843, 2078,  871, 1550,  105,
 182     422,  587,  177, 3094, 3038, 2869, 1574, 1653,
 183    3083,  778, 1159, 3182, 2552, 1483, 2727, 1119,
 184    1739,  644, 2457,  349,  418,  329, 3173, 3254,
 185     817, 1097,  603,  610, 1322, 2044, 1864,  384,
 186    2114, 3193, 1218, 1994, 2455,  220, 2142, 1670,
 187    2144, 1799, 2051,  794, 1819, 2475, 2459,  478,
 188    3221, 3021,  996,  991,  958, 1869, 1522, 1628
 189};
 190
 191
 192#if !defined(WOLFSSL_ARMASM)
 193/* Number-Theoretic Transform.
 194 *
 195 * FIPS 203, Algorithm 9: NTT(f)
 196 * Computes the NTT representation f_hat of the given polynomial f element of
 197 * R_q.
 198 *   1: f_hat <- f
 199 *   2: i <- 1
 200 *   3: for (len <- 128; len >= 2; len <- len/2)
 201 *   4:     for (start <- 0; start < 256; start <- start + 2.len)
 202 *   5:         zeta <- zetas^BitRev_7(i) mod q
 203 *   6:         i <- i + 1
 204 *   7:         for (j <- start; j < start + len; j++)
 205 *   8:             t <- zeta.f[j+len]
 206 *   9:             f_hat[j+len] <- f_hat[j] - t
 207 *  10:             f_hat[j] <- f_hat[j] + t
 208 *  11:         end for
 209 *  12:     end for
 210 *  13: end for
 211 *  14: return f_hat
 212 *
 213 * @param  [in, out]  r  Polynomial to transform.
 214 */
 215static void mlkem_ntt(sword16* r)
 216{
 217#ifdef WOLFSSL_MLKEM_SMALL
 218    unsigned int len;
 219    unsigned int k;
 220    unsigned int j;
 221
 222    /* Step 2 */
 223    k = 1;
 224    /* Step 3 */
 225    for (len = MLKEM_N / 2; len >= 2; len >>= 1) {
 226        unsigned int start;
 227        /* Step 4 */
 228        for (start = 0; start < MLKEM_N; start = j + len) {
 229            /* Step 5, 6*/
 230            sword16 zeta = zetas[k++];
 231            /* Step 7 */
 232            for (j = start; j < start + len; ++j) {
 233                /* Step 8 */
 234                sword32 p = (sword32)zeta * r[j + len];
 235                sword16 t = MLKEM_MONT_RED(p);
 236                sword16 rj = r[j];
 237                /* Step 9 */
 238                r[j + len] = (sword16)(rj - t);
 239                /* Step 10 */
 240                r[j] = (sword16)(rj + t);
 241            }
 242        }
 243    }
 244
 245    /* Reduce coefficients with quick algorithm. */
 246    for (j = 0; j < MLKEM_N; ++j) {
 247        r[j] = MLKEM_BARRETT_RED(r[j]);
 248    }
 249#elif defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
 250    /* Take out the first iteration. */
 251    unsigned int len;
 252    unsigned int k = 1;
 253    unsigned int j;
 254    unsigned int start;
 255    sword16 zeta = zetas[k++];
 256
 257    for (j = 0; j < MLKEM_N / 2; ++j) {
 258        sword32 p = (sword32)zeta * r[j + MLKEM_N / 2];
 259        sword16 t = MLKEM_MONT_RED(p);
 260        sword16 rj = r[j];
 261        r[j + MLKEM_N / 2] = (sword16)(rj - t);
 262        r[j] = (sword16)(rj + t);
 263    }
 264    for (len = MLKEM_N / 4; len >= 2; len >>= 1) {
 265        for (start = 0; start < MLKEM_N; start = j + len) {
 266            zeta = zetas[k++];
 267            for (j = start; j < start + len; ++j) {
 268                sword32 p = (sword32)zeta * r[j + len];
 269                sword16 t = MLKEM_MONT_RED(p);
 270                sword16 rj = r[j];
 271                r[j + len] = (sword16)(rj - t);
 272                r[j] = (sword16)(rj + t);
 273            }
 274        }
 275    }
 276
 277    /* Reduce coefficients with quick algorithm. */
 278    for (j = 0; j < MLKEM_N; ++j) {
 279        r[j] = MLKEM_BARRETT_RED(r[j]);
 280    }
 281#elif defined(WOLFSSL_MLKEM_NTT_UNROLL)
 282    /* Unroll len loop (Step 3). */
 283    unsigned int k = 1;
 284    unsigned int j;
 285    unsigned int start;
 286    sword16 zeta = zetas[k++];
 287
 288    /* len = 128 */
 289    for (j = 0; j < MLKEM_N / 2; ++j) {
 290        sword32 p = (sword32)zeta * r[j + MLKEM_N / 2];
 291        sword16 t = MLKEM_MONT_RED(p);
 292        sword16 rj = r[j];
 293        r[j + MLKEM_N / 2] = rj - t;
 294        r[j] = rj + t;
 295    }
 296    /* len = 64 */
 297    for (start = 0; start < MLKEM_N; start += 2 * 64) {
 298        zeta = zetas[k++];
 299        for (j = 0; j < 64; ++j) {
 300            sword32 p = (sword32)zeta * r[start + j + 64];
 301            sword16 t = MLKEM_MONT_RED(p);
 302            sword16 rj = r[start + j];
 303            r[start + j + 64] = rj - t;
 304            r[start + j] = rj + t;
 305        }
 306    }
 307    /* len = 32 */
 308    for (start = 0; start < MLKEM_N; start += 2 * 32) {
 309        zeta = zetas[k++];
 310        for (j = 0; j < 32; ++j) {
 311            sword32 p = (sword32)zeta * r[start + j + 32];
 312            sword16 t = MLKEM_MONT_RED(p);
 313            sword16 rj = r[start + j];
 314            r[start + j + 32] = rj - t;
 315            r[start + j] = rj + t;
 316        }
 317    }
 318    /* len = 16 */
 319    for (start = 0; start < MLKEM_N; start += 2 * 16) {
 320        zeta = zetas[k++];
 321        for (j = 0; j < 16; ++j) {
 322            sword32 p = (sword32)zeta * r[start + j + 16];
 323            sword16 t = MLKEM_MONT_RED(p);
 324            sword16 rj = r[start + j];
 325            r[start + j + 16] = rj - t;
 326            r[start + j] = rj + t;
 327        }
 328    }
 329    /* len = 8 */
 330    for (start = 0; start < MLKEM_N; start += 2 * 8) {
 331        zeta = zetas[k++];
 332        for (j = 0; j < 8; ++j) {
 333            sword32 p = (sword32)zeta * r[start + j + 8];
 334            sword16 t = MLKEM_MONT_RED(p);
 335            sword16 rj = r[start + j];
 336            r[start + j + 8] = rj - t;
 337            r[start + j] = rj + t;
 338        }
 339    }
 340    /* len = 4 */
 341    for (start = 0; start < MLKEM_N; start += 2 * 4) {
 342        zeta = zetas[k++];
 343        for (j = 0; j < 4; ++j) {
 344            sword32 p = (sword32)zeta * r[start + j + 4];
 345            sword16 t = MLKEM_MONT_RED(p);
 346            sword16 rj = r[start + j];
 347            r[start + j + 4] = rj - t;
 348            r[start + j] = rj + t;
 349        }
 350    }
 351    /* len = 2 */
 352    for (start = 0; start < MLKEM_N; start += 2 * 2) {
 353        zeta = zetas[k++];
 354        for (j = 0; j < 2; ++j) {
 355            sword32 p = (sword32)zeta * r[start + j + 2];
 356            sword16 t = MLKEM_MONT_RED(p);
 357            sword16 rj = r[start + j];
 358            r[start + j + 2] = rj - t;
 359            r[start + j] = rj + t;
 360        }
 361    }
 362    /* Reduce coefficients with quick algorithm. */
 363    for (j = 0; j < MLKEM_N; ++j) {
 364        r[j] = MLKEM_BARRETT_RED(r[j]);
 365    }
 366#else
 367    /* Unroll len (2, 3, 2) and start loops. */
 368    unsigned int j;
 369    sword16 t0;
 370    sword16 t1;
 371    sword16 t2;
 372    sword16 t3;
 373
 374    /* len = 128,64 */
 375    sword16 zeta128 = zetas[1];
 376    sword16 zeta64_0 = zetas[2];
 377    sword16 zeta64_1 = zetas[3];
 378    for (j = 0; j < MLKEM_N / 8; j++) {
 379        sword16 r0 = r[j +   0];
 380        sword16 r1 = r[j +  32];
 381        sword16 r2 = r[j +  64];
 382        sword16 r3 = r[j +  96];
 383        sword16 r4 = r[j + 128];
 384        sword16 r5 = r[j + 160];
 385        sword16 r6 = r[j + 192];
 386        sword16 r7 = r[j + 224];
 387
 388        t0 = MLKEM_MONT_RED((sword32)zeta128 * r4);
 389        t1 = MLKEM_MONT_RED((sword32)zeta128 * r5);
 390        t2 = MLKEM_MONT_RED((sword32)zeta128 * r6);
 391        t3 = MLKEM_MONT_RED((sword32)zeta128 * r7);
 392        r4 = (sword16)(r0 - t0);
 393        r5 = (sword16)(r1 - t1);
 394        r6 = (sword16)(r2 - t2);
 395        r7 = (sword16)(r3 - t3);
 396        r0 = (sword16)(r0 + t0);
 397        r1 = (sword16)(r1 + t1);
 398        r2 = (sword16)(r2 + t2);
 399        r3 = (sword16)(r3 + t3);
 400
 401        t0 = MLKEM_MONT_RED((sword32)zeta64_0 * r2);
 402        t1 = MLKEM_MONT_RED((sword32)zeta64_0 * r3);
 403        t2 = MLKEM_MONT_RED((sword32)zeta64_1 * r6);
 404        t3 = MLKEM_MONT_RED((sword32)zeta64_1 * r7);
 405        r2 = (sword16)(r0 - t0);
 406        r3 = (sword16)(r1 - t1);
 407        r6 = (sword16)(r4 - t2);
 408        r7 = (sword16)(r5 - t3);
 409        r0 = (sword16)(r0 + t0);
 410        r1 = (sword16)(r1 + t1);
 411        r4 = (sword16)(r4 + t2);
 412        r5 = (sword16)(r5 + t3);
 413
 414        r[j +   0] = r0;
 415        r[j +  32] = r1;
 416        r[j +  64] = r2;
 417        r[j +  96] = r3;
 418        r[j + 128] = r4;
 419        r[j + 160] = r5;
 420        r[j + 192] = r6;
 421        r[j + 224] = r7;
 422    }
 423
 424    /* len = 32,16,8 */
 425    for (j = 0; j < MLKEM_N; j += 64) {
 426        unsigned int i;
 427        sword16 zeta32   = zetas[ 4 + j / 64 + 0];
 428        sword16 zeta16_0 = zetas[ 8 + j / 32 + 0];
 429        sword16 zeta16_1 = zetas[ 8 + j / 32 + 1];
 430        sword16 zeta8_0  = zetas[16 + j / 16 + 0];
 431        sword16 zeta8_1  = zetas[16 + j / 16 + 1];
 432        sword16 zeta8_2  = zetas[16 + j / 16 + 2];
 433        sword16 zeta8_3  = zetas[16 + j / 16 + 3];
 434        for (i = 0; i < 8; i++) {
 435            sword16 r0 = r[j + i +  0];
 436            sword16 r1 = r[j + i +  8];
 437            sword16 r2 = r[j + i + 16];
 438            sword16 r3 = r[j + i + 24];
 439            sword16 r4 = r[j + i + 32];
 440            sword16 r5 = r[j + i + 40];
 441            sword16 r6 = r[j + i + 48];
 442            sword16 r7 = r[j + i + 56];
 443
 444            t0 = MLKEM_MONT_RED((sword32)zeta32 * r4);
 445            t1 = MLKEM_MONT_RED((sword32)zeta32 * r5);
 446            t2 = MLKEM_MONT_RED((sword32)zeta32 * r6);
 447            t3 = MLKEM_MONT_RED((sword32)zeta32 * r7);
 448            r4 = (sword16)(r0 - t0);
 449            r5 = (sword16)(r1 - t1);
 450            r6 = (sword16)(r2 - t2);
 451            r7 = (sword16)(r3 - t3);
 452            r0 = (sword16)(r0 + t0);
 453            r1 = (sword16)(r1 + t1);
 454            r2 = (sword16)(r2 + t2);
 455            r3 = (sword16)(r3 + t3);
 456
 457            t0 = MLKEM_MONT_RED((sword32)zeta16_0 * r2);
 458            t1 = MLKEM_MONT_RED((sword32)zeta16_0 * r3);
 459            t2 = MLKEM_MONT_RED((sword32)zeta16_1 * r6);
 460            t3 = MLKEM_MONT_RED((sword32)zeta16_1 * r7);
 461            r2 = (sword16)(r0 - t0);
 462            r3 = (sword16)(r1 - t1);
 463            r6 = (sword16)(r4 - t2);
 464            r7 = (sword16)(r5 - t3);
 465            r0 = (sword16)(r0 + t0);
 466            r1 = (sword16)(r1 + t1);
 467            r4 = (sword16)(r4 + t2);
 468            r5 = (sword16)(r5 + t3);
 469
 470            t0 = MLKEM_MONT_RED((sword32)zeta8_0 * r1);
 471            t1 = MLKEM_MONT_RED((sword32)zeta8_1 * r3);
 472            t2 = MLKEM_MONT_RED((sword32)zeta8_2 * r5);
 473            t3 = MLKEM_MONT_RED((sword32)zeta8_3 * r7);
 474            r1 = (sword16)(r0 - t0);
 475            r3 = (sword16)(r2 - t1);
 476            r5 = (sword16)(r4 - t2);
 477            r7 = (sword16)(r6 - t3);
 478            r0 = (sword16)(r0 + t0);
 479            r2 = (sword16)(r2 + t1);
 480            r4 = (sword16)(r4 + t2);
 481            r6 = (sword16)(r6 + t3);
 482
 483            r[j + i +  0] = r0;
 484            r[j + i +  8] = r1;
 485            r[j + i + 16] = r2;
 486            r[j + i + 24] = r3;
 487            r[j + i + 32] = r4;
 488            r[j + i + 40] = r5;
 489            r[j + i + 48] = r6;
 490            r[j + i + 56] = r7;
 491        }
 492    }
 493
 494    /* len = 4,2 and Final reduction */
 495    for (j = 0; j < MLKEM_N; j += 8) {
 496        sword16 zeta4  = zetas[32 + j / 8 + 0];
 497        sword16 zeta2_0 = zetas[64 + j / 4 + 0];
 498        sword16 zeta2_1 = zetas[64 + j / 4 + 1];
 499        sword16 r0 = r[j + 0];
 500        sword16 r1 = r[j + 1];
 501        sword16 r2 = r[j + 2];
 502        sword16 r3 = r[j + 3];
 503        sword16 r4 = r[j + 4];
 504        sword16 r5 = r[j + 5];
 505        sword16 r6 = r[j + 6];
 506        sword16 r7 = r[j + 7];
 507
 508        t0 = MLKEM_MONT_RED((sword32)zeta4 * r4);
 509        t1 = MLKEM_MONT_RED((sword32)zeta4 * r5);
 510        t2 = MLKEM_MONT_RED((sword32)zeta4 * r6);
 511        t3 = MLKEM_MONT_RED((sword32)zeta4 * r7);
 512        r4 = (sword16)(r0 - t0);
 513        r5 = (sword16)(r1 - t1);
 514        r6 = (sword16)(r2 - t2);
 515        r7 = (sword16)(r3 - t3);
 516        r0 = (sword16)(r0 + t0);
 517        r1 = (sword16)(r1 + t1);
 518        r2 = (sword16)(r2 + t2);
 519        r3 = (sword16)(r3 + t3);
 520
 521        t0 = MLKEM_MONT_RED((sword32)zeta2_0 * r2);
 522        t1 = MLKEM_MONT_RED((sword32)zeta2_0 * r3);
 523        t2 = MLKEM_MONT_RED((sword32)zeta2_1 * r6);
 524        t3 = MLKEM_MONT_RED((sword32)zeta2_1 * r7);
 525        r2 = (sword16)(r0 - t0);
 526        r3 = (sword16)(r1 - t1);
 527        r6 = (sword16)(r4 - t2);
 528        r7 = (sword16)(r5 - t3);
 529        r0 = (sword16)(r0 + t0);
 530        r1 = (sword16)(r1 + t1);
 531        r4 = (sword16)(r4 + t2);
 532        r5 = (sword16)(r5 + t3);
 533
 534        r[j + 0] = MLKEM_BARRETT_RED(r0);
 535        r[j + 1] = MLKEM_BARRETT_RED(r1);
 536        r[j + 2] = MLKEM_BARRETT_RED(r2);
 537        r[j + 3] = MLKEM_BARRETT_RED(r3);
 538        r[j + 4] = MLKEM_BARRETT_RED(r4);
 539        r[j + 5] = MLKEM_BARRETT_RED(r5);
 540        r[j + 6] = MLKEM_BARRETT_RED(r6);
 541        r[j + 7] = MLKEM_BARRETT_RED(r7);
 542    }
 543#endif
 544}
 545
 546#if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
 547    !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
 548/* Zetas for inverse NTT. */
 549const sword16 zetas_inv[MLKEM_N / 2] = {
 550    1701, 1807, 1460, 2371, 2338, 2333,  308,  108,
 551    2851,  870,  854, 1510, 2535, 1278, 1530, 1185,
 552    1659, 1187, 3109,  874, 1335, 2111,  136, 1215,
 553    2945, 1465, 1285, 2007, 2719, 2726, 2232, 2512,
 554      75,  156, 3000, 2911, 2980,  872, 2685, 1590,
 555    2210,  602, 1846,  777,  147, 2170, 2551,  246,
 556    1676, 1755,  460,  291,  235, 3152, 2742, 2907,
 557    3224, 1779, 2458, 1251, 2486, 2774, 2899, 1103,
 558    1275, 2652, 1065, 2881,  725, 1508, 2368,  398,
 559     951,  247, 1421, 3222, 2499,  271,   90,  853,
 560    1860, 3203, 1162, 1618,  666,  320,    8, 2813,
 561    1544,  282, 1838, 1293, 2314,  552, 2677, 2106,
 562    1571,  205, 2918, 1542, 2721, 2597, 2312,  681,
 563     130, 1602, 1871,  829, 2946, 3065, 1325, 2756,
 564    1861, 1474, 1202, 2367, 3147, 1752, 2707,  171,
 565    3127, 3042, 1907, 1836, 1517,  359,  758, 1441
 566};
 567
 568/* Inverse Number-Theoretic Transform.
 569 *
 570 * FIPS 203, Algorithm 10: NTT^-1(f_hat)
 571 * Computes the polynomial f element of R_q that corresponds to the given NTT
 572 * representation f element of T_q.
 573 *   1: f <- f_hat
 574 *   2: i <- 127
 575 *   3: for (len <- 2; len <= 128 ; len <- 2.len)
 576 *   4:     for (start <- 0; start < 256; start <- start + 2.len)
 577 *   5:         zeta <- zetas^BitRev_7(i) mod q
 578 *   6:         i <- i - 1
 579 *   7:         for (j <- start; j < start + len; j++)
 580 *   8:             t <- f[j]
 581 *   9:             f[j] <- t + f[j + len]
 582 *  10:             f[j + len] <- zeta.(f[j+len] - t)
 583 *  11:         end for
 584 *  12:     end for
 585 *  13: end for
 586 *  14: f <- f.3303 mod q
 587 *  15: return f
 588 *
 589 * @param  [in, out]  r  Polynomial to transform.
 590 */
 591static void mlkem_invntt(sword16* r)
 592{
 593#ifdef WOLFSSL_MLKEM_SMALL
 594    unsigned int len;
 595    unsigned int k;
 596    unsigned int j;
 597    sword16 zeta;
 598
 599    /* Step 2 - table reversed */
 600    k = 0;
 601    /* Step 3 */
 602    for (len = 2; len <= MLKEM_N / 2; len <<= 1) {
 603        unsigned int start;
 604        /* Step 4 */
 605        for (start = 0; start < MLKEM_N; start = j + len) {
 606            /* Step 5, 6 */
 607            zeta = zetas_inv[k++];
 608            /* Step 7 */
 609            for (j = start; j < start + len; ++j) {
 610                sword32 p;
 611                /* Step 8 */
 612                sword16 rj = r[j];
 613                sword16 rjl = r[j + len];
 614                /* Step 9 */
 615                sword16 t = (sword16)(rj + rjl);
 616                r[j] = MLKEM_BARRETT_RED(t);
 617                /* Step 10 */
 618                rjl = (sword16)(rj - rjl);
 619                p = (sword32)zeta * rjl;
 620                r[j + len] = MLKEM_MONT_RED(p);
 621            }
 622        }
 623    }
 624
 625    /* Step 14 */
 626    zeta = zetas_inv[127];
 627    for (j = 0; j < MLKEM_N; ++j) {
 628        sword32 p = (sword32)zeta * r[j];
 629        r[j] = MLKEM_MONT_RED(p);
 630    }
 631#elif defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
 632    /* Take out last iteration. */
 633    unsigned int len;
 634    unsigned int k;
 635    unsigned int j;
 636    sword16 zeta;
 637    sword16 zeta2;
 638
 639    k = 0;
 640    for (len = 2; len <= MLKEM_N / 4; len <<= 1) {
 641        unsigned int start;
 642        for (start = 0; start < MLKEM_N; start = j + len) {
 643            zeta = zetas_inv[k++];
 644            for (j = start; j < start + len; ++j) {
 645                sword32 p;
 646                sword16 rj = r[j];
 647                sword16 rjl = r[j + len];
 648                sword16 t = (sword16)(rj + rjl);
 649                r[j] = MLKEM_BARRETT_RED(t);
 650                rjl = (sword16)(rj - rjl);
 651                p = (sword32)zeta * rjl;
 652                r[j + len] = MLKEM_MONT_RED(p);
 653            }
 654        }
 655    }
 656
 657    zeta = zetas_inv[126];
 658    zeta2 = zetas_inv[127];
 659    for (j = 0; j < MLKEM_N / 2; ++j) {
 660        sword32 p;
 661        sword16 rj = r[j];
 662        sword16 rjl = r[j + MLKEM_N / 2];
 663        sword16 t = (sword16)(rj + rjl);
 664        rjl = (sword16)(rj - rjl);
 665        p = (sword32)zeta * rjl;
 666        r[j] = (sword16)t;
 667        r[j + MLKEM_N / 2] = MLKEM_MONT_RED(p);
 668
 669        p = (sword32)zeta2 * r[j];
 670        r[j] = MLKEM_MONT_RED(p);
 671        p = (sword32)zeta2 * r[j + MLKEM_N / 2];
 672        r[j + MLKEM_N / 2] = MLKEM_MONT_RED(p);
 673    }
 674#elif defined(WOLFSSL_MLKEM_INVNTT_UNROLL)
 675    /* Unroll len loop (Step 3). */
 676    unsigned int k;
 677    unsigned int j;
 678    unsigned int start;
 679    sword16 zeta;
 680    sword16 zeta2;
 681
 682    k = 0;
 683    /* len = 2 */
 684    for (start = 0; start < MLKEM_N; start += 2 * 2) {
 685        zeta = zetas_inv[k++];
 686        for (j = 0; j < 2; ++j) {
 687            sword32 p;
 688            sword16 rj = r[start + j];
 689            sword16 rjl = r[start + j + 2];
 690            sword16 t = rj + rjl;
 691            r[start + j] = t;
 692            rjl = rj - rjl;
 693            p = (sword32)zeta * rjl;
 694            r[start + j + 2] = MLKEM_MONT_RED(p);
 695        }
 696    }
 697    /* len = 4 */
 698    for (start = 0; start < MLKEM_N; start += 2 * 4) {
 699        zeta = zetas_inv[k++];
 700        for (j = 0; j < 4; ++j) {
 701            sword32 p;
 702            sword16 rj = r[start + j];
 703            sword16 rjl = r[start + j + 4];
 704            sword16 t = rj + rjl;
 705            r[start + j] = t;
 706            rjl = rj - rjl;
 707            p = (sword32)zeta * rjl;
 708            r[start + j + 4] = MLKEM_MONT_RED(p);
 709        }
 710    }
 711    /* len = 8 */
 712    for (start = 0; start < MLKEM_N; start += 2 * 8) {
 713        zeta = zetas_inv[k++];
 714        for (j = 0; j < 8; ++j) {
 715            sword32 p;
 716            sword16 rj = r[start + j];
 717            sword16 rjl = r[start + j + 8];
 718            sword16 t = rj + rjl;
 719            /* Reduce. */
 720            r[start + j] = MLKEM_BARRETT_RED(t);
 721            rjl = rj - rjl;
 722            p = (sword32)zeta * rjl;
 723            r[start + j + 8] = MLKEM_MONT_RED(p);
 724        }
 725    }
 726    /* len = 16 */
 727    for (start = 0; start < MLKEM_N; start += 2 * 16) {
 728        zeta = zetas_inv[k++];
 729        for (j = 0; j < 16; ++j) {
 730            sword32 p;
 731            sword16 rj = r[start + j];
 732            sword16 rjl = r[start + j + 16];
 733            sword16 t = rj + rjl;
 734            r[start + j] = t;
 735            rjl = rj - rjl;
 736            p = (sword32)zeta * rjl;
 737            r[start + j + 16] = MLKEM_MONT_RED(p);
 738        }
 739    }
 740    /* len = 32 */
 741    for (start = 0; start < MLKEM_N; start += 2 * 32) {
 742        zeta = zetas_inv[k++];
 743        for (j = 0; j < 32; ++j) {
 744            sword32 p;
 745            sword16 rj = r[start + j];
 746            sword16 rjl = r[start + j + 32];
 747            sword16 t = rj + rjl;
 748            r[start + j] = t;
 749            rjl = rj - rjl;
 750            p = (sword32)zeta * rjl;
 751            r[start + j + 32] = MLKEM_MONT_RED(p);
 752        }
 753    }
 754    /* len = 64 */
 755    for (start = 0; start < MLKEM_N; start += 2 * 64) {
 756        zeta = zetas_inv[k++];
 757        for (j = 0; j < 64; ++j) {
 758            sword32 p;
 759            sword16 rj = r[start + j];
 760            sword16 rjl = r[start + j + 64];
 761            sword16 t = rj + rjl;
 762            /* Reduce. */
 763            r[start + j] = MLKEM_BARRETT_RED(t);
 764            rjl = rj - rjl;
 765            p = (sword32)zeta * rjl;
 766            r[start + j + 64] = MLKEM_MONT_RED(p);
 767        }
 768    }
 769    /* len = 128, 256 */
 770    zeta = zetas_inv[126];
 771    zeta2 = zetas_inv[127];
 772    for (j = 0; j < MLKEM_N / 2; ++j) {
 773        sword32 p;
 774        sword16 rj = r[j];
 775        sword16 rjl = r[j + MLKEM_N / 2];
 776        sword16 t = rj + rjl;
 777        rjl = rj - rjl;
 778        p = (sword32)zeta * rjl;
 779        r[j] = t;
 780        r[j + MLKEM_N / 2] = MLKEM_MONT_RED(p);
 781
 782        p = (sword32)zeta2 * r[j];
 783        r[j] = MLKEM_MONT_RED(p);
 784        p = (sword32)zeta2 * r[j + MLKEM_N / 2];
 785        r[j + MLKEM_N / 2] = MLKEM_MONT_RED(p);
 786    }
 787#else
 788    /* Unroll len (2, 3, 3) and start loops. */
 789    unsigned int j;
 790    sword16 t0;
 791    sword16 t1;
 792    sword16 t2;
 793    sword16 t3;
 794    sword16 zeta64_0;
 795    sword16 zeta64_1;
 796    sword16 zeta128;
 797    sword16 zeta256;
 798    sword32 p;
 799
 800    for (j = 0; j < MLKEM_N; j += 8) {
 801        sword16 zeta2_0 = zetas_inv[ 0 + j / 4 + 0];
 802        sword16 zeta2_1 = zetas_inv[ 0 + j / 4 + 1];
 803        sword16 zeta4   = zetas_inv[64 + j / 8 + 0];
 804        sword16 r0 = r[j + 0];
 805        sword16 r1 = r[j + 1];
 806        sword16 r2 = r[j + 2];
 807        sword16 r3 = r[j + 3];
 808        sword16 r4 = r[j + 4];
 809        sword16 r5 = r[j + 5];
 810        sword16 r6 = r[j + 6];
 811        sword16 r7 = r[j + 7];
 812
 813        p = (sword32)zeta2_0 * (sword16)(r0 - r2);
 814        t0 = MLKEM_MONT_RED(p);
 815        p = (sword32)zeta2_0 * (sword16)(r1 - r3);
 816        t1 = MLKEM_MONT_RED(p);
 817        p = (sword32)zeta2_1 * (sword16)(r4 - r6);
 818        t2 = MLKEM_MONT_RED(p);
 819        p = (sword32)zeta2_1 * (sword16)(r5 - r7);
 820        t3 = MLKEM_MONT_RED(p);
 821        r0 = (sword16)(r0 + r2);
 822        r1 = (sword16)(r1 + r3);
 823        r4 = (sword16)(r4 + r6);
 824        r5 = (sword16)(r5 + r7);
 825        r2 = t0;
 826        r3 = t1;
 827        r6 = t2;
 828        r7 = t3;
 829
 830        p = (sword32)zeta4 * (sword16)(r0 - r4);
 831        t0 = MLKEM_MONT_RED(p);
 832        p = (sword32)zeta4 * (sword16)(r1 - r5);
 833        t1 = MLKEM_MONT_RED(p);
 834        p = (sword32)zeta4 * (sword16)(r2 - r6);
 835        t2 = MLKEM_MONT_RED(p);
 836        p = (sword32)zeta4 * (sword16)(r3 - r7);
 837        t3 = MLKEM_MONT_RED(p);
 838        r0 = (sword16)(r0 + r4);
 839        r1 = (sword16)(r1 + r5);
 840        r2 = (sword16)(r2 + r6);
 841        r3 = (sword16)(r3 + r7);
 842        r4 = t0;
 843        r5 = t1;
 844        r6 = t2;
 845        r7 = t3;
 846
 847        r[j + 0] = r0;
 848        r[j + 1] = r1;
 849        r[j + 2] = r2;
 850        r[j + 3] = r3;
 851        r[j + 4] = r4;
 852        r[j + 5] = r5;
 853        r[j + 6] = r6;
 854        r[j + 7] = r7;
 855    }
 856
 857    for (j = 0; j < MLKEM_N; j += 64) {
 858        unsigned int i;
 859        sword16 zeta8_0  = zetas_inv[ 96 + j / 16 + 0];
 860        sword16 zeta8_1  = zetas_inv[ 96 + j / 16 + 1];
 861        sword16 zeta8_2  = zetas_inv[ 96 + j / 16 + 2];
 862        sword16 zeta8_3  = zetas_inv[ 96 + j / 16 + 3];
 863        sword16 zeta16_0 = zetas_inv[112 + j / 32 + 0];
 864        sword16 zeta16_1 = zetas_inv[112 + j / 32 + 1];
 865        sword16 zeta32   = zetas_inv[120 + j / 64 + 0];
 866        for (i = 0; i < 8; i++) {
 867            sword16 r0 = r[j + i +  0];
 868            sword16 r1 = r[j + i +  8];
 869            sword16 r2 = r[j + i + 16];
 870            sword16 r3 = r[j + i + 24];
 871            sword16 r4 = r[j + i + 32];
 872            sword16 r5 = r[j + i + 40];
 873            sword16 r6 = r[j + i + 48];
 874            sword16 r7 = r[j + i + 56];
 875
 876            p = (sword32)zeta8_0 * (sword16)(r0 - r1);
 877            t0 = MLKEM_MONT_RED(p);
 878            p = (sword32)zeta8_1 * (sword16)(r2 - r3);
 879            t1 = MLKEM_MONT_RED(p);
 880            p = (sword32)zeta8_2 * (sword16)(r4 - r5);
 881            t2 = MLKEM_MONT_RED(p);
 882            p = (sword32)zeta8_3 * (sword16)(r6 - r7);
 883            t3 = MLKEM_MONT_RED(p);
 884            r0 = MLKEM_BARRETT_RED(r0 + r1);
 885            r2 = MLKEM_BARRETT_RED(r2 + r3);
 886            r4 = MLKEM_BARRETT_RED(r4 + r5);
 887            r6 = MLKEM_BARRETT_RED(r6 + r7);
 888            r1 = t0;
 889            r3 = t1;
 890            r5 = t2;
 891            r7 = t3;
 892
 893            p = (sword32)zeta16_0 * (sword16)(r0 - r2);
 894            t0 = MLKEM_MONT_RED(p);
 895            p = (sword32)zeta16_0 * (sword16)(r1 - r3);
 896            t1 = MLKEM_MONT_RED(p);
 897            p = (sword32)zeta16_1 * (sword16)(r4 - r6);
 898            t2 = MLKEM_MONT_RED(p);
 899            p = (sword32)zeta16_1 * (sword16)(r5 - r7);
 900            t3 = MLKEM_MONT_RED(p);
 901            r0 = (sword16)(r0 + r2);
 902            r1 = (sword16)(r1 + r3);
 903            r4 = (sword16)(r4 + r6);
 904            r5 = (sword16)(r5 + r7);
 905            r2 = t0;
 906            r3 = t1;
 907            r6 = t2;
 908            r7 = t3;
 909
 910            p = (sword32)zeta32 * (sword16)(r0 - r4);
 911            t0 = MLKEM_MONT_RED(p);
 912            p = (sword32)zeta32 * (sword16)(r1 - r5);
 913            t1 = MLKEM_MONT_RED(p);
 914            p = (sword32)zeta32 * (sword16)(r2 - r6);
 915            t2 = MLKEM_MONT_RED(p);
 916            p = (sword32)zeta32 * (sword16)(r3 - r7);
 917            t3 = MLKEM_MONT_RED(p);
 918            r0 = (sword16)(r0 + r4);
 919            r1 = (sword16)(r1 + r5);
 920            r2 = (sword16)(r2 + r6);
 921            r3 = (sword16)(r3 + r7);
 922            r4 = t0;
 923            r5 = t1;
 924            r6 = t2;
 925            r7 = t3;
 926
 927            r[j + i +  0] = r0;
 928            r[j + i +  8] = r1;
 929            r[j + i + 16] = r2;
 930            r[j + i + 24] = r3;
 931            r[j + i + 32] = r4;
 932            r[j + i + 40] = r5;
 933            r[j + i + 48] = r6;
 934            r[j + i + 56] = r7;
 935        }
 936    }
 937
 938    zeta64_0 = zetas_inv[124];
 939    zeta64_1 = zetas_inv[125];
 940    zeta128  = zetas_inv[126];
 941    zeta256  = zetas_inv[127];
 942    for (j = 0; j < MLKEM_N / 8; j++) {
 943        sword16 r0 = r[j +   0];
 944        sword16 r1 = r[j +  32];
 945        sword16 r2 = r[j +  64];
 946        sword16 r3 = r[j +  96];
 947        sword16 r4 = r[j + 128];
 948        sword16 r5 = r[j + 160];
 949        sword16 r6 = r[j + 192];
 950        sword16 r7 = r[j + 224];
 951
 952        p = (sword32)zeta64_0 * (sword16)(r0 - r2);
 953        t0 = MLKEM_MONT_RED(p);
 954        p = (sword32)zeta64_0 * (sword16)(r1 - r3);
 955        t1 = MLKEM_MONT_RED(p);
 956        p = (sword32)zeta64_1 * (sword16)(r4 - r6);
 957        t2 = MLKEM_MONT_RED(p);
 958        p = (sword32)zeta64_1 * (sword16)(r5 - r7);
 959        t3 = MLKEM_MONT_RED(p);
 960        r0 = MLKEM_BARRETT_RED(r0 + r2);
 961        r1 = MLKEM_BARRETT_RED(r1 + r3);
 962        r4 = MLKEM_BARRETT_RED(r4 + r6);
 963        r5 = MLKEM_BARRETT_RED(r5 + r7);
 964        r2 = t0;
 965        r3 = t1;
 966        r6 = t2;
 967        r7 = t3;
 968
 969        p = (sword32)zeta128 * (sword16)(r0 - r4);
 970        t0 = MLKEM_MONT_RED(p);
 971        p = (sword32)zeta128 * (sword16)(r1 - r5);
 972        t1 = MLKEM_MONT_RED(p);
 973        p = (sword32)zeta128 * (sword16)(r2 - r6);
 974        t2 = MLKEM_MONT_RED(p);
 975        p = (sword32)zeta128 * (sword16)(r3 - r7);
 976        t3 = MLKEM_MONT_RED(p);
 977        r0 = (sword16)(r0 + r4);
 978        r1 = (sword16)(r1 + r5);
 979        r2 = (sword16)(r2 + r6);
 980        r3 = (sword16)(r3 + r7);
 981        r4 = t0;
 982        r5 = t1;
 983        r6 = t2;
 984        r7 = t3;
 985
 986        p = (sword32)zeta256 * r0;
 987        r0 = MLKEM_MONT_RED(p);
 988        p = (sword32)zeta256 * r1;
 989        r1 = MLKEM_MONT_RED(p);
 990        p = (sword32)zeta256 * r2;
 991        r2 = MLKEM_MONT_RED(p);
 992        p = (sword32)zeta256 * r3;
 993        r3 = MLKEM_MONT_RED(p);
 994        p = (sword32)zeta256 * r4;
 995        r4 = MLKEM_MONT_RED(p);
 996        p = (sword32)zeta256 * r5;
 997        r5 = MLKEM_MONT_RED(p);
 998        p = (sword32)zeta256 * r6;
 999        r6 = MLKEM_MONT_RED(p);
1000        p = (sword32)zeta256 * r7;
1001        r7 = MLKEM_MONT_RED(p);
1002
1003        r[j +   0] = r0;
1004        r[j +  32] = r1;
1005        r[j +  64] = r2;
1006        r[j +  96] = r3;
1007        r[j + 128] = r4;
1008        r[j + 160] = r5;
1009        r[j + 192] = r6;
1010        r[j + 224] = r7;
1011    }
1012#endif
1013}
1014#endif
1015
1016/* Multiplication of polynomials in Zq[X]/(X^2-zeta).
1017 *
1018 * Used for multiplication of elements in Rq in NTT domain.
1019 *
1020 * FIPS 203, Algorithm 12: BaseCaseMultiply(a0, a1, b0, b1, zeta)
1021 * Computes the product of two degree-one polynomials with respect to a
1022 * quadratic modulus.
1023 *   1: c0 <- a0.b0 + a1.b1.zeta
1024 *   2: c1 <- a0.b1 + a1.b0
1025 *   3: return (c0, c1)
1026 *
1027 * @param  [out]  r     Result polynomial.
1028 * @param  [in]   a     First factor.
1029 * @param  [in]   b     Second factor.
1030 * @param  [in]   zeta  Integer defining the reduction polynomial.
1031 */
1032static void mlkem_basemul(sword16* r, const sword16* a, const sword16* b,
1033    sword16 zeta)
1034{
1035    sword16 r0;
1036    sword16 a0 = a[0];
1037    sword16 a1 = a[1];
1038    sword16 b0 = b[0];
1039    sword16 b1 = b[1];
1040    sword32 p1;
1041    sword32 p2;
1042
1043    /* Step 1 */
1044    p1   = (sword32)a0 * b0;
1045    p2   = (sword32)a1 * b1;
1046    r0   = MLKEM_MONT_RED(p2);
1047    p2   = (sword32)zeta * r0;
1048    p2  += p1;
1049    r[0] = MLKEM_MONT_RED(p2);
1050
1051    /* Step 2 */
1052    p1   = (sword32)a0 * b1;
1053    p2   = (sword32)a1 * b0;
1054    p1  += p2;
1055    r[1] = MLKEM_MONT_RED(p1);
1056}
1057
1058/* Multiply two polynomials in NTT domain. r = a * b.
1059 *
1060 * FIPS 203, Algorithm 11: MultiplyNTTs(f_hat, g_hat)
1061 * Computes the product (in the ring T_q) of two NTT representations.
1062 *   1: for (i <- 0; i < 128; i++)
1063 *   2:     (h_hat[2i],h_hat[2i+1]) <-
1064 *              BaseCaseMultiply(f_hat[2i],f_hat[2i+1],g_hat[2i],g_hat[2i+1],
1065 *                               zetas^(BitRev_7(i)+1))
1066 *   3: end for
1067 *   4: return h_hat
1068 *
1069 * @param  [out]  r  Result polynomial.
1070 * @param  [in]   a  First polynomial multiplier.
1071 * @param  [in]   b  Second polynomial multiplier.
1072 */
1073static void mlkem_basemul_mont(sword16* r, const sword16* a, const sword16* b)
1074{
1075    const sword16* zeta = zetas + 64;
1076
1077#if defined(WOLFSSL_MLKEM_SMALL)
1078    /* Two multiplications per loop. */
1079    unsigned int i;
1080    /* Step 1 */
1081    for (i = 0; i < MLKEM_N; i += 4, zeta++) {
1082        /* Step 2 */
1083        mlkem_basemul(r + i + 0, a + i + 0, b + i + 0, zeta[0]);
1084        mlkem_basemul(r + i + 2, a + i + 2, b + i + 2, (sword16)(-zeta[0]));
1085    }
1086#elif defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
1087    /* Four multiplications per loop. */
1088    unsigned int i;
1089    for (i = 0; i < MLKEM_N; i += 8, zeta += 2) {
1090        mlkem_basemul(r + i + 0, a + i + 0, b + i + 0, zeta[0]);
1091        mlkem_basemul(r + i + 2, a + i + 2, b + i + 2, (sword16)(-zeta[0]));
1092        mlkem_basemul(r + i + 4, a + i + 4, b + i + 4, zeta[1]);
1093        mlkem_basemul(r + i + 6, a + i + 6, b + i + 6, (sword16)(-zeta[1]));
1094    }
1095#else
1096    /* Eight multiplications per loop. */
1097    unsigned int i;
1098    for (i = 0; i < MLKEM_N; i += 16, zeta += 4) {
1099        mlkem_basemul(r + i +  0, a + i +  0, b + i +  0, zeta[0]);
1100        mlkem_basemul(r + i +  2, a + i +  2, b + i +  2, (sword16)(-zeta[0]));
1101        mlkem_basemul(r + i +  4, a + i +  4, b + i +  4, zeta[1]);
1102        mlkem_basemul(r + i +  6, a + i +  6, b + i +  6, (sword16)(-zeta[1]));
1103        mlkem_basemul(r + i +  8, a + i +  8, b + i +  8, zeta[2]);
1104        mlkem_basemul(r + i + 10, a + i + 10, b + i + 10, (sword16)(-zeta[2]));
1105        mlkem_basemul(r + i + 12, a + i + 12, b + i + 12, zeta[3]);
1106        mlkem_basemul(r + i + 14, a + i + 14, b + i + 14, (sword16)(-zeta[3]));
1107    }
1108#endif
1109}
1110
1111/* Multiply two polynomials in NTT domain and add to result. r += a * b.
1112 *
1113 * FIPS 203, Algorithm 11: MultiplyNTTs(f_hat, g_hat)
1114 * Computes the product (in the ring T_q) of two NTT representations.
1115 *   1: for (i <- 0; i < 128; i++)
1116 *   2:     (h_hat[2i],h_hat[2i+1]) <-
1117 *              BaseCaseMultiply(f_hat[2i],f_hat[2i+1],g_hat[2i],g_hat[2i+1],
1118 *                               zetas^(BitRev_7(i)+1))
1119 *   3: end for
1120 *   4: return h_hat
1121 * Add h_hat to r.
1122 *
1123 * @param  [in, out]  r  Result polynomial.
1124 * @param  [in]       a  First polynomial multiplier.
1125 * @param  [in]       b  Second polynomial multiplier.
1126 */
1127static void mlkem_basemul_mont_add(sword16* r, const sword16* a,
1128    const sword16* b)
1129{
1130    const sword16* zeta = zetas + 64;
1131
1132#if defined(WOLFSSL_MLKEM_SMALL)
1133    /* Two multiplications per loop. */
1134    unsigned int i;
1135    for (i = 0; i < MLKEM_N; i += 4, zeta++) {
1136        sword16 t0[2];
1137        sword16 t2[2];
1138
1139        mlkem_basemul(t0, a + i + 0, b + i + 0, zeta[0]);
1140        mlkem_basemul(t2, a + i + 2, b + i + 2, (sword16)(-zeta[0]));
1141
1142        r[i + 0] = (sword16)(r[i + 0] + t0[0]);
1143        r[i + 1] = (sword16)(r[i + 1] + t0[1]);
1144        r[i + 2] = (sword16)(r[i + 2] + t2[0]);
1145        r[i + 3] = (sword16)(r[i + 3] + t2[1]);
1146    }
1147#elif defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
1148    /* Four multiplications per loop. */
1149    unsigned int i;
1150    for (i = 0; i < MLKEM_N; i += 8, zeta += 2) {
1151        sword16 t0[2];
1152        sword16 t2[2];
1153        sword16 t4[2];
1154        sword16 t6[2];
1155
1156        mlkem_basemul(t0, a + i + 0, b + i + 0, zeta[0]);
1157        mlkem_basemul(t2, a + i + 2, b + i + 2, (sword16)(-zeta[0]));
1158        mlkem_basemul(t4, a + i + 4, b + i + 4, zeta[1]);
1159        mlkem_basemul(t6, a + i + 6, b + i + 6, (sword16)(-zeta[1]));
1160
1161        r[i + 0] = (sword16)(r[i + 0] + t0[0]);
1162        r[i + 1] = (sword16)(r[i + 1] + t0[1]);
1163        r[i + 2] = (sword16)(r[i + 2] + t2[0]);
1164        r[i + 3] = (sword16)(r[i + 3] + t2[1]);
1165        r[i + 4] = (sword16)(r[i + 4] + t4[0]);
1166        r[i + 5] = (sword16)(r[i + 5] + t4[1]);
1167        r[i + 6] = (sword16)(r[i + 6] + t6[0]);
1168        r[i + 7] = (sword16)(r[i + 7] + t6[1]);
1169    }
1170#else
1171    /* Eight multiplications per loop. */
1172    unsigned int i;
1173    for (i = 0; i < MLKEM_N; i += 16, zeta += 4) {
1174        sword16 t0[2];
1175        sword16 t2[2];
1176        sword16 t4[2];
1177        sword16 t6[2];
1178        sword16 t8[2];
1179        sword16 t10[2];
1180        sword16 t12[2];
1181        sword16 t14[2];
1182
1183        mlkem_basemul(t0, a + i + 0, b + i + 0, zeta[0]);
1184        mlkem_basemul(t2, a + i + 2, b + i + 2, (sword16)(-zeta[0]));
1185        mlkem_basemul(t4, a + i + 4, b + i + 4, zeta[1]);
1186        mlkem_basemul(t6, a + i + 6, b + i + 6, (sword16)(-zeta[1]));
1187        mlkem_basemul(t8, a + i + 8, b + i + 8, zeta[2]);
1188        mlkem_basemul(t10, a + i + 10, b + i + 10, (sword16)(-zeta[2]));
1189        mlkem_basemul(t12, a + i + 12, b + i + 12, zeta[3]);
1190        mlkem_basemul(t14, a + i + 14, b + i + 14, (sword16)(-zeta[3]));
1191
1192        r[i + 0] = (sword16)(r[i + 0] + t0[0]);
1193        r[i + 1] = (sword16)(r[i + 1] + t0[1]);
1194        r[i + 2] = (sword16)(r[i + 2] + t2[0]);
1195        r[i + 3] = (sword16)(r[i + 3] + t2[1]);
1196        r[i + 4] = (sword16)(r[i + 4] + t4[0]);
1197        r[i + 5] = (sword16)(r[i + 5] + t4[1]);
1198        r[i + 6] = (sword16)(r[i + 6] + t6[0]);
1199        r[i + 7] = (sword16)(r[i + 7] + t6[1]);
1200        r[i + 8] = (sword16)(r[i + 8] + t8[0]);
1201        r[i + 9] = (sword16)(r[i + 9] + t8[1]);
1202        r[i + 10] = (sword16)(r[i + 10] + t10[0]);
1203        r[i + 11] = (sword16)(r[i + 11] + t10[1]);
1204        r[i + 12] = (sword16)(r[i + 12] + t12[0]);
1205        r[i + 13] = (sword16)(r[i + 13] + t12[1]);
1206        r[i + 14] = (sword16)(r[i + 14] + t14[0]);
1207        r[i + 15] = (sword16)(r[i + 15] + t14[1]);
1208    }
1209#endif
1210}
1211#endif
1212
1213/* Pointwise multiply elements of a and b, into r, and multiply by 2^-16.
1214 *
1215 * @param  [out]  r  Result polynomial.
1216 * @param  [in]   a  First vector polynomial to multiply with.
1217 * @param  [in]   b  Second vector polynomial to multiply with.
1218 * @param  [in]   k  Number of polynomials in vector.
1219 */
1220static void mlkem_pointwise_acc_mont(sword16* r, const sword16* a,
1221    const sword16* b, unsigned int k)
1222{
1223    unsigned int i;
1224
1225    mlkem_basemul_mont(r, a, b);
1226#ifdef WOLFSSL_MLKEM_SMALL
1227    for (i = 1; i < k; ++i) {
1228        mlkem_basemul_mont_add(r, a + i * MLKEM_N, b + i * MLKEM_N);
1229    }
1230#else
1231    for (i = 1; i < k - 1; ++i) {
1232        mlkem_basemul_mont_add(r, a + i * MLKEM_N, b + i * MLKEM_N);
1233    }
1234    mlkem_basemul_mont_add(r, a + (k - 1) * MLKEM_N, b + (k - 1) * MLKEM_N);
1235#endif
1236}
1237
1238/******************************************************************************/
1239
1240/* Initialize ML-KEM implementation.
1241 */
1242void mlkem_init(void)
1243{
1244#if defined(USE_INTEL_SPEEDUP) || (defined(__aarch64__) && \
1245    defined(WOLFSSL_ARMASM))
1246    cpuid_get_flags_ex(&cpuid_flags);
1247#endif
1248}
1249
1250/******************************************************************************/
1251
1252#if defined(__aarch64__) && defined(WOLFSSL_ARMASM)
1253
1254#ifndef WOLFSSL_MLKEM_NO_MAKE_KEY
1255/* Generate a public-private key pair from randomly generated data.
1256 *
1257 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
1258 *   ...
1259 *   16: s_hat <- NTT(s)
1260 *   17: e_hat <- NTT(e)
1261 *   18: t_hat <- A_hat o s_hat + e_hat
1262 *   ...
1263 *
1264 * @param  [in, out]  s  Private key vector of polynomials.
1265 * @param  [out]      t  Public key vector of polynomials.
1266 * @param  [in, out]  e  Error values as a vector of polynomials. Modified.
1267 * @param  [in]       a  Random values in an array of vectors of polynomials.
1268 * @param  [in]       k  Number of polynomials in vector.
1269 */
1270void mlkem_keygen(sword16* s, sword16* t, sword16* e, const sword16* a, int k)
1271{
1272    int i;
1273
1274#ifndef WOLFSSL_AARCH64_NO_SQRDMLSH
1275    if (IS_AARCH64_RDM(cpuid_flags)) {
1276        /* Transform private key. All of result used in public key calculation.
1277         * Step 16: s_hat = NTT(s) */
1278        for (i = 0; i < k; ++i) {
1279            mlkem_ntt_sqrdmlsh(s + i * MLKEM_N);
1280        }
1281
1282        /* For each polynomial in the vectors.
1283         * Step 17, Step 18: Calculate public from A_hat, s_hat and e_hat. */
1284        for (i = 0; i < k; ++i) {
1285            /* Multiply a by private into public polynomial.
1286             * Step 18: ... A_hat o s_hat ... */
1287            mlkem_pointwise_acc_mont(t + i * MLKEM_N, a + i * k * MLKEM_N, s,
1288                (unsigned int)k);
1289            /* Convert public polynomial to Montgomery form.
1290             * Step 18: ... MontRed(A_hat o s_hat) ... */
1291            mlkem_to_mont_sqrdmlsh(t + i * MLKEM_N);
1292            /* Transform error values polynomial.
1293             * Step 17: e_hat = NTT(e) */
1294            mlkem_ntt_sqrdmlsh(e + i * MLKEM_N);
1295            /* Add errors to public key and reduce.
1296             * Step 18: t_hat = BarrettRed(MontRed(A_hat o s_hat) + e_hat) */
1297            mlkem_add_reduce(t + i * MLKEM_N, e + i * MLKEM_N);
1298        }
1299    }
1300    else
1301#endif
1302    {
1303        /* Transform private key. All of result used in public key calculation.
1304         * Step 16: s_hat = NTT(s) */
1305        for (i = 0; i < k; ++i) {
1306            mlkem_ntt(s + i * MLKEM_N);
1307        }
1308
1309        /* For each polynomial in the vectors.
1310         * Step 17, Step 18: Calculate public from A_hat, s_hat and e_hat. */
1311        for (i = 0; i < k; ++i) {
1312            /* Multiply a by private into public polynomial.
1313             * Step 18: ... A_hat o s_hat ... */
1314            mlkem_pointwise_acc_mont(t + i * MLKEM_N, a + i * k * MLKEM_N, s,
1315                (unsigned int)k);
1316            /* Convert public polynomial to Montgomery form.
1317             * Step 18: ... MontRed(A_hat o s_hat) ... */
1318            mlkem_to_mont(t + i * MLKEM_N);
1319            /* Transform error values polynomial.
1320             * Step 17: e_hat = NTT(e) */
1321            mlkem_ntt(e + i * MLKEM_N);
1322            /* Add errors to public key and reduce.
1323             * Step 18: t_hat = BarrettRed(MontRed(A_hat o s_hat) + e_hat) */
1324            mlkem_add_reduce(t + i * MLKEM_N, e + i * MLKEM_N);
1325        }
1326    }
1327}
1328#endif /* WOLFSSL_MLKEM_NO_MAKE_KEY */
1329
1330#if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
1331    !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
1332/* Encapsulate message.
1333 *
1334 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE, m, r)
1335 *   ...
1336 *   Step 18: y_hat <- NTT(y)
1337 *   Step 19: u <- InvNTT(A_hat_trans o y_hat) + e_1
1338 *   ...
1339 *   Step 21: v <- InvNTT(t_hat_trans o y_hat) + e_2 + mu
1340 *   ...
1341 *
1342 * @param  [in]       t   Public key vector of polynomials.
1343 * @param  [out]      u   Vector of polynomials.
1344 * @param  [out]      v   Polynomial.
1345 * @param  [in]       a   Array of vector of polynomials.
1346 * @param  [in, out]  y   Vector of polynomials.
1347 * @param  [in]       e1  Error Vector of polynomials.
1348 * @param  [in]       e2  Error polynomial.
1349 * @param  [in]       m   Message polynomial.
1350 * @param  [in]       k   Number of polynomials in vector.
1351 */
1352void mlkem_encapsulate(const sword16* t, sword16* u, sword16* v,
1353    const sword16* a, sword16* y, const sword16* e1, const sword16* e2,
1354    const sword16* m, int k)
1355{
1356    int i;
1357
1358#ifndef WOLFSSL_AARCH64_NO_SQRDMLSH
1359    if (IS_AARCH64_RDM(cpuid_flags)) {
1360        /* Transform y. All of result used in calculation of u and v.
1361         * Step 18: y_hat <- NTT(y) */
1362        for (i = 0; i < k; ++i) {
1363            mlkem_ntt_sqrdmlsh(y + i * MLKEM_N);
1364        }
1365
1366        /* For each polynomial in the vectors.
1367         * Step 19: u <- InvNTT(A_hat_trans o y_hat) + e_1 */
1368        for (i = 0; i < k; ++i) {
1369            /* Multiply at by y into u polynomial.
1370             * Step 19: ... A_hat_trans o y_hat ... */
1371            mlkem_pointwise_acc_mont(u + i * MLKEM_N, a + i * k * MLKEM_N, y,
1372                (unsigned int)k);
1373            /* Inverse transform u polynomial.
1374             * Step 19: ... InvNTT(A_hat_trans o y_hat) ... */
1375            mlkem_invntt_sqrdmlsh(u + i * MLKEM_N);
1376            /* Add errors to u and reduce.
1377             * Step 19: u <- InvNTT(A_hat_trans o y_hat) + e_1 */
1378            mlkem_add_reduce(u + i * MLKEM_N, e1 + i * MLKEM_N);
1379        }
1380
1381        /* Multiply public key by y into v polynomial.
1382         * Step 21: ... t_hat_trans o y_hat ... */
1383        mlkem_pointwise_acc_mont(v, t, y, (unsigned int)k);
1384        /* Inverse transform v.
1385         * Step 21: ... InvNTT(t_hat_trans o y_hat) ... */
1386        mlkem_invntt_sqrdmlsh(v);
1387    }
1388    else
1389#endif
1390    {
1391        /* Transform y. All of result used in calculation of u and v.
1392         * Step 18: y_hat <- NTT(y) */
1393        for (i = 0; i < k; ++i) {
1394            mlkem_ntt(y + i * MLKEM_N);
1395        }
1396
1397        /* For each polynomial in the vectors.
1398         * Step 19: u <- InvNTT(A_hat_trans o y_hat) + e_1 */
1399        for (i = 0; i < k; ++i) {
1400            /* Multiply at by y into u polynomial.
1401             * Step 19: ... A_hat_trans o y_hat ... */
1402            mlkem_pointwise_acc_mont(u + i * MLKEM_N, a + i * k * MLKEM_N, y,
1403                (unsigned int)k);
1404            /* Inverse transform u polynomial.
1405             * Step 19: ... InvNTT(A_hat_trans o y_hat) ... */
1406            mlkem_invntt(u + i * MLKEM_N);
1407            /* Add errors to u and reduce.
1408             * Step 19: u <- InvNTT(A_hat_trans o y_hat) + e_1 */
1409            mlkem_add_reduce(u + i * MLKEM_N, e1 + i * MLKEM_N);
1410        }
1411
1412        /* Multiply public key by y into v polynomial.
1413         * Step 21: ... t_hat_trans o y_hat ... */
1414        mlkem_pointwise_acc_mont(v, t, y, (unsigned int)k);
1415        /* Inverse transform v.
1416         * Step 21: ... InvNTT(t_hat_trans o y_hat) ... */
1417        mlkem_invntt(v);
1418    }
1419    /* Add errors and message to v and reduce.
1420     * Step 21: v <- InvNTT(t_hat_trans o y_hat) + e_2 + mu */
1421    mlkem_add3_reduce(v, e2, m);
1422}
1423#endif /* !WOLFSSL_MLKEM_NO_ENCAPSULATE || !WOLFSSL_MLKEM_NO_DECAPSULATE */
1424
1425#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE
1426/* Decapsulate message.
1427 *
1428 * FIPS 203, Algorithm 15: K-PKE.Decrypt(dk_PKE,c)
1429 * Uses the decryption key to decrypt a ciphertext.
1430 *   ...
1431 *   6: w <- v' - InvNTT(s_hat_trans o NTT(u'))
1432 *   ...
1433 *
1434 * @param  [in]       s  Decryption key as vector of polynomials.
1435 * @param  [out]      w  Message polynomial.
1436 * @param  [in, out]  u  Vector of polynomials containing error.
1437 * @param  [in]       v  Encapsulated message polynomial.
1438 * @param  [in]       k  Number of polynomials in vector.
1439 */
1440void mlkem_decapsulate(const sword16* s, sword16* w, sword16* u,
1441    const sword16* v, int k)
1442{
1443    int i;
1444
1445#ifndef WOLFSSL_AARCH64_NO_SQRDMLSH
1446    if (IS_AARCH64_RDM(cpuid_flags)) {
1447        /* Transform u. All of result used in calculation of w.
1448         * Step 6: ... NTT(u') */
1449        for (i = 0; i < k; ++i) {
1450            mlkem_ntt_sqrdmlsh(u + i * MLKEM_N);
1451        }
1452
1453        /* Multiply private key by u into w polynomial.
1454         * Step 6: ... s_hat_trans o NTT(u') */
1455        mlkem_pointwise_acc_mont(w, s, u, (unsigned int)k);
1456        /* Inverse transform w.
1457         * Step 6: ... InvNTT(s_hat_trans o NTT(u')) */
1458        mlkem_invntt_sqrdmlsh(w);
1459    }
1460    else
1461#endif
1462    {
1463        /* Transform u. All of result used in calculation of w.
1464         * Step 6: ... NTT(u') */
1465        for (i = 0; i < k; ++i) {
1466            mlkem_ntt(u + i * MLKEM_N);
1467        }
1468
1469        /* Multiply private key by u into w polynomial.
1470         * Step 6: ... s_hat_trans o NTT(u') */
1471        mlkem_pointwise_acc_mont(w, s, u, (unsigned int)k);
1472        /* Inverse transform w.
1473         * Step 6: ... InvNTT(s_hat_trans o NTT(u')) */
1474        mlkem_invntt(w);
1475    }
1476    /* Subtract errors (in w) out of v and reduce into w.
1477     * Step 6: w <- v' - InvNTT(s_hat_trans o NTT(u')) */
1478    mlkem_rsub_reduce(w, v);
1479}
1480#endif /* !WOLFSSL_MLKEM_NO_DECAPSULATE */
1481
1482#else
1483
1484#ifndef WOLFSSL_MLKEM_NO_MAKE_KEY
1485
1486#if !defined(WOLFSSL_MLKEM_SMALL) && !defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
1487/* Number-Theoretic Transform.
1488 *
1489 * FIPS 203, Algorithm 9: NTT(f)
1490 * Computes the NTT representation f_hat of the given polynomial f element of
1491 * R_q.
1492 *   1: f_hat <- f
1493 *   2: i <- 1
1494 *   3: for (len <- 128; len >= 2; len <- len/2)
1495 *   4:     for (start <- 0; start < 256; start <- start + 2.len)
1496 *   5:         zeta <- zetas^BitRev_7(i) mod q
1497 *   6:         i <- i + 1
1498 *   7:         for (j <- start; j < start + len; j++)
1499 *   8:             t <- zeta.f[j+len]
1500 *   9:             f_hat[j+len] <- f_hat[j] - t
1501 *  10:             f_hat[j] <- f_hat[j] + t
1502 *  11:         end for
1503 *  12:     end for
1504 *  13: end for
1505 *  14: return f_hat
1506 *
1507 * @param [in, out]  r  Polynomial to transform.
1508 * @param [in, out]  a  Polynomial to add NTT result to.
1509 */
1510static void mlkem_ntt_add_to(sword16* r, sword16* a)
1511{
1512#if defined(WOLFSSL_MLKEM_NTT_UNROLL)
1513    /* Unroll len loop (Step 3). */
1514    unsigned int k = 1;
1515    unsigned int j;
1516    unsigned int start;
1517    sword16 zeta = zetas[k++];
1518
1519    /* len = 128 */
1520    for (j = 0; j < MLKEM_N / 2; ++j) {
1521        sword32 p = (sword32)zeta * r[j + MLKEM_N / 2];
1522        sword16 t = MLKEM_MONT_RED(p);
1523        sword16 rj = r[j];
1524        r[j + MLKEM_N / 2] = rj - t;
1525        r[j] = rj + t;
1526    }
1527    /* len = 64 */
1528    for (start = 0; start < MLKEM_N; start += 2 * 64) {
1529        zeta = zetas[k++];
1530        for (j = 0; j < 64; ++j) {
1531            sword32 p = (sword32)zeta * r[start + j + 64];
1532            sword16 t = MLKEM_MONT_RED(p);
1533            sword16 rj = r[start + j];
1534            r[start + j + 64] = rj - t;
1535            r[start + j] = rj + t;
1536        }
1537    }
1538    /* len = 32 */
1539    for (start = 0; start < MLKEM_N; start += 2 * 32) {
1540        zeta = zetas[k++];
1541        for (j = 0; j < 32; ++j) {
1542            sword32 p = (sword32)zeta * r[start + j + 32];
1543            sword16 t = MLKEM_MONT_RED(p);
1544            sword16 rj = r[start + j];
1545            r[start + j + 32] = rj - t;
1546            r[start + j] = rj + t;
1547        }
1548    }
1549    /* len = 16 */
1550    for (start = 0; start < MLKEM_N; start += 2 * 16) {
1551        zeta = zetas[k++];
1552        for (j = 0; j < 16; ++j) {
1553            sword32 p = (sword32)zeta * r[start + j + 16];
1554            sword16 t = MLKEM_MONT_RED(p);
1555            sword16 rj = r[start + j];
1556            r[start + j + 16] = rj - t;
1557            r[start + j] = rj + t;
1558        }
1559    }
1560    /* len = 8 */
1561    for (start = 0; start < MLKEM_N; start += 2 * 8) {
1562        zeta = zetas[k++];
1563        for (j = 0; j < 8; ++j) {
1564            sword32 p = (sword32)zeta * r[start + j + 8];
1565            sword16 t = MLKEM_MONT_RED(p);
1566            sword16 rj = r[start + j];
1567            r[start + j + 8] = rj - t;
1568            r[start + j] = rj + t;
1569        }
1570    }
1571    /* len = 4 */
1572    for (start = 0; start < MLKEM_N; start += 2 * 4) {
1573        zeta = zetas[k++];
1574        for (j = 0; j < 4; ++j) {
1575            sword32 p = (sword32)zeta * r[start + j + 4];
1576            sword16 t = MLKEM_MONT_RED(p);
1577            sword16 rj = r[start + j];
1578            r[start + j + 4] = rj - t;
1579            r[start + j] = rj + t;
1580        }
1581    }
1582    /* len = 2 */
1583    for (start = 0; start < MLKEM_N; start += 2 * 2) {
1584        zeta = zetas[k++];
1585        for (j = 0; j < 2; ++j) {
1586            sword32 p = (sword32)zeta * r[start + j + 2];
1587            sword16 t = MLKEM_MONT_RED(p);
1588            sword16 rj = r[start + j];
1589            r[start + j + 2] = rj - t;
1590            r[start + j] = rj + t;
1591        }
1592    }
1593    /* Reduce coefficients with quick algorithm. */
1594    for (j = 0; j < MLKEM_N; ++j) {
1595        sword16 t = a[j] + r[j];
1596        a[j] = MLKEM_BARRETT_RED(t);
1597    }
1598#else /* !WOLFSSL_MLKEM_NTT_UNROLL */
1599    /* Unroll len (2, 3, 2) and start loops. */
1600    unsigned int j;
1601    sword16 t0;
1602    sword16 t1;
1603    sword16 t2;
1604    sword16 t3;
1605
1606    /* len = 128,64 */
1607    sword16 zeta128 = zetas[1];
1608    sword16 zeta64_0 = zetas[2];
1609    sword16 zeta64_1 = zetas[3];
1610    for (j = 0; j < MLKEM_N / 8; j++) {
1611        sword16 r0 = r[j +   0];
1612        sword16 r1 = r[j +  32];
1613        sword16 r2 = r[j +  64];
1614        sword16 r3 = r[j +  96];
1615        sword16 r4 = r[j + 128];
1616        sword16 r5 = r[j + 160];
1617        sword16 r6 = r[j + 192];
1618        sword16 r7 = r[j + 224];
1619
1620        t0 = MLKEM_MONT_RED((sword32)zeta128 * r4);
1621        t1 = MLKEM_MONT_RED((sword32)zeta128 * r5);
1622        t2 = MLKEM_MONT_RED((sword32)zeta128 * r6);
1623        t3 = MLKEM_MONT_RED((sword32)zeta128 * r7);
1624        r4 = (sword16)(r0 - t0);
1625        r5 = (sword16)(r1 - t1);
1626        r6 = (sword16)(r2 - t2);
1627        r7 = (sword16)(r3 - t3);
1628        r0 = (sword16)(r0 + t0);
1629        r1 = (sword16)(r1 + t1);
1630        r2 = (sword16)(r2 + t2);
1631        r3 = (sword16)(r3 + t3);
1632
1633        t0 = MLKEM_MONT_RED((sword32)zeta64_0 * r2);
1634        t1 = MLKEM_MONT_RED((sword32)zeta64_0 * r3);
1635        t2 = MLKEM_MONT_RED((sword32)zeta64_1 * r6);
1636        t3 = MLKEM_MONT_RED((sword32)zeta64_1 * r7);
1637        r2 = (sword16)(r0 - t0);
1638        r3 = (sword16)(r1 - t1);
1639        r6 = (sword16)(r4 - t2);
1640        r7 = (sword16)(r5 - t3);
1641        r0 = (sword16)(r0 + t0);
1642        r1 = (sword16)(r1 + t1);
1643        r4 = (sword16)(r4 + t2);
1644        r5 = (sword16)(r5 + t3);
1645
1646        r[j +   0] = r0;
1647        r[j +  32] = r1;
1648        r[j +  64] = r2;
1649        r[j +  96] = r3;
1650        r[j + 128] = r4;
1651        r[j + 160] = r5;
1652        r[j + 192] = r6;
1653        r[j + 224] = r7;
1654    }
1655
1656    /* len = 32,16,8 */
1657    for (j = 0; j < MLKEM_N; j += 64) {
1658        unsigned int i;
1659        sword16 zeta32   = zetas[ 4 + j / 64 + 0];
1660        sword16 zeta16_0 = zetas[ 8 + j / 32 + 0];
1661        sword16 zeta16_1 = zetas[ 8 + j / 32 + 1];
1662        sword16 zeta8_0  = zetas[16 + j / 16 + 0];
1663        sword16 zeta8_1  = zetas[16 + j / 16 + 1];
1664        sword16 zeta8_2  = zetas[16 + j / 16 + 2];
1665        sword16 zeta8_3  = zetas[16 + j / 16 + 3];
1666        for (i = 0; i < 8; i++) {
1667            sword16 r0 = r[j + i +  0];
1668            sword16 r1 = r[j + i +  8];
1669            sword16 r2 = r[j + i + 16];
1670            sword16 r3 = r[j + i + 24];
1671            sword16 r4 = r[j + i + 32];
1672            sword16 r5 = r[j + i + 40];
1673            sword16 r6 = r[j + i + 48];
1674            sword16 r7 = r[j + i + 56];
1675
1676            t0 = MLKEM_MONT_RED((sword32)zeta32 * r4);
1677            t1 = MLKEM_MONT_RED((sword32)zeta32 * r5);
1678            t2 = MLKEM_MONT_RED((sword32)zeta32 * r6);
1679            t3 = MLKEM_MONT_RED((sword32)zeta32 * r7);
1680            r4 = (sword16)(r0 - t0);
1681            r5 = (sword16)(r1 - t1);
1682            r6 = (sword16)(r2 - t2);
1683            r7 = (sword16)(r3 - t3);
1684            r0 = (sword16)(r0 + t0);
1685            r1 = (sword16)(r1 + t1);
1686            r2 = (sword16)(r2 + t2);
1687            r3 = (sword16)(r3 + t3);
1688
1689            t0 = MLKEM_MONT_RED((sword32)zeta16_0 * r2);
1690            t1 = MLKEM_MONT_RED((sword32)zeta16_0 * r3);
1691            t2 = MLKEM_MONT_RED((sword32)zeta16_1 * r6);
1692            t3 = MLKEM_MONT_RED((sword32)zeta16_1 * r7);
1693            r2 = (sword16)(r0 - t0);
1694            r3 = (sword16)(r1 - t1);
1695            r6 = (sword16)(r4 - t2);
1696            r7 = (sword16)(r5 - t3);
1697            r0 = (sword16)(r0 + t0);
1698            r1 = (sword16)(r1 + t1);
1699            r4 = (sword16)(r4 + t2);
1700            r5 = (sword16)(r5 + t3);
1701
1702            t0 = MLKEM_MONT_RED((sword32)zeta8_0 * r1);
1703            t1 = MLKEM_MONT_RED((sword32)zeta8_1 * r3);
1704            t2 = MLKEM_MONT_RED((sword32)zeta8_2 * r5);
1705            t3 = MLKEM_MONT_RED((sword32)zeta8_3 * r7);
1706            r1 = (sword16)(r0 - t0);
1707            r3 = (sword16)(r2 - t1);
1708            r5 = (sword16)(r4 - t2);
1709            r7 = (sword16)(r6 - t3);
1710            r0 = (sword16)(r0 + t0);
1711            r2 = (sword16)(r2 + t1);
1712            r4 = (sword16)(r4 + t2);
1713            r6 = (sword16)(r6 + t3);
1714
1715            r[j + i +  0] = r0;
1716            r[j + i +  8] = r1;
1717            r[j + i + 16] = r2;
1718            r[j + i + 24] = r3;
1719            r[j + i + 32] = r4;
1720            r[j + i + 40] = r5;
1721            r[j + i + 48] = r6;
1722            r[j + i + 56] = r7;
1723        }
1724    }
1725
1726    /* len = 4,2 and Final reduction */
1727    for (j = 0; j < MLKEM_N; j += 8) {
1728        sword16 zeta4  = zetas[32 + j / 8 + 0];
1729        sword16 zeta2_0 = zetas[64 + j / 4 + 0];
1730        sword16 zeta2_1 = zetas[64 + j / 4 + 1];
1731        sword16 r0 = r[j + 0];
1732        sword16 r1 = r[j + 1];
1733        sword16 r2 = r[j + 2];
1734        sword16 r3 = r[j + 3];
1735        sword16 r4 = r[j + 4];
1736        sword16 r5 = r[j + 5];
1737        sword16 r6 = r[j + 6];
1738        sword16 r7 = r[j + 7];
1739
1740        t0 = MLKEM_MONT_RED((sword32)zeta4 * r4);
1741        t1 = MLKEM_MONT_RED((sword32)zeta4 * r5);
1742        t2 = MLKEM_MONT_RED((sword32)zeta4 * r6);
1743        t3 = MLKEM_MONT_RED((sword32)zeta4 * r7);
1744        r4 = (sword16)(r0 - t0);
1745        r5 = (sword16)(r1 - t1);
1746        r6 = (sword16)(r2 - t2);
1747        r7 = (sword16)(r3 - t3);
1748        r0 = (sword16)(r0 + t0);
1749        r1 = (sword16)(r1 + t1);
1750        r2 = (sword16)(r2 + t2);
1751        r3 = (sword16)(r3 + t3);
1752
1753        t0 = MLKEM_MONT_RED((sword32)zeta2_0 * r2);
1754        t1 = MLKEM_MONT_RED((sword32)zeta2_0 * r3);
1755        t2 = MLKEM_MONT_RED((sword32)zeta2_1 * r6);
1756        t3 = MLKEM_MONT_RED((sword32)zeta2_1 * r7);
1757        r2 = (sword16)(r0 - t0);
1758        r3 = (sword16)(r1 - t1);
1759        r6 = (sword16)(r4 - t2);
1760        r7 = (sword16)(r5 - t3);
1761        r0 = (sword16)(r0 + t0);
1762        r1 = (sword16)(r1 + t1);
1763        r4 = (sword16)(r4 + t2);
1764        r5 = (sword16)(r5 + t3);
1765
1766        r0 = (sword16)(r0 + a[j + 0]);
1767        r1 = (sword16)(r1 + a[j + 1]);
1768        r2 = (sword16)(r2 + a[j + 2]);
1769        r3 = (sword16)(r3 + a[j + 3]);
1770        r4 = (sword16)(r4 + a[j + 4]);
1771        r5 = (sword16)(r5 + a[j + 5]);
1772        r6 = (sword16)(r6 + a[j + 6]);
1773        r7 = (sword16)(r7 + a[j + 7]);
1774
1775        a[j + 0] = MLKEM_BARRETT_RED(r0);
1776        a[j + 1] = MLKEM_BARRETT_RED(r1);
1777        a[j + 2] = MLKEM_BARRETT_RED(r2);
1778        a[j + 3] = MLKEM_BARRETT_RED(r3);
1779        a[j + 4] = MLKEM_BARRETT_RED(r4);
1780        a[j + 5] = MLKEM_BARRETT_RED(r5);
1781        a[j + 6] = MLKEM_BARRETT_RED(r6);
1782        a[j + 7] = MLKEM_BARRETT_RED(r7);
1783    }
1784#endif /* !WOLFSSL_MLKEM_NTT_UNROLL */
1785}
1786#endif /* !WOLFSSL_MLKEM_SMALL && !WOLFSSL_MLKEM_NO_LARGE_CODE */
1787
1788#ifndef WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM
1789/* Generate a public-private key pair from randomly generated data.
1790 *
1791 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
1792 *   ...
1793 *   16: s_hat <- NTT(s)
1794 *   17: e_hat <- NTT(e)
1795 *   18: t_hat <- A_hat o s_hat + e_hat
1796 *   ...
1797 *
1798 * @param  [in, out]  s  Private key vector of polynomials.
1799 * @param  [out]      t  Public key vector of polynomials.
1800 * @param  [in, out]  e  Error values as a vector of polynomials. Modified.
1801 * @param  [in]       a  Random values in an array of vectors of polynomials.
1802 * @param  [in]       k  Number of polynomials in vector.
1803 */
1804static void mlkem_keygen_c(sword16* s, sword16* t, sword16* e, const sword16* a,
1805    int k)
1806{
1807    int i;
1808
1809    /* Transform private key. All of result used in public key calculation
1810     * Step 16: s_hat = NTT(s) */
1811    for (i = 0; i < k; ++i) {
1812        mlkem_ntt(s + i * MLKEM_N);
1813    }
1814
1815    /* For each polynomial in the vectors.
1816     * Step 17, Step 18: Calculate public from A_hat, s_hat and e_hat. */
1817    for (i = 0; i < k; ++i) {
1818        int j;
1819
1820        /* Multiply a by private into public polynomial.
1821         * Step 18: ... A_hat o s_hat ... */
1822        mlkem_pointwise_acc_mont(t + i * MLKEM_N, a + i * k * MLKEM_N, s,
1823            (unsigned int)k);
1824        /* Convert public polynomial to Montgomery form.
1825         * Step 18: ... MontRed(A_hat o s_hat) ... */
1826        for (j = 0; j < MLKEM_N; ++j) {
1827            sword32 n = t[i * MLKEM_N + j] * (sword32)MLKEM_F;
1828            t[i * MLKEM_N + j] = MLKEM_MONT_RED(n);
1829        }
1830        /* Transform error values polynomial.
1831         * Step 17: e_hat = NTT(e) */
1832#if defined(WOLFSSL_MLKEM_SMALL) || defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
1833        mlkem_ntt(e + i * MLKEM_N);
1834        /* Add errors to public key and reduce.
1835         * Step 18: t_hat = BarrettRed(MontRed(A_hat o s_hat) + e_hat) */
1836        for (j = 0; j < MLKEM_N; ++j) {
1837            sword16 n = (sword16)(t[i * MLKEM_N + j] + e[i * MLKEM_N + j]);
1838            t[i * MLKEM_N + j] = MLKEM_BARRETT_RED(n);
1839        }
1840#else
1841        /* Add errors to public key and reduce.
1842         * Step 18: t_hat = BarrettRed(MontRed(A_hat o s_hat) + e_hat) */
1843        mlkem_ntt_add_to(e + i * MLKEM_N, t + i * MLKEM_N);
1844#endif
1845    }
1846}
1847
1848/* Generate a public-private key pair from randomly generated data.
1849 *
1850 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
1851 *   ...
1852 *   16: s_hat <- NTT(s)
1853 *   17: e_hat <- NTT(e)
1854 *   18: t_hat <- A_hat o s_hat + e_hat
1855 *   ...
1856 *
1857 * @param  [in, out]  s  Private key vector of polynomials.
1858 * @param  [out]      t  Public key vector of polynomials.
1859 * @param  [in, out]  e  Error values as a vector of polynomials. Modified.
1860 * @param  [in]       a  Random values in an array of vectors of polynomials.
1861 * @param  [in]       k  Number of polynomials in vector.
1862 */
1863void mlkem_keygen(sword16* s, sword16* t, sword16* e, const sword16* a, int k)
1864{
1865#ifdef USE_INTEL_SPEEDUP
1866    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
1867        /* Alg 13: Steps 16-18 */
1868        mlkem_keygen_avx2(s, t, e, a, k);
1869        RESTORE_VECTOR_REGISTERS();
1870    }
1871    else
1872#endif
1873    {
1874        /* Alg 13: Steps 16-18 */
1875        mlkem_keygen_c(s, t, e, a, k);
1876    }
1877}
1878
1879#else /* WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM */
1880
1881/* Generate a public-private key pair from randomly generated data.
1882 *
1883 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
1884 *   3: for (i <- 0; i < k; i++)                         > generate matrix A_hat
1885 *   ... (generate A[i])
1886 *   7: end for
1887 *   ...
1888 *  13:      e[i] <- SamplePolyCBD_eta_1(PRF_eta_1(sigma, N))
1889 *   ...
1890 *  16: s_hat <- NTT(s)
1891 *  17: e_hat <- NTT(e)
1892 *  18: t_hat <- A_hat o s_hat + e_hat
1893 *   ...
1894 *
1895 * @param  [in, out]  s      Private key vector of polynomials.
1896 * @param  [out]      t      Public key vector of polynomials.
1897 * @param  [in, out]  prf    XOF object.
1898 * @param  [in]       tv     Temporary vector of polynomials.
1899 * @param  [in]       k      Number of polynomials in vector.
1900 * @param  [in]       rho    Random seed to generate matrix A from.
1901 * @param  [in, out]  sigma  Random seed to generate noise from.
1902 * @return  0 on success.
1903 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
1904 *          WOLFSSL_SMALL_STACK is defined.
1905 * @return  Other negative value when a hash error occurred.
1906 */
1907int mlkem_keygen_seeds(sword16* s, sword16* t, MLKEM_PRF_T* prf,
1908    sword16* tv, int k, byte* rho, byte* sigma)
1909{
1910    int i;
1911    int ret = 0;
1912    sword16* ai = tv;
1913    sword16* e = tv;
1914
1915    /* Transform private key. All of result used in public key calculation
1916     * Step 16: s_hat = NTT(s) */
1917    for (i = 0; i < k; ++i) {
1918        mlkem_ntt(s + i * MLKEM_N);
1919    }
1920
1921    /* For each polynomial in the vectors.
1922     * Step 17, Step 18: Calculate public from A_hat, s_hat and e_hat. */
1923    for (i = 0; i < k; ++i) {
1924        int j;
1925
1926        /* Generate a vector of matrix A.
1927         * Steps 4-6: generate A[i] */
1928        ret = mlkem_gen_matrix_i(prf, ai, k, rho, i, 0);
1929        if (ret != 0) {
1930           break;
1931        }
1932
1933        /* Multiply a by private into public polynomial.
1934         * Step 18: ... A_hat o s_hat ... */
1935        mlkem_pointwise_acc_mont(t + i * MLKEM_N, ai, s, (unsigned int)k);
1936        /* Convert public polynomial to Montgomery form.
1937         * Step 18: ... MontRed(A_hat o s_hat) ... */
1938        for (j = 0; j < MLKEM_N; ++j) {
1939            sword32 n = t[i * MLKEM_N + j] * (sword32)MLKEM_F;
1940            t[i * MLKEM_N + j] = MLKEM_MONT_RED(n);
1941        }
1942
1943        /* Generate noise using PRF.
1944         * Step 13: e[i] <- SamplePolyCBD_eta_1(PRF_eta_1(sigma, N)) */
1945        ret = mlkem_get_noise_i(prf, k, e, sigma, i, 1);
1946        if (ret != 0) {
1947           break;
1948        }
1949        /* Transform error values polynomial.
1950         * Step 17: e_hat = NTT(e) */
1951#if defined(WOLFSSL_MLKEM_SMALL) || defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
1952        mlkem_ntt(e);
1953        /* Add errors to public key and reduce.
1954         * Step 18: t_hat = BarrettRed(MontRed(A_hat o s_hat) + e_hat) */
1955        for (j = 0; j < MLKEM_N; ++j) {
1956            sword16 n = (sword16)(t[i * MLKEM_N + j] + e[j]);
1957            t[i * MLKEM_N + j] = MLKEM_BARRETT_RED(n);
1958        }
1959#else
1960        /* Add errors to public key and reduce.
1961         * Step 18: t_hat = BarrettRed(MontRed(A_hat o s_hat) + e_hat) */
1962        mlkem_ntt_add_to(e, t + i * MLKEM_N);
1963#endif
1964    }
1965
1966    return ret;
1967}
1968
1969#endif /* WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM */
1970#endif /* !WOLFSSL_MLKEM_NO_MAKE_KEY */
1971
1972#if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
1973    !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
1974#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM
1975/* Encapsulate message.
1976 *
1977 * @param  [in]       pub  Public key vector of polynomials.
1978 * @param  [out]      u    Vector of polynomials.
1979 * @param  [out]      v    Polynomial.
1980 * @param  [in]       a    Array of vector of polynomials.
1981 * @param  [in, out]  y    Vector of polynomials.
1982 * @param  [in]       e1   Error Vector of polynomials.
1983 * @param  [in]       e2   Error polynomial.
1984 * @param  [in]       m    Message polynomial.
1985 * @param  [in]       k    Number of polynomials in vector.
1986 */
1987static void mlkem_encapsulate_c(const sword16* pub, sword16* u, sword16* v,
1988    const sword16* a, sword16* y, const sword16* e1, const sword16* e2,
1989    const sword16* m, int k)
1990{
1991    int i;
1992
1993    /* Transform y. All of result used in calculation of u and v. */
1994    for (i = 0; i < k; ++i) {
1995        mlkem_ntt(y + i * MLKEM_N);
1996    }
1997
1998    /* For each polynomial in the vectors. */
1999    for (i = 0; i < k; ++i) {
2000        int j;
2001
2002        /* Multiply at by y into u polynomial. */
2003        mlkem_pointwise_acc_mont(u + i * MLKEM_N, a + i * k * MLKEM_N, y,
2004            (unsigned int)k);
2005        /* Inverse transform u polynomial. */
2006        mlkem_invntt(u + i * MLKEM_N);
2007        /* Add errors to u and reduce. */
2008#if defined(WOLFSSL_MLKEM_SMALL) || defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
2009        for (j = 0; j < MLKEM_N; ++j) {
2010            sword16 t = (sword16)(u[i * MLKEM_N + j] + e1[i * MLKEM_N + j]);
2011            u[i * MLKEM_N + j] = MLKEM_BARRETT_RED(t);
2012        }
2013#else
2014        for (j = 0; j < MLKEM_N; j += 8) {
2015            sword16 t0 = (sword16)(u[i * MLKEM_N + j + 0] +
2016                                   e1[i * MLKEM_N + j + 0]);
2017            sword16 t1 = (sword16)(u[i * MLKEM_N + j + 1] +
2018                                   e1[i * MLKEM_N + j + 1]);
2019            sword16 t2 = (sword16)(u[i * MLKEM_N + j + 2] +
2020                                   e1[i * MLKEM_N + j + 2]);
2021            sword16 t3 = (sword16)(u[i * MLKEM_N + j + 3] +
2022                                   e1[i * MLKEM_N + j + 3]);
2023            sword16 t4 = (sword16)(u[i * MLKEM_N + j + 4] +
2024                                   e1[i * MLKEM_N + j + 4]);
2025            sword16 t5 = (sword16)(u[i * MLKEM_N + j + 5] +
2026                                   e1[i * MLKEM_N + j + 5]);
2027            sword16 t6 = (sword16)(u[i * MLKEM_N + j + 6] +
2028                                   e1[i * MLKEM_N + j + 6]);
2029            sword16 t7 = (sword16)(u[i * MLKEM_N + j + 7] +
2030                                   e1[i * MLKEM_N + j + 7]);
2031            u[i * MLKEM_N + j + 0] = MLKEM_BARRETT_RED(t0);
2032            u[i * MLKEM_N + j + 1] = MLKEM_BARRETT_RED(t1);
2033            u[i * MLKEM_N + j + 2] = MLKEM_BARRETT_RED(t2);
2034            u[i * MLKEM_N + j + 3] = MLKEM_BARRETT_RED(t3);
2035            u[i * MLKEM_N + j + 4] = MLKEM_BARRETT_RED(t4);
2036            u[i * MLKEM_N + j + 5] = MLKEM_BARRETT_RED(t5);
2037            u[i * MLKEM_N + j + 6] = MLKEM_BARRETT_RED(t6);
2038            u[i * MLKEM_N + j + 7] = MLKEM_BARRETT_RED(t7);
2039        }
2040#endif
2041    }
2042
2043    /* Multiply public key by y into v polynomial. */
2044    mlkem_pointwise_acc_mont(v, pub, y, (unsigned int)k);
2045    /* Inverse transform v. */
2046    mlkem_invntt(v);
2047    /* Add errors and message to v and reduce. */
2048    for (i = 0; i < MLKEM_N; ++i) {
2049        sword16 t = (sword16)(v[i] + e2[i] + m[i]);
2050        v[i] = MLKEM_BARRETT_RED(t);
2051    }
2052}
2053
2054/* Encapsulate message.
2055 *
2056 * @param  [in]       pub  Public key vector of polynomials.
2057 * @param  [out]      u    Vector of polynomials.
2058 * @param  [out]      v    Polynomial.
2059 * @param  [in]       a    Array of vector of polynomials.
2060 * @param  [in, out]  y    Vector of polynomials.
2061 * @param  [in]       e1   Error Vector of polynomials.
2062 * @param  [in]       e2   Error polynomial.
2063 * @param  [in]       m    Message polynomial.
2064 * @param  [in]       k    Number of polynomials in vector.
2065 */
2066void mlkem_encapsulate(const sword16* pub, sword16* u, sword16* v,
2067    const sword16* a, sword16* y, const sword16* e1, const sword16* e2,
2068    const sword16* m, int k)
2069{
2070#ifdef USE_INTEL_SPEEDUP
2071    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
2072        mlkem_encapsulate_avx2(pub, u, v, a, y, e1, e2, m, k);
2073        RESTORE_VECTOR_REGISTERS();
2074    }
2075    else
2076#endif
2077    {
2078        mlkem_encapsulate_c(pub, u, v, a, y, e1, e2, m, k);
2079    }
2080}
2081
2082#else
2083
2084/* Encapsulate message.
2085 *
2086 * @param  [in]       pub    Public key vector of polynomials.
2087 * @param  [in, out]  prf    XOF object.
2088 * @param  [out]      u      Vector of polynomials.
2089 * @param  [in, out]  tp     Polynomial.
2090 * @param  [in, out]  y      Vector of polynomials.
2091 * @param  [in]       k      Number of polynomials in vector.
2092 * @param  [in]       msg    Message to encapsulate.
2093 * @param  [in]       seed   Random seed to generate matrix A from.
2094 * @param  [in, out]  coins  Random seed to generate noise from.
2095 * @return  0 on success.
2096 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
2097 *          WOLFSSL_SMALL_STACK is defined.
2098 * @return  Other negative value when a hash error occurred.
2099 */
2100int mlkem_encapsulate_seeds(const sword16* pub, MLKEM_PRF_T* prf, sword16* u,
2101    sword16* tp, sword16* y, int k, const byte* msg, byte* seed, byte* coins)
2102{
2103    int ret = 0;
2104    int i;
2105    sword16* a = tp;
2106    sword16* e1 = tp;
2107    sword16* v = tp;
2108    sword16* e2 = tp + MLKEM_N;
2109    sword16* m = y;
2110
2111    /* Transform y. All of result used in calculation of u and v. */
2112    for (i = 0; i < k; ++i) {
2113        mlkem_ntt(y + i * MLKEM_N);
2114    }
2115
2116    /* For each polynomial in the vectors. */
2117    for (i = 0; i < k; ++i) {
2118        int j;
2119
2120        /* Generate a vector of matrix A. */
2121        ret = mlkem_gen_matrix_i(prf, a, k, seed, i, 1);
2122        if (ret != 0) {
2123           break;
2124        }
2125
2126        /* Multiply at by y into u polynomial. */
2127        mlkem_pointwise_acc_mont(u + i * MLKEM_N, a, y, (unsigned int)k);
2128        /* Inverse transform u polynomial. */
2129        mlkem_invntt(u + i * MLKEM_N);
2130
2131        /* Generate noise using PRF. */
2132        ret = mlkem_get_noise_i(prf, k, e1, coins, i, 0);
2133        if (ret != 0) {
2134           break;
2135        }
2136        /* Add errors to u and reduce. */
2137#if defined(WOLFSSL_MLKEM_SMALL) || defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
2138        for (j = 0; j < MLKEM_N; ++j) {
2139            sword16 t = (sword16)(u[i * MLKEM_N + j] + e1[j]);
2140            u[i * MLKEM_N + j] = MLKEM_BARRETT_RED(t);
2141        }
2142#else
2143        for (j = 0; j < MLKEM_N; j += 8) {
2144            sword16 t0 = (sword16)(u[i * MLKEM_N + j + 0] + e1[j + 0]);
2145            sword16 t1 = (sword16)(u[i * MLKEM_N + j + 1] + e1[j + 1]);
2146            sword16 t2 = (sword16)(u[i * MLKEM_N + j + 2] + e1[j + 2]);
2147            sword16 t3 = (sword16)(u[i * MLKEM_N + j + 3] + e1[j + 3]);
2148            sword16 t4 = (sword16)(u[i * MLKEM_N + j + 4] + e1[j + 4]);
2149            sword16 t5 = (sword16)(u[i * MLKEM_N + j + 5] + e1[j + 5]);
2150            sword16 t6 = (sword16)(u[i * MLKEM_N + j + 6] + e1[j + 6]);
2151            sword16 t7 = (sword16)(u[i * MLKEM_N + j + 7] + e1[j + 7]);
2152            u[i * MLKEM_N + j + 0] = MLKEM_BARRETT_RED(t0);
2153            u[i * MLKEM_N + j + 1] = MLKEM_BARRETT_RED(t1);
2154            u[i * MLKEM_N + j + 2] = MLKEM_BARRETT_RED(t2);
2155            u[i * MLKEM_N + j + 3] = MLKEM_BARRETT_RED(t3);
2156            u[i * MLKEM_N + j + 4] = MLKEM_BARRETT_RED(t4);
2157            u[i * MLKEM_N + j + 5] = MLKEM_BARRETT_RED(t5);
2158            u[i * MLKEM_N + j + 6] = MLKEM_BARRETT_RED(t6);
2159            u[i * MLKEM_N + j + 7] = MLKEM_BARRETT_RED(t7);
2160        }
2161#endif
2162    }
2163
2164    /* Multiply public key by y into v polynomial. */
2165    mlkem_pointwise_acc_mont(v, pub, y, (unsigned int)k);
2166    /* Inverse transform v. */
2167    mlkem_invntt(v);
2168
2169    mlkem_from_msg(m, msg);
2170
2171    /* Generate noise using PRF. */
2172    coins[WC_ML_KEM_SYM_SZ] = (byte)(2 * k);
2173    ret = mlkem_get_noise_eta2_c(prf, e2, coins);
2174    if (ret == 0) {
2175        /* Add errors and message to v and reduce. */
2176    #if defined(WOLFSSL_MLKEM_SMALL) || defined(WOLFSSL_MLKEM_NO_LARGE_CODE)
2177        for (i = 0; i < MLKEM_N; ++i) {
2178            sword16 t = (sword16)(v[i] + e2[i] + m[i]);
2179            v[i] = MLKEM_BARRETT_RED(t);
2180        }
2181    #else
2182        for (i = 0; i < MLKEM_N; i += 8) {
2183            sword16 t0 = (sword16)(v[i + 0] + e2[i + 0] + m[i + 0]);
2184            sword16 t1 = (sword16)(v[i + 1] + e2[i + 1] + m[i + 1]);
2185            sword16 t2 = (sword16)(v[i + 2] + e2[i + 2] + m[i + 2]);
2186            sword16 t3 = (sword16)(v[i + 3] + e2[i + 3] + m[i + 3]);
2187            sword16 t4 = (sword16)(v[i + 4] + e2[i + 4] + m[i + 4]);
2188            sword16 t5 = (sword16)(v[i + 5] + e2[i + 5] + m[i + 5]);
2189            sword16 t6 = (sword16)(v[i + 6] + e2[i + 6] + m[i + 6]);
2190            sword16 t7 = (sword16)(v[i + 7] + e2[i + 7] + m[i + 7]);
2191            v[i + 0] = MLKEM_BARRETT_RED(t0);
2192            v[i + 1] = MLKEM_BARRETT_RED(t1);
2193            v[i + 2] = MLKEM_BARRETT_RED(t2);
2194            v[i + 3] = MLKEM_BARRETT_RED(t3);
2195            v[i + 4] = MLKEM_BARRETT_RED(t4);
2196            v[i + 5] = MLKEM_BARRETT_RED(t5);
2197            v[i + 6] = MLKEM_BARRETT_RED(t6);
2198            v[i + 7] = MLKEM_BARRETT_RED(t7);
2199        }
2200    #endif
2201    }
2202
2203    return ret;
2204}
2205#endif
2206#endif /* !WOLFSSL_MLKEM_NO_ENCAPSULATE || !WOLFSSL_MLKEM_NO_DECAPSULATE */
2207
2208#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE
2209
2210/* Decapsulate message.
2211 *
2212 * FIPS 203, Algorithm 15: K-PKE.Decrypt(dk_PKE,c)
2213 * Uses the decryption key to decrypt a ciphertext.
2214 *   ...
2215 *   6: w <- v' - InvNTT(s_hat_trans o NTT(u'))
2216 *   ...
2217 *
2218 * @param  [in]       s  Private key vector of polynomials.
2219 * @param  [out]      w  Message polynomial.
2220 * @param  [in, out]  u  Vector of polynomials containing error.
2221 * @param  [in]       v  Encapsulated message polynomial.
2222 * @param  [in]       k  Number of polynomials in vector.
2223 */
2224static void mlkem_decapsulate_c(const sword16* s, sword16* w, sword16* u,
2225    const sword16* v, int k)
2226{
2227    int i;
2228
2229    /* Transform u. All of result used in calculation of w.
2230     * Step 6: ... NTT(u') */
2231    for (i = 0; i < k; ++i) {
2232        mlkem_ntt(u + i * MLKEM_N);
2233    }
2234
2235    /* Multiply private key by u into w polynomial.
2236     * Step 6: ... s_hat_trans o NTT(u') */
2237    mlkem_pointwise_acc_mont(w, s, u, (unsigned int)k);
2238    /* Inverse transform w.
2239     * Step 6: ... InvNTT(s_hat_trans o NTT(u')) */
2240    mlkem_invntt(w);
2241    /* Subtract errors (in w) out of v and reduce into w.
2242     * Step 6: w <- v' - InvNTT(s_hat_trans o NTT(u')) */
2243    for (i = 0; i < MLKEM_N; ++i) {
2244        sword16 t = (sword16)(v[i] - w[i]);
2245        w[i] = MLKEM_BARRETT_RED(t);
2246    }
2247}
2248
2249/* Decapsulate message.
2250 *
2251 * FIPS 203, Algorithm 15: K-PKE.Decrypt(dk_PKE,c)
2252 * Uses the decryption key to decrypt a ciphertext.
2253 *   ...
2254 *   6: w <- v' - InvNTT(s_hat_trans o NTT(u'))
2255 *   ...
2256 *
2257 * @param  [in]       s  Private key vector of polynomials.
2258 * @param  [out]      w  Message polynomial.
2259 * @param  [in, out]  u  Vector of polynomials containing error.
2260 * @param  [in]       v  Encapsulated message polynomial.
2261 * @param  [in]       k  Number of polynomials in vector.
2262 */
2263void mlkem_decapsulate(const sword16* s, sword16* w, sword16* u,
2264    const sword16* v, int k)
2265{
2266#ifdef USE_INTEL_SPEEDUP
2267    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
2268        mlkem_decapsulate_avx2(s, w, u, v, k);
2269        RESTORE_VECTOR_REGISTERS();
2270    }
2271    else
2272#endif
2273    {
2274        mlkem_decapsulate_c(s, w, u, v, k);
2275    }
2276}
2277
2278#endif /* !WOLFSSL_MLKEM_NO_DECAPSULATE */
2279#endif
2280
2281/******************************************************************************/
2282
2283#if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
2284#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
2285/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
2286 *
2287 * Seed used with XOF to generate random bytes.
2288 *
2289 * @param  [out]  a           Matrix of uniform integers.
2290 * @param  [in]   seed        Bytes to seed XOF generation.
2291 * @param  [in]   transposed  Whether A or A^T is generated.
2292 * @return  0 on success.
2293 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
2294 *          WOLFSSL_SMALL_STACK is defined.
2295 */
2296static int mlkem_gen_matrix_k2_avx2(sword16* a, byte* seed, int transposed)
2297{
2298    int i;
2299#ifdef WOLFSSL_SMALL_STACK
2300    byte *rand = NULL;
2301    word64 *state = NULL;
2302#else
2303    byte rand[4 * GEN_MATRIX_SIZE + 4];
2304    word64 state[25 * 4];
2305#endif
2306    unsigned int ctr0;
2307    unsigned int ctr1;
2308    unsigned int ctr2;
2309    unsigned int ctr3;
2310    byte* p;
2311
2312#ifdef WOLFSSL_SMALL_STACK
2313    rand = (byte*)XMALLOC(4 * GEN_MATRIX_SIZE + 4, NULL,
2314                          DYNAMIC_TYPE_TMP_BUFFER);
2315    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
2316                          DYNAMIC_TYPE_TMP_BUFFER);
2317    if ((rand == NULL) || (state == NULL)) {
2318        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2319        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2320        return MEMORY_E;
2321    }
2322#endif
2323
2324    /* Loading 64 bits, only using 48 bits. Loading 4 bytes more than used. */
2325    rand[4 * GEN_MATRIX_SIZE + 0] = 0xff;
2326    rand[4 * GEN_MATRIX_SIZE + 1] = 0xff;
2327    rand[4 * GEN_MATRIX_SIZE + 2] = 0xff;
2328    rand[4 * GEN_MATRIX_SIZE + 3] = 0xff;
2329
2330    if (!transposed) {
2331        state[4*4 + 0] = 0x1f0000 + 0x000;
2332        state[4*4 + 1] = 0x1f0000 + 0x001;
2333        state[4*4 + 2] = 0x1f0000 + 0x100;
2334        state[4*4 + 3] = 0x1f0000 + 0x101;
2335    }
2336    else {
2337        state[4*4 + 0] = 0x1f0000 + 0x000;
2338        state[4*4 + 1] = 0x1f0000 + 0x100;
2339        state[4*4 + 2] = 0x1f0000 + 0x001;
2340        state[4*4 + 3] = 0x1f0000 + 0x101;
2341    }
2342
2343    sha3_128_blocksx4_seed_avx2(state, seed);
2344    mlkem_redistribute_21_rand_avx2(state, rand + 0 * GEN_MATRIX_SIZE,
2345        rand + 1 * GEN_MATRIX_SIZE, rand + 2 * GEN_MATRIX_SIZE,
2346        rand + 3 * GEN_MATRIX_SIZE);
2347    for (i = SHA3_128_BYTES; i < GEN_MATRIX_SIZE; i += SHA3_128_BYTES) {
2348        sha3_blocksx4_avx2(state);
2349        mlkem_redistribute_21_rand_avx2(state, rand + i + 0 * GEN_MATRIX_SIZE,
2350            rand + i + 1 * GEN_MATRIX_SIZE, rand + i + 2 * GEN_MATRIX_SIZE,
2351            rand + i + 3 * GEN_MATRIX_SIZE);
2352    }
2353
2354    /* Sample random bytes to create a polynomial. */
2355    p = rand;
2356    ctr0 = mlkem_rej_uniform_n_avx2(a + 0 * MLKEM_N, MLKEM_N, p,
2357        GEN_MATRIX_SIZE);
2358    p += GEN_MATRIX_SIZE;
2359    ctr1 = mlkem_rej_uniform_n_avx2(a + 1 * MLKEM_N, MLKEM_N, p,
2360        GEN_MATRIX_SIZE);
2361    p += GEN_MATRIX_SIZE;
2362    ctr2 = mlkem_rej_uniform_n_avx2(a + 2 * MLKEM_N, MLKEM_N, p,
2363        GEN_MATRIX_SIZE);
2364    p += GEN_MATRIX_SIZE;
2365    ctr3 = mlkem_rej_uniform_n_avx2(a + 3 * MLKEM_N, MLKEM_N, p,
2366        GEN_MATRIX_SIZE);
2367    /* Create more blocks if too many rejected. */
2368    while ((ctr0 < MLKEM_N) || (ctr1 < MLKEM_N) || (ctr2 < MLKEM_N) ||
2369           (ctr3 < MLKEM_N)) {
2370        sha3_blocksx4_avx2(state);
2371        mlkem_redistribute_21_rand_avx2(state, rand + 0 * GEN_MATRIX_SIZE,
2372            rand + 1 * GEN_MATRIX_SIZE, rand + 2 * GEN_MATRIX_SIZE,
2373            rand + 3 * GEN_MATRIX_SIZE);
2374
2375        p = rand;
2376        ctr0 += mlkem_rej_uniform_avx2(a + 0 * MLKEM_N + ctr0, MLKEM_N - ctr0,
2377            p, XOF_BLOCK_SIZE);
2378        p += GEN_MATRIX_SIZE;
2379        ctr1 += mlkem_rej_uniform_avx2(a + 1 * MLKEM_N + ctr1, MLKEM_N - ctr1,
2380            p, XOF_BLOCK_SIZE);
2381        p += GEN_MATRIX_SIZE;
2382        ctr2 += mlkem_rej_uniform_avx2(a + 2 * MLKEM_N + ctr2, MLKEM_N - ctr2,
2383            p, XOF_BLOCK_SIZE);
2384        p += GEN_MATRIX_SIZE;
2385        ctr3 += mlkem_rej_uniform_avx2(a + 3 * MLKEM_N + ctr3, MLKEM_N - ctr3,
2386            p, XOF_BLOCK_SIZE);
2387    }
2388
2389    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2390    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2391
2392    return 0;
2393}
2394#endif
2395
2396#if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
2397/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
2398 *
2399 * Seed used with XOF to generate random bytes.
2400 *
2401 * @param  [out]  a           Matrix of uniform integers.
2402 * @param  [in]   seed        Bytes to seed XOF generation.
2403 * @param  [in]   transposed  Whether A or A^T is generated.
2404 * @return  0 on success.
2405 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
2406 *          WOLFSSL_SMALL_STACK is defined.
2407 */
2408static int mlkem_gen_matrix_k3_avx2(sword16* a, byte* seed, int transposed)
2409{
2410    int i;
2411    int k;
2412#ifdef WOLFSSL_SMALL_STACK
2413    byte *rand = NULL;
2414    word64 *state = NULL;
2415#else
2416    byte rand[4 * GEN_MATRIX_SIZE + 4];
2417    word64 state[25 * 4];
2418#endif
2419    unsigned int ctr0;
2420    unsigned int ctr1;
2421    unsigned int ctr2;
2422    unsigned int ctr3;
2423    byte* p;
2424
2425#ifdef WOLFSSL_SMALL_STACK
2426    rand = (byte*)XMALLOC(4 * GEN_MATRIX_SIZE + 4, NULL,
2427                          DYNAMIC_TYPE_TMP_BUFFER);
2428    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
2429                          DYNAMIC_TYPE_TMP_BUFFER);
2430    if ((rand == NULL) || (state == NULL)) {
2431        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2432        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2433        return MEMORY_E;
2434    }
2435#endif
2436
2437    /* Loading 64 bits, only using 48 bits. Loading 4 bytes more than used. */
2438    rand[4 * GEN_MATRIX_SIZE + 0] = 0xff;
2439    rand[4 * GEN_MATRIX_SIZE + 1] = 0xff;
2440    rand[4 * GEN_MATRIX_SIZE + 2] = 0xff;
2441    rand[4 * GEN_MATRIX_SIZE + 3] = 0xff;
2442
2443    for (k = 0; k < 2; k++) {
2444        for (i = 0; i < 4; i++) {
2445            if (!transposed) {
2446                state[4*4 + i] = (word32)(0x1f0000 + (((k*4+i)/3) << 8) +
2447                                          ((k*4+i)%3));
2448            }
2449            else {
2450                state[4*4 + i] = (word32)(0x1f0000 + (((k*4+i)%3) << 8) +
2451                                          ((k*4+i)/3));
2452
2453            }
2454        }
2455
2456        sha3_128_blocksx4_seed_avx2(state, seed);
2457        mlkem_redistribute_21_rand_avx2(state,
2458            rand + 0 * GEN_MATRIX_SIZE, rand + 1 * GEN_MATRIX_SIZE,
2459            rand + 2 * GEN_MATRIX_SIZE, rand + 3 * GEN_MATRIX_SIZE);
2460        for (i = SHA3_128_BYTES; i < GEN_MATRIX_SIZE; i += SHA3_128_BYTES) {
2461            sha3_blocksx4_avx2(state);
2462            mlkem_redistribute_21_rand_avx2(state,
2463                rand + i + 0 * GEN_MATRIX_SIZE, rand + i + 1 * GEN_MATRIX_SIZE,
2464                rand + i + 2 * GEN_MATRIX_SIZE, rand + i + 3 * GEN_MATRIX_SIZE);
2465        }
2466
2467        /* Sample random bytes to create a polynomial. */
2468        p = rand;
2469        ctr0 = mlkem_rej_uniform_n_avx2(a + 0 * MLKEM_N, MLKEM_N, p,
2470            GEN_MATRIX_SIZE);
2471        p += GEN_MATRIX_SIZE;
2472        ctr1 = mlkem_rej_uniform_n_avx2(a + 1 * MLKEM_N, MLKEM_N, p,
2473            GEN_MATRIX_SIZE);
2474        p += GEN_MATRIX_SIZE;
2475        ctr2 = mlkem_rej_uniform_n_avx2(a + 2 * MLKEM_N, MLKEM_N, p,
2476            GEN_MATRIX_SIZE);
2477        p += GEN_MATRIX_SIZE;
2478        ctr3 = mlkem_rej_uniform_n_avx2(a + 3 * MLKEM_N, MLKEM_N, p,
2479            GEN_MATRIX_SIZE);
2480        /* Create more blocks if too many rejected. */
2481        while ((ctr0 < MLKEM_N) || (ctr1 < MLKEM_N) || (ctr2 < MLKEM_N) ||
2482               (ctr3 < MLKEM_N)) {
2483            sha3_blocksx4_avx2(state);
2484            mlkem_redistribute_21_rand_avx2(state, rand + 0 * GEN_MATRIX_SIZE,
2485                rand + 1 * GEN_MATRIX_SIZE, rand + 2 * GEN_MATRIX_SIZE,
2486                rand + 3 * GEN_MATRIX_SIZE);
2487
2488            p = rand;
2489            ctr0 += mlkem_rej_uniform_avx2(a + 0 * MLKEM_N + ctr0,
2490                MLKEM_N - ctr0, p, XOF_BLOCK_SIZE);
2491            p += GEN_MATRIX_SIZE;
2492            ctr1 += mlkem_rej_uniform_avx2(a + 1 * MLKEM_N + ctr1,
2493                MLKEM_N - ctr1, p, XOF_BLOCK_SIZE);
2494            p += GEN_MATRIX_SIZE;
2495            ctr2 += mlkem_rej_uniform_avx2(a + 2 * MLKEM_N + ctr2,
2496                MLKEM_N - ctr2, p, XOF_BLOCK_SIZE);
2497            p += GEN_MATRIX_SIZE;
2498            ctr3 += mlkem_rej_uniform_avx2(a + 3 * MLKEM_N + ctr3,
2499                MLKEM_N - ctr3, p, XOF_BLOCK_SIZE);
2500        }
2501
2502        a += 4 * MLKEM_N;
2503    }
2504
2505    readUnalignedWords64(state, seed, 4);
2506    /* Transposed value same as not. */
2507    state[4] = 0x1f0000 + (2 << 8) + 2;
2508    XMEMSET(state + 5, 0, sizeof(*state) * (25 - 5));
2509    state[20] = W64LIT(0x8000000000000000);
2510    for (i = 0; i < GEN_MATRIX_SIZE; i += SHA3_128_BYTES) {
2511#ifndef WC_SHA3_NO_ASM
2512        if (IS_INTEL_BMI2(cpuid_flags)) {
2513            sha3_block_bmi2(state);
2514        }
2515        else if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0))
2516        {
2517            sha3_block_avx2(state);
2518            RESTORE_VECTOR_REGISTERS();
2519        }
2520        else
2521#endif /* !WC_SHA3_NO_ASM */
2522        {
2523            BlockSha3(state);
2524        }
2525        XMEMCPY(rand + i, state, SHA3_128_BYTES);
2526    }
2527    ctr0 = mlkem_rej_uniform_n_avx2(a, MLKEM_N, rand, GEN_MATRIX_SIZE);
2528    while (ctr0 < MLKEM_N) {
2529#ifndef WC_SHA3_NO_ASM
2530        if (IS_INTEL_BMI2(cpuid_flags)) {
2531            sha3_block_bmi2(state);
2532        }
2533        else if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0))
2534        {
2535            sha3_block_avx2(state);
2536            RESTORE_VECTOR_REGISTERS();
2537        }
2538        else
2539#endif /* !WC_SHA3_NO_ASM */
2540        {
2541            BlockSha3(state);
2542        }
2543        XMEMCPY(rand, state, SHA3_128_BYTES);
2544        ctr0 += mlkem_rej_uniform_avx2(a + ctr0, MLKEM_N - ctr0, rand,
2545            XOF_BLOCK_SIZE);
2546    }
2547
2548    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2549    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2550
2551    return 0;
2552}
2553#endif
2554#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
2555/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
2556 *
2557 * Seed used with XOF to generate random bytes.
2558 *
2559 * @param  [out]  a           Matrix of uniform integers.
2560 * @param  [in]   seed        Bytes to seed XOF generation.
2561 * @param  [in]   transposed  Whether A or A^T is generated.
2562 * @return  0 on success.
2563 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
2564 *          WOLFSSL_SMALL_STACK is defined.
2565 */
2566static int mlkem_gen_matrix_k4_avx2(sword16* a, byte* seed, int transposed)
2567{
2568    int i;
2569    int k;
2570#ifdef WOLFSSL_SMALL_STACK
2571    byte *rand = NULL;
2572    word64 *state = NULL;
2573#else
2574    byte rand[4 * GEN_MATRIX_SIZE + 4];
2575    word64 state[25 * 4];
2576#endif
2577    unsigned int ctr0;
2578    unsigned int ctr1;
2579    unsigned int ctr2;
2580    unsigned int ctr3;
2581    byte* p;
2582
2583#ifdef WOLFSSL_SMALL_STACK
2584    rand = (byte*)XMALLOC(4 * GEN_MATRIX_SIZE + 4, NULL,
2585                          DYNAMIC_TYPE_TMP_BUFFER);
2586    state = (word64*)XMALLOC(sizeof(word64) * 25 * 4, NULL,
2587                          DYNAMIC_TYPE_TMP_BUFFER);
2588    if ((rand == NULL) || (state == NULL)) {
2589        XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2590        XFREE(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2591        return MEMORY_E;
2592    }
2593#endif
2594
2595    /* Loading 64 bits, only using 48 bits. Loading 4 bytes more than used. */
2596    rand[4 * GEN_MATRIX_SIZE + 0] = 0xff;
2597    rand[4 * GEN_MATRIX_SIZE + 1] = 0xff;
2598    rand[4 * GEN_MATRIX_SIZE + 2] = 0xff;
2599    rand[4 * GEN_MATRIX_SIZE + 3] = 0xff;
2600
2601    for (k = 0; k < 4; k++) {
2602        for (i = 0; i < 4; i++) {
2603            if (!transposed) {
2604                state[4*4 + i] = (word32)(0x1f0000 + (k << 8) + i);
2605            }
2606            else {
2607                state[4*4 + i] = (word32)(0x1f0000 + (i << 8) + k);
2608            }
2609        }
2610
2611        sha3_128_blocksx4_seed_avx2(state, seed);
2612        mlkem_redistribute_21_rand_avx2(state,
2613            rand + 0 * GEN_MATRIX_SIZE, rand + 1 * GEN_MATRIX_SIZE,
2614            rand + 2 * GEN_MATRIX_SIZE, rand + 3 * GEN_MATRIX_SIZE);
2615        for (i = SHA3_128_BYTES; i < GEN_MATRIX_SIZE; i += SHA3_128_BYTES) {
2616            sha3_blocksx4_avx2(state);
2617            mlkem_redistribute_21_rand_avx2(state,
2618                rand + i + 0 * GEN_MATRIX_SIZE, rand + i + 1 * GEN_MATRIX_SIZE,
2619                rand + i + 2 * GEN_MATRIX_SIZE, rand + i + 3 * GEN_MATRIX_SIZE);
2620        }
2621
2622        /* Sample random bytes to create a polynomial. */
2623        p = rand;
2624        ctr0 = mlkem_rej_uniform_n_avx2(a + 0 * MLKEM_N, MLKEM_N, p,
2625            GEN_MATRIX_SIZE);
2626        p += GEN_MATRIX_SIZE;
2627        ctr1 = mlkem_rej_uniform_n_avx2(a + 1 * MLKEM_N, MLKEM_N, p,
2628            GEN_MATRIX_SIZE);
2629        p += GEN_MATRIX_SIZE;
2630        ctr2 = mlkem_rej_uniform_n_avx2(a + 2 * MLKEM_N, MLKEM_N, p,
2631            GEN_MATRIX_SIZE);
2632        p += GEN_MATRIX_SIZE;
2633        ctr3 = mlkem_rej_uniform_n_avx2(a + 3 * MLKEM_N, MLKEM_N, p,
2634            GEN_MATRIX_SIZE);
2635        /* Create more blocks if too many rejected. */
2636        while ((ctr0 < MLKEM_N) || (ctr1 < MLKEM_N) || (ctr2 < MLKEM_N) ||
2637               (ctr3 < MLKEM_N)) {
2638            sha3_blocksx4_avx2(state);
2639            mlkem_redistribute_21_rand_avx2(state, rand + 0 * GEN_MATRIX_SIZE,
2640                rand + 1 * GEN_MATRIX_SIZE, rand + 2 * GEN_MATRIX_SIZE,
2641                rand + 3 * GEN_MATRIX_SIZE);
2642
2643            p = rand;
2644            ctr0 += mlkem_rej_uniform_avx2(a + 0 * MLKEM_N + ctr0,
2645                MLKEM_N - ctr0, p, XOF_BLOCK_SIZE);
2646            p += GEN_MATRIX_SIZE;
2647            ctr1 += mlkem_rej_uniform_avx2(a + 1 * MLKEM_N + ctr1,
2648                MLKEM_N - ctr1, p, XOF_BLOCK_SIZE);
2649            p += GEN_MATRIX_SIZE;
2650            ctr2 += mlkem_rej_uniform_avx2(a + 2 * MLKEM_N + ctr2,
2651                MLKEM_N - ctr2, p, XOF_BLOCK_SIZE);
2652            p += GEN_MATRIX_SIZE;
2653            ctr3 += mlkem_rej_uniform_avx2(a + 3 * MLKEM_N + ctr3,
2654                MLKEM_N - ctr3, p, XOF_BLOCK_SIZE);
2655        }
2656
2657        a += 4 * MLKEM_N;
2658    }
2659
2660    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2661    WC_FREE_VAR_EX(state, NULL, DYNAMIC_TYPE_TMP_BUFFER);
2662
2663    return 0;
2664}
2665#endif /* WOLFSSL_KYBER1024 || WOLFSSL_WC_ML_KEM_1024 */
2666#elif defined(WOLFSSL_ARMASM) && defined(__aarch64__)
2667#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
2668/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
2669 *
2670 * Seed used with XOF to generate random bytes.
2671 *
2672 * @param  [out]  a           Matrix of uniform integers.
2673 * @param  [in]   seed        Bytes to seed XOF generation.
2674 * @param  [in]   transposed  Whether A or A^T is generated.
2675 * @return  0 on success.
2676 */
2677static int mlkem_gen_matrix_k2_aarch64(sword16* a, byte* seed, int transposed)
2678{
2679    word64 state[3 * 25];
2680    word64* st = (word64*)state;
2681    unsigned int ctr0;
2682    unsigned int ctr1;
2683    unsigned int ctr2;
2684    byte* p;
2685
2686    if (!transposed) {
2687        state[0*25 + 4] = 0x1f0000 + (0 << 8) + 0;
2688        state[1*25 + 4] = 0x1f0000 + (0 << 8) + 1;
2689        state[2*25 + 4] = 0x1f0000 + (1 << 8) + 0;
2690    }
2691    else {
2692        state[0*25 + 4] = 0x1f0000 + (0 << 8) + 0;
2693        state[1*25 + 4] = 0x1f0000 + (1 << 8) + 0;
2694        state[2*25 + 4] = 0x1f0000 + (0 << 8) + 1;
2695    }
2696
2697    mlkem_shake128_blocksx3_seed_neon(state, seed);
2698    /* Sample random bytes to create a polynomial. */
2699    p = (byte*)st;
2700    ctr0 = mlkem_rej_uniform_neon(a + 0 * MLKEM_N, MLKEM_N, p, XOF_BLOCK_SIZE);
2701    p += 25 * 8;
2702    ctr1 = mlkem_rej_uniform_neon(a + 1 * MLKEM_N, MLKEM_N, p, XOF_BLOCK_SIZE);
2703    p += 25 * 8;
2704    ctr2 = mlkem_rej_uniform_neon(a + 2 * MLKEM_N, MLKEM_N, p, XOF_BLOCK_SIZE);
2705    while ((ctr0 < MLKEM_N) || (ctr1 < MLKEM_N) || (ctr2 < MLKEM_N)) {
2706        mlkem_sha3_blocksx3_neon(st);
2707
2708        p = (byte*)st;
2709        ctr0 += mlkem_rej_uniform_neon(a + 0 * MLKEM_N + ctr0, MLKEM_N - ctr0,
2710            p, XOF_BLOCK_SIZE);
2711        p += 25 * 8;
2712        ctr1 += mlkem_rej_uniform_neon(a + 1 * MLKEM_N + ctr1, MLKEM_N - ctr1,
2713            p, XOF_BLOCK_SIZE);
2714        p += 25 * 8;
2715        ctr2 += mlkem_rej_uniform_neon(a + 2 * MLKEM_N + ctr2, MLKEM_N - ctr2,
2716            p, XOF_BLOCK_SIZE);
2717    }
2718
2719    a += 3 * MLKEM_N;
2720
2721    readUnalignedWords64(state, seed, 4);
2722    /* Transposed value same as not. */
2723    state[4] = 0x1f0000 + (1 << 8) + 1;
2724    XMEMSET(state + 5, 0, sizeof(*state) * (25 - 5));
2725    state[20] = W64LIT(0x8000000000000000);
2726    BlockSha3(state);
2727    p = (byte*)state;
2728    ctr0 = mlkem_rej_uniform_neon(a, MLKEM_N, p, XOF_BLOCK_SIZE);
2729    while (ctr0 < MLKEM_N) {
2730        BlockSha3(state);
2731        ctr0 += mlkem_rej_uniform_neon(a + ctr0, MLKEM_N - ctr0, p,
2732            XOF_BLOCK_SIZE);
2733    }
2734
2735    return 0;
2736}
2737#endif
2738
2739#if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
2740/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
2741 *
2742 * Seed used with XOF to generate random bytes.
2743 *
2744 * @param  [out]  a           Matrix of uniform integers.
2745 * @param  [in]   seed        Bytes to seed XOF generation.
2746 * @param  [in]   transposed  Whether A or A^T is generated.
2747 * @return  0 on success.
2748 */
2749static int mlkem_gen_matrix_k3_aarch64(sword16* a, byte* seed, int transposed)
2750{
2751    int i;
2752    int k;
2753    word64 state[3 * 25];
2754    word64* st = (word64*)state;
2755    unsigned int ctr0;
2756    unsigned int ctr1;
2757    unsigned int ctr2;
2758    byte* p;
2759
2760    for (k = 0; k < 3; k++) {
2761        for (i = 0; i < 3; i++) {
2762            if (!transposed) {
2763                state[i*25 + 4] = 0x1f0000 + ((k << 8) + i);
2764            }
2765            else {
2766                state[i*25 + 4] = 0x1f0000 + ((i << 8) + k);
2767            }
2768        }
2769
2770        mlkem_shake128_blocksx3_seed_neon(state, seed);
2771        /* Sample random bytes to create a polynomial. */
2772        p = (byte*)st;
2773        ctr0 = mlkem_rej_uniform_neon(a + 0 * MLKEM_N, MLKEM_N, p,
2774            XOF_BLOCK_SIZE);
2775        p += 25 * 8;
2776        ctr1 = mlkem_rej_uniform_neon(a + 1 * MLKEM_N, MLKEM_N, p,
2777            XOF_BLOCK_SIZE);
2778        p += 25 * 8;
2779        ctr2 = mlkem_rej_uniform_neon(a + 2 * MLKEM_N, MLKEM_N, p,
2780            XOF_BLOCK_SIZE);
2781        /* Create more blocks if too many rejected. */
2782        while ((ctr0 < MLKEM_N) || (ctr1 < MLKEM_N) || (ctr2 < MLKEM_N)) {
2783            mlkem_sha3_blocksx3_neon(st);
2784
2785            p = (byte*)st;
2786            ctr0 += mlkem_rej_uniform_neon(a + 0 * MLKEM_N + ctr0,
2787                MLKEM_N - ctr0, p, XOF_BLOCK_SIZE);
2788            p += 25 * 8;
2789            ctr1 += mlkem_rej_uniform_neon(a + 1 * MLKEM_N + ctr1,
2790                MLKEM_N - ctr1, p, XOF_BLOCK_SIZE);
2791            p += 25 * 8;
2792            ctr2 += mlkem_rej_uniform_neon(a + 2 * MLKEM_N + ctr2,
2793                MLKEM_N - ctr2, p, XOF_BLOCK_SIZE);
2794        }
2795
2796        a += 3 * MLKEM_N;
2797    }
2798
2799    return 0;
2800}
2801#endif
2802
2803#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
2804/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
2805 *
2806 * Seed used with XOF to generate random bytes.
2807 *
2808 * @param  [out]  a           Matrix of uniform integers.
2809 * @param  [in]   seed        Bytes to seed XOF generation.
2810 * @param  [in]   transposed  Whether A or A^T is generated.
2811 * @return  0 on success.
2812 */
2813static int mlkem_gen_matrix_k4_aarch64(sword16* a, byte* seed, int transposed)
2814{
2815    int i;
2816    int k;
2817    word64 state[3 * 25];
2818    word64* st = (word64*)state;
2819    unsigned int ctr0;
2820    unsigned int ctr1;
2821    unsigned int ctr2;
2822    byte* p;
2823
2824    for (k = 0; k < 5; k++) {
2825        for (i = 0; i < 3; i++) {
2826            byte bi = ((k * 3) + i) / 4;
2827            byte bj = ((k * 3) + i) % 4;
2828            if (!transposed) {
2829                state[i*25 + 4] = 0x1f0000 + (bi << 8) + bj;
2830            }
2831            else {
2832                state[i*25 + 4] = 0x1f0000 + (bj << 8) + bi;
2833            }
2834        }
2835
2836        mlkem_shake128_blocksx3_seed_neon(state, seed);
2837        /* Sample random bytes to create a polynomial. */
2838        p = (byte*)st;
2839        ctr0 = mlkem_rej_uniform_neon(a + 0 * MLKEM_N, MLKEM_N, p,
2840            XOF_BLOCK_SIZE);
2841        p += 25 * 8;
2842        ctr1 = mlkem_rej_uniform_neon(a + 1 * MLKEM_N, MLKEM_N, p,
2843            XOF_BLOCK_SIZE);
2844        p += 25 * 8;
2845        ctr2 = mlkem_rej_uniform_neon(a + 2 * MLKEM_N, MLKEM_N, p,
2846            XOF_BLOCK_SIZE);
2847        /* Create more blocks if too many rejected. */
2848        while ((ctr0 < MLKEM_N) || (ctr1 < MLKEM_N) || (ctr2 < MLKEM_N)) {
2849            mlkem_sha3_blocksx3_neon(st);
2850
2851            p = (byte*)st;
2852            ctr0 += mlkem_rej_uniform_neon(a + 0 * MLKEM_N + ctr0,
2853                MLKEM_N - ctr0, p, XOF_BLOCK_SIZE);
2854            p += 25 * 8;
2855            ctr1 += mlkem_rej_uniform_neon(a + 1 * MLKEM_N + ctr1,
2856                MLKEM_N - ctr1, p, XOF_BLOCK_SIZE);
2857            p += 25 * 8;
2858            ctr2 += mlkem_rej_uniform_neon(a + 2 * MLKEM_N + ctr2,
2859                MLKEM_N - ctr2, p, XOF_BLOCK_SIZE);
2860        }
2861
2862        a += 3 * MLKEM_N;
2863    }
2864
2865    readUnalignedWords64(state, seed, 4);
2866    /* Transposed value same as not. */
2867    state[4] = 0x1f0000 + (3 << 8) + 3;
2868    XMEMSET(state + 5, 0, sizeof(*state) * (25 - 5));
2869    state[20] = W64LIT(0x8000000000000000);
2870    BlockSha3(state);
2871    p = (byte*)state;
2872    ctr0 = mlkem_rej_uniform_neon(a, MLKEM_N, p, XOF_BLOCK_SIZE);
2873    while (ctr0 < MLKEM_N) {
2874        BlockSha3(state);
2875        ctr0 += mlkem_rej_uniform_neon(a + ctr0, MLKEM_N - ctr0, p,
2876            XOF_BLOCK_SIZE);
2877    }
2878
2879    return 0;
2880}
2881#endif
2882#endif /* USE_INTEL_SPEEDUP */
2883
2884#if !(defined(WOLFSSL_ARMASM) && defined(__aarch64__))
2885/* Absorb the seed data for squeezing out pseudo-random data.
2886 *
2887 * FIPS 203, Section 4.1:
2888 * 1. XOF.init() = SHAKE128.Init().
2889 * 2. XOF.Absorb(ctx,str) = SHAKE128.Absorb(ctx,str).
2890 *
2891 * @param  [in, out]  shake128  SHAKE-128 object.
2892 * @param  [in]       seed      Data to absorb.
2893 * @param  [in]       len       Length of data to absorb in bytes.
2894 * @return  0 on success always.
2895 */
2896static int mlkem_xof_absorb(wc_Shake* shake128, const byte* seed, int len)
2897{
2898    int ret;
2899
2900    ret = wc_InitShake128(shake128, NULL, INVALID_DEVID);
2901    if (ret == 0) {
2902        ret = wc_Shake128_Absorb(shake128, seed, (word32)len);
2903    }
2904
2905    return ret;
2906}
2907
2908/* Squeeze the state to produce pseudo-random data.
2909 *
2910 * FIPS 203, Section 4.1:
2911 * 3. XOF.Squeeze(ctx,l) = SHAKE128.Squeeze(ctx,8.l).
2912 *
2913 * @param  [in, out]  shake128  SHAKE-128 object.
2914 * @param  [out]      out       Buffer to write to.
2915 * @param  [in]       blocks    Number of blocks to write.
2916 * @return  0 on success always.
2917 */
2918static int mlkem_xof_squeezeblocks(wc_Shake* shake128, byte* out, int blocks)
2919{
2920    return wc_Shake128_SqueezeBlocks(shake128, out, (word32)blocks);
2921}
2922#endif
2923
2924/* New/Initialize SHA-3 object.
2925 *
2926 * FIPS 203, Section 4.1:
2927 * H(s) := SHA3-256(s)
2928 *
2929 * @param  [in, out]  hash    SHA-3 object.
2930 * @param  [in]       heap    Dynamic memory allocator hint.
2931 * @param  [in]       devId   Device id.
2932 * @return  0 on success always.
2933 */
2934int mlkem_hash_new(wc_Sha3* hash, void* heap, int devId)
2935{
2936    return wc_InitSha3_256(hash, heap, devId);
2937}
2938
2939/* Free SHA-3 object.
2940 *
2941 * FIPS 203, Section 4.1:
2942 * H(s) := SHA3-256(s)
2943 *
2944 * @param  [in, out]  hash  SHA-3 object.
2945 */
2946void mlkem_hash_free(wc_Sha3* hash)
2947{
2948    wc_Sha3_256_Free(hash);
2949}
2950
2951/* Hash data using SHA3-256 with SHA-3 object.
2952 *
2953 * FIPS 203, Section 4.1:
2954 * H(s) := SHA3-256(s)
2955 *
2956 * @param  [in, out]  hash     SHA-3 object.
2957 * @param  [in]       data     Data to be hashed.
2958 * @param  [in]       dataLen  Length of data in bytes.
2959 * @param  [out]      out      Hash of data.
2960 * @return  0 on success.
2961 */
2962int mlkem_hash256(wc_Sha3* hash, const byte* data, word32 dataLen, byte* out)
2963{
2964    int ret;
2965
2966    /* Process all data. */
2967    ret = wc_Sha3_256_Update(hash, data, dataLen);
2968    if (ret == 0) {
2969        /* Calculate Hash of data passed in and re-initialize. */
2970        ret = wc_Sha3_256_Final(hash, out);
2971    }
2972
2973    return ret;
2974}
2975
2976/* Hash one or two blocks of data using SHA3-512 with SHA-3 object.
2977 *
2978 * FIPS 203, Section 4.1:
2979 * G(s) := SHA3-512(s)
2980 *
2981 * @param  [in, out]  hash      SHA-3 object.
2982 * @param  [in]       data1     First block of data to be hashed.
2983 * @param  [in]       data1Len  Length of first block of data in bytes.
2984 * @param  [in]       data2     Second block of data to be hashed. May be NULL.
2985 * @param  [in]       data2Len  Length of second block of data in bytes.
2986 * @param  [out]      out       Hash of all data.
2987 * @return  0 on success.
2988 */
2989int mlkem_hash512(wc_Sha3* hash, const byte* data1, word32 data1Len,
2990    const byte* data2, word32 data2Len, byte* out)
2991{
2992    int ret;
2993
2994    /* Process first block of data. */
2995    ret = wc_Sha3_512_Update(hash, data1, data1Len);
2996    /* Check if there is a second block of data. */
2997    if ((ret == 0) && (data2 != NULL) && (data2Len > 0)) {
2998        /* Process second block of data. */
2999        ret = wc_Sha3_512_Update(hash, data2, data2Len);
3000    }
3001    if (ret == 0) {
3002        /* Calculate Hash of data passed in and re-initialize. */
3003        ret = wc_Sha3_512_Final(hash, out);
3004    }
3005
3006    return ret;
3007}
3008
3009/* Initialize SHAKE-256 object.
3010 *
3011 * @param  [in, out]  prf  SHAKE-256 object.
3012 */
3013void mlkem_prf_init(wc_Shake* prf)
3014{
3015    wc_InitShake256(prf, NULL, 0);
3016}
3017
3018/* New/Initialize SHAKE-256 object.
3019 *
3020 * FIPS 203, Section 4.1, 4.3:
3021 * PRF_eta(s,b) := SHAKE256(s||b,8.64.eta)
3022 *
3023 * @param  [in, out]  prf    SHAKE-256 object.
3024 * @param  [in]       heap   Dynamic memory allocator hint.
3025 * @param  [in]       devId  Device id.
3026 * @return  0 on success always.
3027 */
3028int mlkem_prf_new(wc_Shake* prf, void* heap, int devId)
3029{
3030    return wc_InitShake256(prf, heap, devId);
3031}
3032
3033/* Free SHAKE-256 object.
3034 *
3035 * FIPS 203, Section 4.1, 4.3:
3036 * PRF_eta(s,b) := SHAKE256(s||b,8.64.eta)
3037 *
3038 * @param  [in, out]  prf  SHAKE-256 object.
3039 */
3040void mlkem_prf_free(wc_Shake* prf)
3041{
3042    wc_Shake256_Free(prf);
3043}
3044
3045#if !(defined(WOLFSSL_ARMASM) && defined(__aarch64__))
3046/* Create pseudo-random data from the key using SHAKE-256.
3047 *
3048 * FIPS 203, Section 4.1, 4.3:
3049 * PRF_eta(s,b) := SHAKE256(s||b,8.64.eta)
3050 *
3051 * @param  [in, out]  shake256  SHAKE-256 object.
3052 * @param  [out]      out       Buffer to write to.
3053 * @param  [in]       outLen    Number of bytes to write.
3054 * @param  [in]       key       Data to derive from. Must be:
3055 *                                WC_ML_KEM_SYM_SZ + 1 bytes in length.
3056 * @return  0 on success always.
3057 */
3058static int mlkem_prf(wc_Shake* shake256, byte* out, unsigned int outLen,
3059    const byte* key)
3060{
3061#ifdef USE_INTEL_SPEEDUP
3062    word64 state[25];
3063
3064    (void)shake256;
3065
3066    /* Put first WC_ML_KEM_SYM_SZ bytes of key into blank state. */
3067    readUnalignedWords64(state, key, WC_ML_KEM_SYM_SZ / sizeof(word64));
3068    /* Last byte in with end of content marker. */
3069    state[WC_ML_KEM_SYM_SZ / 8] = 0x1f00 | key[WC_ML_KEM_SYM_SZ];
3070    /* Set rest of state to 0. */
3071    XMEMSET(state + WC_ML_KEM_SYM_SZ / 8 + 1, 0,
3072        (25 - WC_ML_KEM_SYM_SZ / 8 - 1) * sizeof(word64));
3073    /* ... except for rate marker. */
3074    state[WC_SHA3_256_COUNT - 1] = W64LIT(0x8000000000000000);
3075
3076    /* Generate as much output as is required. */
3077    while (outLen > 0) {
3078        /* Get as much of an output block as is needed. */
3079        unsigned int len = min(outLen, WC_SHA3_256_BLOCK_SIZE);
3080
3081        /* Perform a block operation on the state for next block of output. */
3082#ifndef WC_SHA3_NO_ASM
3083        if (IS_INTEL_BMI2(cpuid_flags)) {
3084            sha3_block_bmi2(state);
3085        }
3086        else if (IS_INTEL_AVX2(cpuid_flags) &&
3087                 (SAVE_VECTOR_REGISTERS2() == 0)) {
3088            sha3_block_avx2(state);
3089            RESTORE_VECTOR_REGISTERS();
3090        }
3091        else
3092#endif /* !WC_SHA3_NO_ASM */
3093        {
3094            BlockSha3(state);
3095        }
3096
3097        /* Copy the state as output. */
3098        XMEMCPY(out, state, len);
3099        /* Update output pointer and length. */
3100        out += len;
3101        outLen -= len;
3102    }
3103
3104    return 0;
3105#else
3106    int ret;
3107
3108    /* Process all data. */
3109    ret = wc_Shake256_Update(shake256, key, WC_ML_KEM_SYM_SZ + 1);
3110    if (ret == 0) {
3111        /* Calculate Hash of data passed in and re-initialize. */
3112        ret = wc_Shake256_Final(shake256, out, outLen);
3113    }
3114
3115    return ret;
3116#endif
3117}
3118#endif
3119
3120#ifdef WOLFSSL_MLKEM_KYBER
3121#ifdef USE_INTEL_SPEEDUP
3122/* Create pseudo-random key from the seed using SHAKE-256.
3123 *
3124 * @param  [in]  seed      Data to derive from.
3125 * @param  [in]  seedLen   Length of data to derive from in bytes.
3126 * @param  [out] out       Buffer to write to.
3127 * @param  [in]  outLen    Number of bytes to derive.
3128 * @return  0 on success always.
3129 */
3130int mlkem_kdf(const byte* seed, int seedLen, byte* out, int outLen)
3131{
3132    word64 state[25];
3133    word32 len64 = seedLen / 8;
3134
3135    readUnalignedWords64(state, seed, len64);
3136    state[len64] = 0x1f;
3137    XMEMSET(state + len64 + 1, 0, (25 - len64 - 1) * sizeof(word64));
3138    state[WC_SHA3_256_COUNT - 1] = W64LIT(0x8000000000000000);
3139
3140#ifndef WC_SHA3_NO_ASM
3141    if (IS_INTEL_BMI2(cpuid_flags)) {
3142        sha3_block_bmi2(state);
3143    }
3144    else if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
3145        sha3_block_avx2(state);
3146        RESTORE_VECTOR_REGISTERS();
3147    }
3148    else
3149#endif
3150    {
3151        BlockSha3(state);
3152    }
3153    XMEMCPY(out, state, outLen);
3154
3155    return 0;
3156}
3157#endif
3158
3159#if defined(WOLFSSL_ARMASM) && defined(__aarch64__)
3160/* Create pseudo-random key from the seed using SHAKE-256.
3161 *
3162 * @param  [in]  seed      Data to derive from.
3163 * @param  [in]  seedLen   Length of data to derive from in bytes.
3164 * @param  [out] out       Buffer to write to.
3165 * @param  [in]  outLen    Number of bytes to derive.
3166 * @return  0 on success always.
3167 */
3168int mlkem_kdf(const byte* seed, int seedLen, byte* out, int outLen)
3169{
3170    word64 state[25];
3171    word32 len64 = seedLen / 8;
3172
3173    readUnalignedWords64(state, seed, len64);
3174    state[len64] = 0x1f;
3175    XMEMSET(state + len64 + 1, 0, (25 - len64 - 1) * sizeof(word64));
3176    state[WC_SHA3_256_COUNT - 1] = W64LIT(0x8000000000000000);
3177
3178    BlockSha3(state);
3179    XMEMCPY(out, state, outLen);
3180
3181    return 0;
3182}
3183#endif
3184#endif
3185
3186#ifndef WOLFSSL_NO_ML_KEM
3187/* Derive the secret from z and cipher text.
3188 *
3189 * @param [in, out]  prf   SHAKE-256 object.
3190 * @param [in]       z     Implicit rejection value.
3191 * @param [in]       ct    Cipher text.
3192 * @param [in]       ctSz  Length of cipher text in bytes.
3193 * @param [out]      ss    Shared secret.
3194 * @return  0 on success.
3195 * @return  MEMORY_E when dynamic memory allocation failed.
3196 * @return  Other negative value when a hash error occurred.
3197 */
3198int mlkem_derive_secret(wc_Shake* prf, const byte* z, const byte* ct,
3199    word32 ctSz, byte* ss)
3200{
3201    int ret;
3202
3203#ifdef USE_INTEL_SPEEDUP
3204    XMEMCPY(prf->t, z, WC_ML_KEM_SYM_SZ);
3205    XMEMCPY(prf->t + WC_ML_KEM_SYM_SZ, ct,
3206        WC_SHA3_256_COUNT * 8 - WC_ML_KEM_SYM_SZ);
3207    prf->i = WC_ML_KEM_SYM_SZ + WC_SHA3_256_COUNT * 8 - WC_ML_KEM_SYM_SZ;
3208    ct += WC_SHA3_256_COUNT * 8 - WC_ML_KEM_SYM_SZ;
3209    ctSz -= WC_SHA3_256_COUNT * 8 - WC_ML_KEM_SYM_SZ;
3210    ret = wc_Shake256_Update(prf, ct, ctSz);
3211    if (ret == 0) {
3212        ret = wc_Shake256_Final(prf, ss, WC_ML_KEM_SS_SZ);
3213    }
3214#else
3215    ret = wc_InitShake256(prf, NULL, INVALID_DEVID);
3216    if (ret == 0) {
3217        ret = wc_Shake256_Update(prf, z, WC_ML_KEM_SYM_SZ);
3218    }
3219    if (ret == 0) {
3220        ret = wc_Shake256_Update(prf, ct, ctSz);
3221    }
3222    if (ret == 0) {
3223        ret = wc_Shake256_Final(prf, ss, WC_ML_KEM_SS_SZ);
3224    }
3225#endif
3226
3227    return ret;
3228}
3229#endif
3230
3231#if !defined(WOLFSSL_ARMASM)
3232/* Rejection sampling on uniform random bytes to generate uniform random
3233 * integers mod q.
3234 *
3235 * FIPS 203, Algorithm 7: SampleNTT(B)
3236 * Takes a 32-byte seed and two indices as input and outputs a pseudorandom
3237 * element of T_q.
3238 *   ...
3239 *   4: while j < 256 do
3240 *   5:     (ctx,C) <- XOF.Squeeze(ctx,3)
3241 *   6:     d1 <- C[0] + 256.(C[1] mod 16)
3242 *   7:     d2 <- lower(C[1] / 16) + 16.C[2]
3243 *   8:     if d1 < q then
3244 *   9:         a_hat[j] <- d1
3245 *  10:         j <- j + 1
3246 *  11:     end if
3247 *  12:     if d2 < q and j < 256 then
3248 *  13:         a_hat[j] <- d2
3249 *  14:         j <- j + 1
3250 *  15:     end if
3251 *  16: end while
3252 *  ...
3253 *
3254 * @param  [out]  p     Uniform random integers mod q.
3255 * @param  [in]   len   Maximum number of integers.
3256 * @param  [in]   r     Uniform random bytes buffer.
3257 * @param  [in]   rLen  Length of random data in buffer.
3258 * @return  Number of integers sampled.
3259 */
3260static unsigned int mlkem_rej_uniform_c(sword16* p, unsigned int len,
3261    const byte* r, unsigned int rLen)
3262{
3263    unsigned int i;
3264    unsigned int j;
3265
3266#if defined(WOLFSSL_MLKEM_SMALL) || !defined(WC_64BIT_CPU) || \
3267    defined(BIG_ENDIAN_ORDER)
3268    /* Keep sampling until max number of integers reached or buffer is used up.
3269     * Step 4. */
3270    for (i = 0, j = 0; (i < len) && (j <= rLen - 3); j += 3) {
3271        /* Step 5 - Now using 3 bytes of what the caller generated. */
3272        /* Use 24 bits (3 bytes) as two 12 bits integers. */
3273        /* Step 6. */
3274        sword16 v0 = ((r[0] >> 0) | ((word16)r[1] << 8)) & 0xFFF;
3275        /* Step 7. */
3276        sword16 v1 = ((r[1] >> 4) | ((word16)r[2] << 4)) & 0xFFF;
3277
3278        /* Reject first 12-bit integer if greater than or equal to q.
3279         * Step 8 */
3280        if (v0 < MLKEM_Q) {
3281            /* Steps 9-10 */
3282            p[i++] = v0;
3283        }
3284        /* Check second if we don't have enough integers yet.
3285         * Reject second 12-bit integer if greater than or equal to q.
3286         * Step 12 */
3287        if ((i < len) && (v1 < MLKEM_Q)) {
3288            /* Steps 13-14 */
3289            p[i++] = v1;
3290        }
3291
3292        /* Move over used bytes. */
3293        r += 3;
3294    }
3295#else
3296    /* Unroll loops. Minimal work per loop. */
3297    unsigned int minJ;
3298
3299    /* Calculate minimum number of 6 byte data blocks to get all required
3300     * numbers assuming no rejections. */
3301    minJ = len / 4 * 6;
3302    if (minJ > rLen)
3303        minJ = rLen;
3304    i = 0;
3305    for (j = 0; j < minJ; j += 6) {
3306        /* Use 48 bits (6 bytes) as four 12-bit integers. */
3307        word64 r_word = readUnalignedWord64(r);
3308        sword16 v0 =  r_word        & 0xfff;
3309        sword16 v1 = (r_word >> 12) & 0xfff;
3310        sword16 v2 = (r_word >> 24) & 0xfff;
3311        sword16 v3 = (r_word >> 36) & 0xfff;
3312
3313        p[i] = v0;
3314        i += (v0 < MLKEM_Q);
3315        p[i] = v1;
3316        i += (v1 < MLKEM_Q);
3317        p[i] = v2;
3318        i += (v2 < MLKEM_Q);
3319        p[i] = v3;
3320        i += (v3 < MLKEM_Q);
3321
3322        /* Move over used bytes. */
3323        r += 6;
3324    }
3325    /* Check whether we have all the numbers we need. */
3326    if (j < rLen) {
3327        /* Keep trying until we have fewer than 4 numbers to find or data is
3328         * used up. */
3329        for (; (i + 4 < len) && (j < rLen); j += 6) {
3330            /* Use 48 bits (6 bytes) as four 12-bit integers. */
3331            word64 r_word = readUnalignedWord64(r);
3332            sword16 v0 =  r_word        & 0xfff;
3333            sword16 v1 = (r_word >> 12) & 0xfff;
3334            sword16 v2 = (r_word >> 24) & 0xfff;
3335            sword16 v3 = (r_word >> 36) & 0xfff;
3336
3337            p[i] = v0;
3338            i += (v0 < MLKEM_Q);
3339            p[i] = v1;
3340            i += (v1 < MLKEM_Q);
3341            p[i] = v2;
3342            i += (v2 < MLKEM_Q);
3343            p[i] = v3;
3344            i += (v3 < MLKEM_Q);
3345
3346            /* Move over used bytes. */
3347            r += 6;
3348        }
3349        /* Keep trying until we have all the numbers we need or the data is used
3350         * up. */
3351        for (; (i < len) && (j < rLen); j += 6) {
3352            /* Use 48 bits (6 bytes) as four 12-bit integers. */
3353            word64 r_word = readUnalignedWord64(r);
3354            sword16 v0 =  r_word        & 0xfff;
3355            sword16 v1 = (r_word >> 12) & 0xfff;
3356            sword16 v2 = (r_word >> 24) & 0xfff;
3357            sword16 v3 = (r_word >> 36) & 0xfff;
3358
3359            /* Reject first 12-bit integer if greater than or equal to q. */
3360            if (v0 < MLKEM_Q) {
3361                p[i++] = v0;
3362            }
3363            /* Check second if we don't have enough integers yet.
3364             * Reject second 12-bit integer if greater than or equal to q. */
3365            if ((i < len) && (v1 < MLKEM_Q)) {
3366                p[i++] = v1;
3367            }
3368            /* Check third if we don't have enough integers yet.
3369             * Reject third 12-bit integer if greater than or equal to q. */
3370            if ((i < len) && (v2 < MLKEM_Q)) {
3371                p[i++] = v2;
3372            }
3373            /* Check fourth if we don't have enough integers yet.
3374             * Reject fourth 12-bit integer if greater than or equal to q. */
3375            if ((i < len) && (v3 < MLKEM_Q)) {
3376                p[i++] = v3;
3377            }
3378
3379            /* Move over used bytes. */
3380            r += 6;
3381        }
3382    }
3383#endif
3384
3385    return i;
3386}
3387#endif
3388
3389#if !defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) || \
3390    !defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM)
3391
3392#if !(defined(WOLFSSL_ARMASM) && defined(__aarch64__))
3393/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
3394 *
3395 * Seed used with XOF to generate random bytes.
3396 *
3397 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
3398 *   ...
3399 *   3: for (i <- 0; i < k; i++)
3400 *   4:     for (j <- 0; j < k; j++)
3401 *   5:         A_hat[i,j] <- SampleNTT(rho||j||i)
3402 *   6:     end for
3403 *   7: end for
3404 *   ...
3405 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
3406 *   ...
3407 *   4: for (i <- 0; i < k; i++)
3408 *   5:     for (j <- 0; j < k; j++)
3409 *   6:         A_hat[i,j] <- SampleNTT(rho||j||i)  (Transposed is rho||i||j)
3410 *   7:     end for
3411 *   8: end for
3412 *   ...
3413 * FIPS 203, Algorithm 7: SampleNTT(B)
3414 * Takes a 32-byte seed and two indices as input and outputs a pseudorandom
3415 * element of T_q.
3416 *   1: ctx <- XOF.init()
3417 *   2: ctx <- XOF.Absorb(ctx,B)
3418 *   3: j <- 0
3419 *   4: while j < 256 do
3420 *   5:     (ctx,C) <- XOF.Squeeze(ctx,3)
3421 *   ...
3422 *  16: end while
3423 *  17: return a_hat
3424 *
3425 * @param  [in, out]  prf         XOF object.
3426 * @param  [out]      a           Matrix of uniform integers.
3427 * @param  [in]       k           Number of dimensions. k x k polynomials.
3428 * @param  [in]       seed        Bytes to seed XOF generation.
3429 * @param  [in]       transposed  Whether A or A^T is generated.
3430 * @return  0 on success.
3431 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
3432 *          WOLFSSL_SMALL_STACK is defined.
3433 */
3434static int mlkem_gen_matrix_c(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed,
3435    int transposed)
3436{
3437#if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
3438    byte* rand;
3439#else
3440    byte rand[GEN_MATRIX_SIZE + 2];
3441#endif
3442    byte extSeed[WC_ML_KEM_SYM_SZ + 2];
3443    int ret = 0;
3444    int i;
3445
3446    /* Copy seed into buffer that has space for i and j to be appended. */
3447    XMEMCPY(extSeed, seed, WC_ML_KEM_SYM_SZ);
3448
3449#if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
3450    /* Allocate large amount of memory to hold random bytes to be sampled. */
3451    rand = (byte*)XMALLOC(GEN_MATRIX_SIZE + 2, NULL, DYNAMIC_TYPE_TMP_BUFFER);
3452    if (rand == NULL) {
3453        ret = MEMORY_E;
3454    }
3455#endif
3456
3457#if !defined(WOLFSSL_MLKEM_SMALL) && defined(WC_64BIT_CPU)
3458    /* Loading 64 bits, only using 48 bits. Loading 2 bytes more than used. */
3459    if (ret == 0) {
3460        rand[GEN_MATRIX_SIZE+0] = 0xff;
3461        rand[GEN_MATRIX_SIZE+1] = 0xff;
3462    }
3463#endif
3464
3465    /* Generate each vector of polynomials.
3466     * Alg 13, Step 3. Alg 14, Step 4. */
3467    for (i = 0; (ret == 0) && (i < k); i++, a += k * MLKEM_N) {
3468        int j;
3469        /* Generate each polynomial in vector from seed with indices.
3470         * Alg 13, Step 4. Alg 14, Step 5. */
3471        for (j = 0; (ret == 0) && (j < k); j++) {
3472            if (transposed) {
3473                /* Alg 14, Step 6: .. rho||i||j ... */
3474                extSeed[WC_ML_KEM_SYM_SZ + 0] = (byte)i;
3475                extSeed[WC_ML_KEM_SYM_SZ + 1] = (byte)j;
3476            }
3477            else {
3478                /* Alg 13, Step 5: .. rho||j||i ... */
3479                extSeed[WC_ML_KEM_SYM_SZ + 0] = (byte)j;
3480                extSeed[WC_ML_KEM_SYM_SZ + 1] = (byte)i;
3481            }
3482            /* Absorb the index specific seed.
3483             * Alg 7, Step 1-2 */
3484            ret = mlkem_xof_absorb(prf, extSeed, sizeof(extSeed));
3485            if (ret == 0) {
3486                /* Create data based on the seed.
3487                 * Alg 7, Step 5. Generating enough to, on average, be able to
3488                 * get enough valid values. */
3489                ret = mlkem_xof_squeezeblocks(prf, rand, GEN_MATRIX_NBLOCKS);
3490            }
3491            if (ret == 0) {
3492                unsigned int ctr;
3493
3494                /* Sample random bytes to create a polynomial.
3495                 * Alg 7, Step 3 - implicitly counter is 0.
3496                 * Alg 7, Step 4-16. */
3497                ctr = mlkem_rej_uniform_c(a + j * MLKEM_N, MLKEM_N, rand,
3498                    GEN_MATRIX_SIZE);
3499                /* Create more blocks if too many rejected.
3500                 * Alg 7, Step 4. */
3501                while (ctr < MLKEM_N) {
3502                    /* Alg 7, Step 5. */
3503                    mlkem_xof_squeezeblocks(prf, rand, 1);
3504                    /* Alg 7, Step 4-16. */
3505                    ctr += mlkem_rej_uniform_c(a + j * MLKEM_N + ctr,
3506                        MLKEM_N - ctr, rand, XOF_BLOCK_SIZE);
3507                }
3508            }
3509        }
3510    }
3511
3512#if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
3513    /* Dispose of temporary buffer. */
3514    XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
3515#endif
3516
3517    return ret;
3518}
3519#endif
3520
3521/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
3522 *
3523 * Seed used with XOF to generate random bytes.
3524 *
3525 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d), Steps 3-7
3526 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r), Steps 4-8
3527 *
3528 * @param  [in, out]  prf         XOF object.
3529 * @param  [out]      a           Matrix of uniform integers.
3530 * @param  [in]       k           Number of dimensions. k x k polynomials.
3531 * @param  [in]       seed        Bytes to seed XOF generation.
3532 * @param  [in]       transposed  Whether A or A^T is generated.
3533 * @return  0 on success.
3534 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
3535 *          WOLFSSL_SMALL_STACK is defined.
3536 */
3537int mlkem_gen_matrix(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed,
3538    int transposed)
3539{
3540    int ret;
3541
3542#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
3543    if (k == WC_ML_KEM_512_K) {
3544#if defined(WOLFSSL_ARMASM) && defined(__aarch64__)
3545        ret = mlkem_gen_matrix_k2_aarch64(a, seed, transposed);
3546#else
3547    #if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
3548        if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
3549            ret = mlkem_gen_matrix_k2_avx2(a, seed, transposed);
3550            RESTORE_VECTOR_REGISTERS();
3551        }
3552        else
3553    #endif
3554        {
3555            ret = mlkem_gen_matrix_c(prf, a, WC_ML_KEM_512_K, seed, transposed);
3556        }
3557#endif
3558    }
3559    else
3560#endif
3561#if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
3562    if (k == WC_ML_KEM_768_K) {
3563#if defined(WOLFSSL_ARMASM) && defined(__aarch64__)
3564        ret = mlkem_gen_matrix_k3_aarch64(a, seed, transposed);
3565#else
3566    #if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
3567        if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
3568            ret = mlkem_gen_matrix_k3_avx2(a, seed, transposed);
3569            RESTORE_VECTOR_REGISTERS();
3570        }
3571        else
3572    #endif
3573        {
3574            ret = mlkem_gen_matrix_c(prf, a, WC_ML_KEM_768_K, seed, transposed);
3575        }
3576#endif
3577    }
3578    else
3579#endif
3580#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
3581    if (k == WC_ML_KEM_1024_K) {
3582#if defined(WOLFSSL_ARMASM) && defined(__aarch64__)
3583        ret = mlkem_gen_matrix_k4_aarch64(a, seed, transposed);
3584#else
3585    #if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
3586        if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
3587            ret = mlkem_gen_matrix_k4_avx2(a, seed, transposed);
3588            RESTORE_VECTOR_REGISTERS();
3589        }
3590        else
3591    #endif
3592        {
3593            ret = mlkem_gen_matrix_c(prf, a, WC_ML_KEM_1024_K, seed,
3594                transposed);
3595        }
3596#endif
3597    }
3598    else
3599#endif
3600    {
3601        ret = BAD_STATE_E;
3602    }
3603
3604    (void)prf;
3605
3606    return ret;
3607}
3608
3609#endif
3610
3611#if defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) || \
3612    defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM)
3613
3614/* Deterministically generate a matrix (or transpose) of uniform integers mod q.
3615 *
3616 * Seed used with XOF to generate random bytes.
3617 *
3618 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
3619 * ...
3620 * 4:     for (j <- 0; j < k; j++)
3621 * 5:         A_hat[i,j] <- SampleNTT(rho||j||i)
3622 * 6:     end for
3623 * ...
3624 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
3625 * ...
3626 * 5:     for (j <- 0; j < k; j++)
3627 * 6:         A_hat[i,j] <- SampleNTT(rho||j||i)  (Transposed is rho||i||j)
3628 * 7:     end for
3629 * ...
3630 *
3631 * @param  [in, out]  prf         XOF object.
3632 * @param  [out]      a           Matrix of uniform integers.
3633 * @param  [in]       k           Number of dimensions. k x k polynomials.
3634 * @param  [in]       seed        Bytes to seed XOF generation.
3635 * @param  [in]       i           Index of vector to generate.
3636 * @param  [in]       transposed  Whether A or A^T is generated.
3637 * @return  0 on success.
3638 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
3639 *          WOLFSSL_SMALL_STACK is defined.
3640 */
3641static int mlkem_gen_matrix_i(MLKEM_PRF_T* prf, sword16* a, int k, byte* seed,
3642    int i, int transposed)
3643{
3644#if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
3645    byte* rand;
3646#else
3647    byte rand[GEN_MATRIX_SIZE + 2];
3648#endif
3649    byte extSeed[WC_ML_KEM_SYM_SZ + 2];
3650    int ret = 0;
3651    int j;
3652
3653    XMEMCPY(extSeed, seed, WC_ML_KEM_SYM_SZ);
3654
3655#if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
3656    /* Allocate large amount of memory to hold random bytes to be sampled. */
3657    rand = (byte*)XMALLOC(GEN_MATRIX_SIZE + 2, NULL, DYNAMIC_TYPE_TMP_BUFFER);
3658    if (rand == NULL) {
3659        ret = MEMORY_E;
3660    }
3661#endif
3662
3663#if !defined(WOLFSSL_MLKEM_SMALL) && defined(WC_64BIT_CPU)
3664    /* Loading 64 bits, only using 48 bits. Loading 2 bytes more than used. */
3665    if (ret == 0) {
3666        rand[GEN_MATRIX_SIZE+0] = 0xff;
3667        rand[GEN_MATRIX_SIZE+1] = 0xff;
3668    }
3669#endif
3670
3671    /* Generate each polynomial in vector from seed with indices.
3672     * Alg 13, Step 4. Alg 14, Step 5. */
3673    for (j = 0; (ret == 0) && (j < k); j++) {
3674        if (transposed) {
3675            /* Alg 14, Step 6: .. rho||i||j ... */
3676            extSeed[WC_ML_KEM_SYM_SZ + 0] = (byte)i;
3677            extSeed[WC_ML_KEM_SYM_SZ + 1] = (byte)j;
3678        }
3679        else {
3680            /* Alg 13, Step 5: .. rho||j||i ... */
3681            extSeed[WC_ML_KEM_SYM_SZ + 0] = (byte)j;
3682            extSeed[WC_ML_KEM_SYM_SZ + 1] = (byte)i;
3683        }
3684        /* Absorb the index specific seed.
3685         * Alg 7, Step 1-2 */
3686        ret = mlkem_xof_absorb(prf, extSeed, sizeof(extSeed));
3687        if (ret == 0) {
3688            /* Create data based on the seed.
3689             * Alg 7, Step 5. Generating enough to, on average, be able to get
3690             * enough valid values. */
3691            ret = mlkem_xof_squeezeblocks(prf, rand, GEN_MATRIX_NBLOCKS);
3692        }
3693        if (ret == 0) {
3694            unsigned int ctr;
3695
3696            /* Sample random bytes to create a polynomial.
3697             * Alg 7, Step 3 - implicitly counter is 0.
3698             * Alg 7, Step 4-16. */
3699            ctr = mlkem_rej_uniform_c(a + j * MLKEM_N, MLKEM_N, rand,
3700                GEN_MATRIX_SIZE);
3701            /* Create more blocks if too many rejected.
3702             * Alg 7, Step 4. */
3703            while (ctr < MLKEM_N) {
3704                /* Alg 7, Step 5. */
3705                mlkem_xof_squeezeblocks(prf, rand, 1);
3706                /* Alg 7, Step 4-16. */
3707                ctr += mlkem_rej_uniform_c(a + j * MLKEM_N + ctr,
3708                    MLKEM_N - ctr, rand, XOF_BLOCK_SIZE);
3709            }
3710        }
3711    }
3712
3713#if defined(WOLFSSL_SMALL_STACK) && !defined(WOLFSSL_NO_MALLOC)
3714    /* Dispose of temporary buffer. */
3715    XFREE(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
3716#endif
3717
3718    return ret;
3719}
3720
3721#endif
3722
3723
3724/******************************************************************************/
3725
3726/* Subtract one 2 bit value from another out of a larger number.
3727 *
3728 * FIPS 203, Algorithm 8: SamplePolyCBD_eta(B)
3729 * Takes a seed as input and outputs a pseudorandom sample from the distribution
3730 * D_eta(R_q).
3731 *
3732 * @param  [in]  d  Value containing sequential 2 bit values.
3733 * @param  [in]  i  Start index of the two values in 2 bits each.
3734 * @return  Difference of the two values with range -2..2.
3735 */
3736#define ETA2_SUB(d, i) \
3737    (sword16)(((sword16)(((d) >> ((i) * 4 + 0)) & 0x3)) - \
3738              ((sword16)(((d) >> ((i) * 4 + 2)) & 0x3)))
3739
3740/* Compute polynomial with coefficients distributed according to a centered
3741 * binomial distribution with parameter eta2 from uniform random bytes.
3742 *
3743 * FIPS 203, Algorithm 8: SamplePolyCBD_eta(B)
3744 * Takes a seed as input and outputs a pseudorandom sample from the distribution
3745 * D_eta(R_q).
3746 *
3747 * @param [out]  p  Polynomial computed.
3748 * @param [in]   r  Random bytes.
3749 */
3750static void mlkem_cbd_eta2(sword16* p, const byte* r)
3751{
3752    unsigned int i;
3753
3754#ifndef WORD64_AVAILABLE
3755    /* Calculate eight integer coefficients at a time. */
3756    for (i = 0; i < MLKEM_N; i += 8) {
3757    #ifdef WOLFSSL_MLKEM_SMALL
3758        unsigned int j;
3759    #endif
3760        /* Take the next 4 bytes, little endian, as a 32 bit value. */
3761    #ifdef BIG_ENDIAN_ORDER
3762        word32 t = ByteReverseWord32(*(word32*)r);
3763    #else
3764        word32 t = *(word32*)r;
3765    #endif
3766        word32 d;
3767        /* Add second bits to first. */
3768        d  = (t >> 0) & 0x55555555;
3769        d += (t >> 1) & 0x55555555;
3770        /* Values 0, 1 or 2 in consecutive 2 bits.
3771         * 0 - 1/4, 1 - 2/4, 2 - 1/4. */
3772
3773    #ifdef WOLFSSL_MLKEM_SMALL
3774        for (j = 0; j < 8; j++) {
3775            p[i + j] = ETA2_SUB(d, j);
3776        }
3777    #else
3778        p[i + 0] = ETA2_SUB(d, 0);
3779        p[i + 1] = ETA2_SUB(d, 1);
3780        p[i + 2] = ETA2_SUB(d, 2);
3781        p[i + 3] = ETA2_SUB(d, 3);
3782        p[i + 4] = ETA2_SUB(d, 4);
3783        p[i + 5] = ETA2_SUB(d, 5);
3784        p[i + 6] = ETA2_SUB(d, 6);
3785        p[i + 7] = ETA2_SUB(d, 7);
3786    #endif
3787        /* -2 - 1/16, -1 - 4/16, 0 - 6/16, 1 - 4/16, 2 - 1/16  */
3788
3789        /* Move over used bytes. */
3790        r += 4;
3791    }
3792#else
3793    /* Calculate sixteen integer coefficients at a time. */
3794    for (i = 0; i < MLKEM_N; i += 16) {
3795    #ifdef WOLFSSL_MLKEM_SMALL
3796        unsigned int j;
3797    #endif
3798        /* Take the next 8 bytes, little endian, as a 64 bit value. */
3799    #ifdef BIG_ENDIAN_ORDER
3800        word64 t = ByteReverseWord64(readUnalignedWord64(r));
3801    #else
3802        word64 t = readUnalignedWord64(r);
3803    #endif
3804        word64 d;
3805        /* Add second bits to first. */
3806        d  = (t >> 0) & 0x5555555555555555L;
3807        d += (t >> 1) & 0x5555555555555555L;
3808        /* Values 0, 1 or 2 in consecutive 2 bits.
3809         * 0 - 1/4, 1 - 2/4, 2 - 1/4. */
3810
3811    #ifdef WOLFSSL_MLKEM_SMALL
3812        for (j = 0; j < 16; j++) {
3813            p[i + j] = ETA2_SUB(d, j);
3814        }
3815    #else
3816        p[i +  0] = ETA2_SUB(d,  0);
3817        p[i +  1] = ETA2_SUB(d,  1);
3818        p[i +  2] = ETA2_SUB(d,  2);
3819        p[i +  3] = ETA2_SUB(d,  3);
3820        p[i +  4] = ETA2_SUB(d,  4);
3821        p[i +  5] = ETA2_SUB(d,  5);
3822        p[i +  6] = ETA2_SUB(d,  6);
3823        p[i +  7] = ETA2_SUB(d,  7);
3824        p[i +  8] = ETA2_SUB(d,  8);
3825        p[i +  9] = ETA2_SUB(d,  9);
3826        p[i + 10] = ETA2_SUB(d, 10);
3827        p[i + 11] = ETA2_SUB(d, 11);
3828        p[i + 12] = ETA2_SUB(d, 12);
3829        p[i + 13] = ETA2_SUB(d, 13);
3830        p[i + 14] = ETA2_SUB(d, 14);
3831        p[i + 15] = ETA2_SUB(d, 15);
3832    #endif
3833        /* -2 - 1/16, -1 - 4/16, 0 - 6/16, 1 - 4/16, 2 - 1/16  */
3834
3835        /* Move over used bytes. */
3836        r += 8;
3837    }
3838#endif
3839}
3840
3841#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
3842/* Subtract one 3 bit value from another out of a larger number.
3843 *
3844 * FIPS 203, Algorithm 8: SamplePolyCBD_eta(B)
3845 * Takes a seed as input and outputs a pseudorandom sample from the distribution
3846 * D_eta(R_q).
3847 *
3848 * @param  [in]  d  Value containing sequential 3 bit values.
3849 * @param  [in]  i  Start index of the two values in 3 bits each.
3850 * @return  Difference of the two values with range -3..3.
3851 */
3852#define ETA3_SUB(d, i) \
3853    (sword16)(((sword16)(((d) >> ((i) * 6 + 0)) & 0x7)) - \
3854              ((sword16)(((d) >> ((i) * 6 + 3)) & 0x7)))
3855
3856/* Compute polynomial with coefficients distributed according to a centered
3857 * binomial distribution with parameter eta3 from uniform random bytes.
3858 *
3859 * FIPS 203, Algorithm 8: SamplePolyCBD_eta(B)
3860 * Takes a seed as input and outputs a pseudorandom sample from the distribution
3861 * D_eta(R_q).
3862 *
3863 * @param [out]  p  Polynomial computed.
3864 * @param [in]   r  Random bytes.
3865 */
3866static void mlkem_cbd_eta3(sword16* p, const byte* r)
3867{
3868    unsigned int i;
3869
3870#if defined(WOLFSSL_SMALL_STACK) || defined(WOLFSSL_MLKEM_NO_LARGE_CODE) || \
3871    defined(BIG_ENDIAN_ORDER)
3872#ifndef WORD64_AVAILABLE
3873    /* Calculate four integer coefficients at a time. */
3874    for (i = 0; i < MLKEM_N; i += 4) {
3875    #ifdef WOLFSSL_MLKEM_SMALL
3876        unsigned int j;
3877    #endif
3878        /* Take the next 3 bytes, little endian, as a 24 bit value. */
3879        word32 t = (((word32)(r[0])) <<  0) |
3880                   (((word32)(r[1])) <<  8) |
3881                   (((word32)(r[2])) << 16);
3882        word32 d;
3883        /* Add second and third bits to first. */
3884        d  = (t >> 0) & 0x00249249;
3885        d += (t >> 1) & 0x00249249;
3886        d += (t >> 2) & 0x00249249;
3887        /* Values 0, 1, 2 or 3 in consecutive 3 bits.
3888         * 0 - 1/8, 1 - 3/8, 2 - 3/8, 3 - 1/8. */
3889
3890    #ifdef WOLFSSL_MLKEM_SMALL
3891        for (j = 0; j < 4; j++) {
3892            p[i + j] = ETA3_SUB(d, j);
3893        }
3894    #else
3895        p[i + 0] = ETA3_SUB(d, 0);
3896        p[i + 1] = ETA3_SUB(d, 1);
3897        p[i + 2] = ETA3_SUB(d, 2);
3898        p[i + 3] = ETA3_SUB(d, 3);
3899    #endif
3900        /* -3-1/64, -2-6/64, -1-15/64, 0-20/64, 1-15/64, 2-6/64, 3-1/64 */
3901
3902        /* Move over used bytes. */
3903        r += 3;
3904    }
3905#else
3906    /* Calculate eight integer coefficients at a time. */
3907    for (i = 0; i < MLKEM_N; i += 8) {
3908    #ifdef WOLFSSL_MLKEM_SMALL
3909        unsigned int j;
3910    #endif
3911        /* Take the next 6 bytes, little endian, as a 48 bit value. */
3912        word64 t = (((word64)(r[0])) <<  0) |
3913                   (((word64)(r[1])) <<  8) |
3914                   (((word64)(r[2])) << 16) |
3915                   (((word64)(r[3])) << 24) |
3916                   (((word64)(r[4])) << 32) |
3917                   (((word64)(r[5])) << 40);
3918        word64 d;
3919        /* Add second and third bits to first. */
3920        d  = (t >> 0) & 0x0000249249249249L;
3921        d += (t >> 1) & 0x0000249249249249L;
3922        d += (t >> 2) & 0x0000249249249249L;
3923        /* Values 0, 1, 2 or 3 in consecutive 3 bits.
3924         * 0 - 1/8, 1 - 3/8, 2 - 3/8, 3 - 1/8. */
3925
3926    #ifdef WOLFSSL_MLKEM_SMALL
3927        for (j = 0; j < 8; j++) {
3928            p[i + j] = ETA3_SUB(d, j);
3929        }
3930    #else
3931        p[i + 0] = ETA3_SUB(d, 0);
3932        p[i + 1] = ETA3_SUB(d, 1);
3933        p[i + 2] = ETA3_SUB(d, 2);
3934        p[i + 3] = ETA3_SUB(d, 3);
3935        p[i + 4] = ETA3_SUB(d, 4);
3936        p[i + 5] = ETA3_SUB(d, 5);
3937        p[i + 6] = ETA3_SUB(d, 6);
3938        p[i + 7] = ETA3_SUB(d, 7);
3939    #endif
3940        /* -3-1/64, -2-6/64, -1-15/64, 0-20/64, 1-15/64, 2-6/64, 3-1/64 */
3941
3942        /* Move over used bytes. */
3943        r += 6;
3944    }
3945#endif /* WORD64_AVAILABLE */
3946#else
3947    /* Calculate eight integer coefficients at a time. */
3948    for (i = 0; i < MLKEM_N; i += 16) {
3949        const word32* r32 = (const word32*)r;
3950        /* Take the next 12 bytes, little endian, as 24 bit values. */
3951        word32 t0 =   r32[0]                          & 0xffffff;
3952        word32 t1 = ((r32[0] >> 24) | (r32[1] <<  8)) & 0xffffff;
3953        word32 t2 = ((r32[1] >> 16) | (r32[2] << 16)) & 0xffffff;
3954        word32 t3 =   r32[2] >>  8                              ;
3955        word32 d0;
3956        word32 d1;
3957        word32 d2;
3958        word32 d3;
3959
3960        /* Add second and third bits to first. */
3961        d0  = (t0 >> 0) & 0x00249249;
3962        d0 += (t0 >> 1) & 0x00249249;
3963        d0 += (t0 >> 2) & 0x00249249;
3964        d1  = (t1 >> 0) & 0x00249249;
3965        d1 += (t1 >> 1) & 0x00249249;
3966        d1 += (t1 >> 2) & 0x00249249;
3967        d2  = (t2 >> 0) & 0x00249249;
3968        d2 += (t2 >> 1) & 0x00249249;
3969        d2 += (t2 >> 2) & 0x00249249;
3970        d3  = (t3 >> 0) & 0x00249249;
3971        d3 += (t3 >> 1) & 0x00249249;
3972        d3 += (t3 >> 2) & 0x00249249;
3973        /* Values 0, 1, 2 or 3 in consecutive 3 bits.
3974         * 0 - 1/8, 1 - 3/8, 2 - 3/8, 3 - 1/8. */
3975
3976        p[i +  0] = ETA3_SUB(d0, 0);
3977        p[i +  1] = ETA3_SUB(d0, 1);
3978        p[i +  2] = ETA3_SUB(d0, 2);
3979        p[i +  3] = ETA3_SUB(d0, 3);
3980        p[i +  4] = ETA3_SUB(d1, 0);
3981        p[i +  5] = ETA3_SUB(d1, 1);
3982        p[i +  6] = ETA3_SUB(d1, 2);
3983        p[i +  7] = ETA3_SUB(d1, 3);
3984        p[i +  8] = ETA3_SUB(d2, 0);
3985        p[i +  9] = ETA3_SUB(d2, 1);
3986        p[i + 10] = ETA3_SUB(d2, 2);
3987        p[i + 11] = ETA3_SUB(d2, 3);
3988        p[i + 12] = ETA3_SUB(d3, 0);
3989        p[i + 13] = ETA3_SUB(d3, 1);
3990        p[i + 14] = ETA3_SUB(d3, 2);
3991        p[i + 15] = ETA3_SUB(d3, 3);
3992        /* -3-1/64, -2-6/64, -1-15/64, 0-20/64, 1-15/64, 2-6/64, 3-1/64 */
3993
3994        /* Move over used bytes. */
3995        r += 12;
3996    }
3997#endif /* WOLFSSL_SMALL_STACK || WOLFSSL_MLKEM_NO_LARGE_CODE ||
3998        * BIG_ENDIAN_ORDER */
3999}
4000#endif
4001
4002#if !(defined(__aarch64__) && defined(WOLFSSL_ARMASM))
4003
4004/* Get noise/error by calculating random bytes and sampling to a binomial
4005 * distribution.
4006 *
4007 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
4008 *   ...
4009 *   9:     s[i] <- SamplePolyCBD_eta_1(PRF_eta_1(sigma, N))
4010 *   ...
4011 *  13:     e[i] <- SamplePolyCBD_eta_1(PRF_eta_1(sigma, N))
4012 *   ...
4013 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
4014 *   ...
4015 *  10:     y[i] <- SamplePolyCBD_eta_1(PRF_eta_1(r, N))
4016 *   ...
4017 *
4018 * @param  [in, out]  prf   Pseudo-random function object.
4019 * @param  [out]      p     Polynomial.
4020 * @param  [in]       seed  Seed to use when calculating random.
4021 * @param  [in]       eta1  Size of noise/error integers.
4022 * @return  0 on success.
4023 */
4024static int mlkem_get_noise_eta1_c(MLKEM_PRF_T* prf, sword16* p,
4025    const byte* seed, byte eta1)
4026{
4027    int ret;
4028
4029    (void)eta1;
4030
4031#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
4032    if (eta1 == MLKEM_CBD_ETA3) {
4033        byte rand[ETA3_RAND_SIZE];
4034
4035        /* Calculate random bytes from seed with PRF. */
4036        ret = mlkem_prf(prf, rand, sizeof(rand), seed);
4037        if (ret == 0) {
4038            /* Sample for values in range -3..3 from 3 bits of random. */
4039            mlkem_cbd_eta3(p, rand);
4040         }
4041    }
4042    else
4043#endif
4044    {
4045        byte rand[ETA2_RAND_SIZE];
4046
4047        /* Calculate random bytes from seed with PRF. */
4048        ret = mlkem_prf(prf, rand, sizeof(rand), seed);
4049        if (ret == 0) {
4050            /* Sample for values in range -2..2 from 2 bits of random. */
4051            mlkem_cbd_eta2(p, rand);
4052        }
4053    }
4054
4055    return ret;
4056}
4057
4058/* Get noise/error by calculating random bytes and sampling to a binomial
4059 * distribution. Values -2..2
4060 *
4061 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
4062 *   ...
4063 *  14:     e1[i] <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4064 *   ...
4065 *  17:     e2 <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4066 *   ...
4067 *
4068 * @param  [in, out]  prf   Pseudo-random function object.
4069 * @param  [out]      p     Polynomial.
4070 * @param  [in]       seed  Seed to use when calculating random.
4071 * @return  0 on success.
4072 */
4073static int mlkem_get_noise_eta2_c(MLKEM_PRF_T* prf, sword16* p,
4074    const byte* seed)
4075{
4076    int ret;
4077    byte rand[ETA2_RAND_SIZE];
4078
4079    /* Calculate random bytes from seed with PRF. */
4080    ret = mlkem_prf(prf, rand, sizeof(rand), seed);
4081    if (ret == 0) {
4082        mlkem_cbd_eta2(p, rand);
4083    }
4084
4085    return ret;
4086}
4087
4088#endif
4089
4090#if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
4091#define PRF_RAND_SZ   (2 * SHA3_256_BYTES)
4092
4093#if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768) || \
4094    defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
4095/* Get the noise/error by calculating random bytes.
4096 *
4097 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
4098 *   ...
4099 *  14:     e1[i] <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4100 *   ...
4101 *  17:     e2 <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4102 *   ...
4103 *
4104 * @param  [out]  rand  Random number byte array.
4105 * @param  [in]   seed  Seed to generate random from.
4106 * @param  [in]   o     Offset of seed count.
4107 */
4108static void mlkem_get_noise_x4_eta2_avx2(byte* rand, byte* seed, byte o)
4109{
4110    int i;
4111    word64 state[25 * 4];
4112
4113    for (i = 0; i < 4; i++) {
4114        state[4*4 + i] = (word32)(0x1f00 + i + o);
4115    }
4116
4117    sha3_256_blocksx4_seed_avx2(state, seed);
4118    mlkem_redistribute_16_rand_avx2(state, rand + 0 * ETA2_RAND_SIZE,
4119        rand + 1 * ETA2_RAND_SIZE, rand + 2 * ETA2_RAND_SIZE,
4120        rand + 3 * ETA2_RAND_SIZE);
4121}
4122#endif
4123
4124#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512) || \
4125    defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
4126/* Get noise/error by calculating random bytes and sampling to a binomial
4127 * distribution. Values -2..2
4128 *
4129 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
4130 *   ...
4131 *  14:     e1[i] <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4132 *   ...
4133 *  17:     e2 <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4134 *   ...
4135 *
4136 * @param  [in, out]  prf   Pseudo-random function object.
4137 * @param  [out]      p     Polynomial.
4138 * @param  [in]       seed  Seed to use when calculating random.
4139 * @return  0 on success.
4140 */
4141static int mlkem_get_noise_eta2_avx2(MLKEM_PRF_T* prf, sword16* p,
4142    const byte* seed)
4143{
4144    word64 state[25];
4145
4146    (void)prf;
4147
4148    /* Put first WC_ML_KEM_SYM_SZ bytes of key into blank state. */
4149    readUnalignedWords64(state, seed, WC_ML_KEM_SYM_SZ / sizeof(word64));
4150    /* Last byte in with end of content marker. */
4151    state[WC_ML_KEM_SYM_SZ / 8] = 0x1f00 | seed[WC_ML_KEM_SYM_SZ];
4152    /* Set rest of state to 0. */
4153    XMEMSET(state + WC_ML_KEM_SYM_SZ / 8 + 1, 0,
4154        (25 - WC_ML_KEM_SYM_SZ / 8 - 1) * sizeof(word64));
4155    /* ... except for rate marker. */
4156    state[WC_SHA3_256_COUNT - 1] = W64LIT(0x8000000000000000);
4157
4158    /* Perform a block operation on the state for next block of output. */
4159#ifndef WC_SHA3_NO_ASM
4160    if (IS_INTEL_BMI2(cpuid_flags)) {
4161        sha3_block_bmi2(state);
4162    }
4163    else if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
4164        sha3_block_avx2(state);
4165        RESTORE_VECTOR_REGISTERS();
4166    }
4167    else
4168#endif /* !WC_SHA3_NO_ASM */
4169    {
4170        BlockSha3(state);
4171    }
4172    mlkem_cbd_eta2_avx2(p, (byte*)state);
4173
4174    return 0;
4175}
4176#endif
4177
4178#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
4179/* Get the noise/error by calculating random bytes.
4180 *
4181 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
4182 *   ...
4183 *   9:     s[i] <- SamplePolyCBD_eta_1(PRF_eta_1(sigma, N))
4184 *   ...
4185 *  13:     e[i] <- SamplePolyCBD_eta_1(PRF_eta_1(sigma, N))
4186 *   ...
4187 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
4188 *   ...
4189 *  10:     y[i] <- SamplePolyCBD_eta_1(PRF_eta_1(r, N))
4190 *   ...
4191 *
4192 * @param  [out]  rand  Random number byte array.
4193 * @param  [in]   seed  Seed to generate random from.
4194 */
4195static void mlkem_get_noise_x4_eta3_avx2(byte* rand, byte* seed)
4196{
4197    word64 state[25 * 4];
4198    int i;
4199
4200    state[4*4 + 0] = 0x1f00 + 0;
4201    state[4*4 + 1] = 0x1f00 + 1;
4202    state[4*4 + 2] = 0x1f00 + 2;
4203    state[4*4 + 3] = 0x1f00 + 3;
4204
4205    sha3_256_blocksx4_seed_avx2(state, seed);
4206    mlkem_redistribute_17_rand_avx2(state, rand + 0 * PRF_RAND_SZ,
4207        rand + 1 * PRF_RAND_SZ, rand + 2 * PRF_RAND_SZ,
4208        rand + 3 * PRF_RAND_SZ);
4209    i = SHA3_256_BYTES;
4210    sha3_blocksx4_avx2(state);
4211    mlkem_redistribute_8_rand_avx2(state, rand + i + 0 * PRF_RAND_SZ,
4212        rand + i + 1 * PRF_RAND_SZ, rand + i + 2 * PRF_RAND_SZ,
4213        rand + i + 3 * PRF_RAND_SZ);
4214}
4215
4216/* Get the noise/error by calculating random bytes and sampling to a binomial
4217 * distribution.
4218 *
4219 * @param  [in, out]  prf   Pseudo-random function object.
4220 * @param  [out]      vec1  First Vector of polynomials.
4221 * @param  [out]      vec2  Second Vector of polynomials.
4222 * @param  [out]      poly  Polynomial.
4223 * @param  [in, out]  seed  Seed to use when calculating random.
4224 * @return  0 on success.
4225 * @return  MEMORY_E when dynamic memory allocation fails. Only possible when
4226 *          WOLFSSL_SMALL_STACK is defined.
4227 */
4228static int mlkem_get_noise_k2_avx2(MLKEM_PRF_T* prf, sword16* vec1,
4229    sword16* vec2, sword16* poly, byte* seed)
4230{
4231    int ret = 0;
4232    WC_DECLARE_VAR(rand, byte, 4 * PRF_RAND_SZ, 0);
4233
4234    WC_ALLOC_VAR_EX(rand, byte, 4 * PRF_RAND_SZ, NULL, DYNAMIC_TYPE_TMP_BUFFER,
4235        return MEMORY_E);
4236
4237    mlkem_get_noise_x4_eta3_avx2(rand, seed);
4238    mlkem_cbd_eta3_avx2(vec1          , rand + 0 * PRF_RAND_SZ);
4239    mlkem_cbd_eta3_avx2(vec1 + MLKEM_N, rand + 1 * PRF_RAND_SZ);
4240    if (poly == NULL) {
4241        mlkem_cbd_eta3_avx2(vec2          , rand + 2 * PRF_RAND_SZ);
4242        mlkem_cbd_eta3_avx2(vec2 + MLKEM_N, rand + 3 * PRF_RAND_SZ);
4243    }
4244    else {
4245        mlkem_cbd_eta2_avx2(vec2          , rand + 2 * PRF_RAND_SZ);
4246        mlkem_cbd_eta2_avx2(vec2 + MLKEM_N, rand + 3 * PRF_RAND_SZ);
4247
4248        seed[WC_ML_KEM_SYM_SZ] = 4;
4249        ret = mlkem_get_noise_eta2_avx2(prf, poly, seed);
4250    }
4251
4252    WC_FREE_VAR_EX(rand, NULL, DYNAMIC_TYPE_TMP_BUFFER);
4253
4254    return ret;
4255}
4256#endif
4257
4258#if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
4259/* Get the noise/error by calculating random bytes and sampling to a binomial
4260 * distribution.
4261 *
4262 * @param  [out]      vec1  First Vector of polynomials.
4263 * @param  [out]      vec2  Second Vector of polynomials.
4264 * @param  [out]      poly  Polynomial.
4265 * @param  [in]       seed  Seed to use when calculating random.
4266 * @return  0 on success.
4267 */
4268static int mlkem_get_noise_k3_avx2(sword16* vec1, sword16* vec2, sword16* poly,
4269    byte* seed)
4270{
4271    byte rand[4 * ETA2_RAND_SIZE];
4272
4273    mlkem_get_noise_x4_eta2_avx2(rand, seed, 0);
4274    mlkem_cbd_eta2_avx2(vec1              , rand + 0 * ETA2_RAND_SIZE);
4275    mlkem_cbd_eta2_avx2(vec1 + 1 * MLKEM_N, rand + 1 * ETA2_RAND_SIZE);
4276    mlkem_cbd_eta2_avx2(vec1 + 2 * MLKEM_N, rand + 2 * ETA2_RAND_SIZE);
4277    mlkem_cbd_eta2_avx2(vec2              , rand + 3 * ETA2_RAND_SIZE);
4278    mlkem_get_noise_x4_eta2_avx2(rand, seed, 4);
4279    mlkem_cbd_eta2_avx2(vec2 + 1 * MLKEM_N, rand + 0 * ETA2_RAND_SIZE);
4280    mlkem_cbd_eta2_avx2(vec2 + 2 * MLKEM_N, rand + 1 * ETA2_RAND_SIZE);
4281    if (poly != NULL) {
4282        mlkem_cbd_eta2_avx2(poly, rand + 2 * ETA2_RAND_SIZE);
4283    }
4284
4285    return 0;
4286}
4287#endif
4288
4289#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
4290/* Get the noise/error by calculating random bytes and sampling to a binomial
4291 * distribution.
4292 *
4293 * @param  [in, out]  prf   Pseudo-random function object.
4294 * @param  [out]      vec1  First Vector of polynomials.
4295 * @param  [out]      vec2  Second Vector of polynomials.
4296 * @param  [out]      poly  Polynomial.
4297 * @param  [in, out]  seed  Seed to use when calculating random.
4298 * @return  0 on success.
4299 */
4300static int mlkem_get_noise_k4_avx2(MLKEM_PRF_T* prf, sword16* vec1,
4301    sword16* vec2, sword16* poly, byte* seed)
4302{
4303    int ret = 0;
4304    byte rand[4 * ETA2_RAND_SIZE];
4305
4306    (void)prf;
4307
4308    mlkem_get_noise_x4_eta2_avx2(rand, seed, 0);
4309    mlkem_cbd_eta2_avx2(vec1              , rand + 0 * ETA2_RAND_SIZE);
4310    mlkem_cbd_eta2_avx2(vec1 + 1 * MLKEM_N, rand + 1 * ETA2_RAND_SIZE);
4311    mlkem_cbd_eta2_avx2(vec1 + 2 * MLKEM_N, rand + 2 * ETA2_RAND_SIZE);
4312    mlkem_cbd_eta2_avx2(vec1 + 3 * MLKEM_N, rand + 3 * ETA2_RAND_SIZE);
4313    mlkem_get_noise_x4_eta2_avx2(rand, seed, 4);
4314    mlkem_cbd_eta2_avx2(vec2              , rand + 0 * ETA2_RAND_SIZE);
4315    mlkem_cbd_eta2_avx2(vec2 + 1 * MLKEM_N, rand + 1 * ETA2_RAND_SIZE);
4316    mlkem_cbd_eta2_avx2(vec2 + 2 * MLKEM_N, rand + 2 * ETA2_RAND_SIZE);
4317    mlkem_cbd_eta2_avx2(vec2 + 3 * MLKEM_N, rand + 3 * ETA2_RAND_SIZE);
4318    if (poly != NULL) {
4319        seed[WC_ML_KEM_SYM_SZ] = 8;
4320        ret = mlkem_get_noise_eta2_avx2(prf, poly, seed);
4321    }
4322
4323    return ret;
4324}
4325#endif
4326#endif /* USE_INTEL_SPEEDUP */
4327
4328#if defined(__aarch64__) && defined(WOLFSSL_ARMASM)
4329
4330#define PRF_RAND_SZ   (2 * SHA3_256_BYTES)
4331
4332/* Get the noise/error by calculating random bytes.
4333 *
4334 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
4335 *   ...
4336 *  14:     e1[i] <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4337 *   ...
4338 *  17:     e2 <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4339 *   ...
4340 *
4341 * @param  [out]  rand  Random number byte array.
4342 * @param  [in]   seed  Seed to generate random from.
4343 * @param  [in]   o     Offset of seed count.
4344 */
4345static void mlkem_get_noise_x3_eta2_aarch64(byte* rand, byte* seed, byte o)
4346{
4347    word64* state = (word64*)rand;
4348
4349    state[0*25 + 4] = 0x1f00 + 0 + o;
4350    state[1*25 + 4] = 0x1f00 + 1 + o;
4351    state[2*25 + 4] = 0x1f00 + 2 + o;
4352
4353    mlkem_shake256_blocksx3_seed_neon(state, seed);
4354}
4355
4356#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
4357/* Get the noise/error by calculating random bytes.
4358 *
4359 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
4360 *   ...
4361 *   9:     s[i] <- SamplePolyCBD_eta_1(PRF_eta_1(sigma, N))
4362 *   ...
4363 *  13:     e[i] <- SamplePolyCBD_eta_1(PRF_eta_1(sigma, N))
4364 *   ...
4365 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
4366 *   ...
4367 *  10:     y[i] <- SamplePolyCBD_eta_1(PRF_eta_1(r, N))
4368 *   ...
4369 *
4370 * @param  [out]  rand  Random number byte array.
4371 * @param  [in]   seed  Seed to generate random from.
4372 * @param  [in]   o     Offset of seed count.
4373 */
4374static void mlkem_get_noise_x3_eta3_aarch64(byte* rand, byte* seed, byte o)
4375{
4376    word64 state[3 * 25];
4377
4378    state[0*25 + 4] = 0x1f00 + 0 + o;
4379    state[1*25 + 4] = 0x1f00 + 1 + o;
4380    state[2*25 + 4] = 0x1f00 + 2 + o;
4381
4382    mlkem_shake256_blocksx3_seed_neon(state, seed);
4383    XMEMCPY(rand + 0 * ETA3_RAND_SIZE, state + 0*25, SHA3_256_BYTES);
4384    XMEMCPY(rand + 1 * ETA3_RAND_SIZE, state + 1*25, SHA3_256_BYTES);
4385    XMEMCPY(rand + 2 * ETA3_RAND_SIZE, state + 2*25, SHA3_256_BYTES);
4386    mlkem_sha3_blocksx3_neon(state);
4387    rand += SHA3_256_BYTES;
4388    XMEMCPY(rand + 0 * ETA3_RAND_SIZE, state + 0*25,
4389        ETA3_RAND_SIZE - SHA3_256_BYTES);
4390    XMEMCPY(rand + 1 * ETA3_RAND_SIZE, state + 1*25,
4391        ETA3_RAND_SIZE - SHA3_256_BYTES);
4392    XMEMCPY(rand + 2 * ETA3_RAND_SIZE, state + 2*25,
4393        ETA3_RAND_SIZE - SHA3_256_BYTES);
4394}
4395
4396/* Get the noise/error by calculating random bytes.
4397 *
4398 * FIPS 203, Algorithm 13: K-PKE.KeyGen(d)
4399 *   ...
4400 *  13:     e[i] <- SamplePolyCBD_eta_1(PRF_eta_1(sigma, N))
4401 *   ...
4402 *
4403 * @param  [out]  rand  Random number byte array.
4404 * @param  [in]   seed  Seed to generate random from.
4405 * @param  [in]   o     Offset of seed count.
4406 */
4407static void mlkem_get_noise_eta3_aarch64(byte* rand, byte* seed, byte o)
4408{
4409    word64 state[25];
4410
4411    state[0] = ((word64*)seed)[0];
4412    state[1] = ((word64*)seed)[1];
4413    state[2] = ((word64*)seed)[2];
4414    state[3] = ((word64*)seed)[3];
4415    state[4] = 0x1f00 + o;
4416    XMEMSET(state + 5, 0, sizeof(*state) * (25 - 5));
4417    state[16] = W64LIT(0x8000000000000000);
4418    BlockSha3(state);
4419    XMEMCPY(rand                 , state, SHA3_256_BYTES);
4420    BlockSha3(state);
4421    XMEMCPY(rand + SHA3_256_BYTES, state, ETA3_RAND_SIZE - SHA3_256_BYTES);
4422}
4423
4424/* Get the noise/error by calculating random bytes and sampling to a binomial
4425 * distribution.
4426 *
4427 * @param  [out]      vec1  First Vector of polynomials.
4428 * @param  [out]      vec2  Second Vector of polynomials.
4429 * @param  [out]      poly  Polynomial.
4430 * @param  [in]       seed  Seed to use when calculating random.
4431 * @return  0 on success.
4432 */
4433static int mlkem_get_noise_k2_aarch64(sword16* vec1, sword16* vec2,
4434    sword16* poly, byte* seed)
4435{
4436    int ret = 0;
4437    byte rand[3 * 25 * 8];
4438
4439    mlkem_get_noise_x3_eta3_aarch64(rand, seed, 0);
4440    mlkem_cbd_eta3(vec1          , rand + 0 * ETA3_RAND_SIZE);
4441    mlkem_cbd_eta3(vec1 + MLKEM_N, rand + 1 * ETA3_RAND_SIZE);
4442    if (poly == NULL) {
4443        mlkem_cbd_eta3(vec2          , rand + 2 * ETA3_RAND_SIZE);
4444        mlkem_get_noise_eta3_aarch64(rand, seed, 3);
4445        mlkem_cbd_eta3(vec2 + MLKEM_N, rand                     );
4446    }
4447    else {
4448        mlkem_get_noise_x3_eta2_aarch64(rand, seed, 2);
4449        mlkem_cbd_eta2(vec2          , rand + 0 * 25 * 8);
4450        mlkem_cbd_eta2(vec2 + MLKEM_N, rand + 1 * 25 * 8);
4451        mlkem_cbd_eta2(poly          , rand + 2 * 25 * 8);
4452    }
4453
4454    return ret;
4455}
4456#endif
4457
4458#if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
4459/* Get the noise/error by calculating random bytes.
4460 *
4461 * FIPS 203, Algorithm 14: K-PKE.Encrypt(ek_PKE,m,r)
4462 *   ...
4463 *  14:     e1[i] <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4464 *   ...
4465 *  17:     e2 <- SamplePolyCBD_eta_2(PRF_eta_2(r, N))
4466 *   ...
4467 *
4468 * @param  [out]  rand  Random number byte array.
4469 * @param  [in]   seed  Seed to generate random from.
4470 * @param  [in]   o     Offset of seed count.
4471 */
4472static void mlkem_get_noise_eta2_aarch64(byte* rand, byte* seed, byte o)
4473{
4474    word64* state = (word64*)rand;
4475
4476    state[0] = ((word64*)seed)[0];
4477    state[1] = ((word64*)seed)[1];
4478    state[2] = ((word64*)seed)[2];
4479    state[3] = ((word64*)seed)[3];
4480    /* Transposed value same as not. */
4481    state[4] = 0x1f00 + o;
4482    XMEMSET(state + 5, 0, sizeof(*state) * (25 - 5));
4483    state[16] = W64LIT(0x8000000000000000);
4484    BlockSha3(state);
4485}
4486
4487/* Get the noise/error by calculating random bytes and sampling to a binomial
4488 * distribution.
4489 *
4490 * @param  [out]      vec1  First Vector of polynomials.
4491 * @param  [out]      vec2  Second Vector of polynomials.
4492 * @param  [out]      poly  Polynomial.
4493 * @param  [in]       seed  Seed to use when calculating random.
4494 * @return  0 on success.
4495 */
4496static int mlkem_get_noise_k3_aarch64(sword16* vec1, sword16* vec2,
4497     sword16* poly, byte* seed)
4498{
4499    byte rand[3 * 25 * 8];
4500
4501    mlkem_get_noise_x3_eta2_aarch64(rand, seed, 0);
4502    mlkem_cbd_eta2(vec1              , rand + 0 * 25 * 8);
4503    mlkem_cbd_eta2(vec1 + 1 * MLKEM_N, rand + 1 * 25 * 8);
4504    mlkem_cbd_eta2(vec1 + 2 * MLKEM_N, rand + 2 * 25 * 8);
4505    mlkem_get_noise_x3_eta2_aarch64(rand, seed, 3);
4506    mlkem_cbd_eta2(vec2              , rand + 0 * 25 * 8);
4507    mlkem_cbd_eta2(vec2 + 1 * MLKEM_N, rand + 1 * 25 * 8);
4508    mlkem_cbd_eta2(vec2 + 2 * MLKEM_N, rand + 2 * 25 * 8);
4509    if (poly != NULL) {
4510        mlkem_get_noise_eta2_aarch64(rand, seed, 6);
4511        mlkem_cbd_eta2(poly              , rand + 0 * 25 * 8);
4512    }
4513
4514    return 0;
4515}
4516#endif
4517
4518#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
4519/* Get the noise/error by calculating random bytes and sampling to a binomial
4520 * distribution.
4521 *
4522 * @param  [out]      vec1  First Vector of polynomials.
4523 * @param  [out]      vec2  Second Vector of polynomials.
4524 * @param  [out]      poly  Polynomial.
4525 * @param  [in]       seed  Seed to use when calculating random.
4526 * @return  0 on success.
4527 */
4528static int mlkem_get_noise_k4_aarch64(sword16* vec1, sword16* vec2,
4529    sword16* poly, byte* seed)
4530{
4531    int ret = 0;
4532    byte rand[3 * 25 * 8];
4533
4534    mlkem_get_noise_x3_eta2_aarch64(rand, seed, 0);
4535    mlkem_cbd_eta2(vec1              , rand + 0 * 25 * 8);
4536    mlkem_cbd_eta2(vec1 + 1 * MLKEM_N, rand + 1 * 25 * 8);
4537    mlkem_cbd_eta2(vec1 + 2 * MLKEM_N, rand + 2 * 25 * 8);
4538    mlkem_get_noise_x3_eta2_aarch64(rand, seed, 3);
4539    mlkem_cbd_eta2(vec1 + 3 * MLKEM_N, rand + 0 * 25 * 8);
4540    mlkem_cbd_eta2(vec2              , rand + 1 * 25 * 8);
4541    mlkem_cbd_eta2(vec2 + 1 * MLKEM_N, rand + 2 * 25 * 8);
4542    mlkem_get_noise_x3_eta2_aarch64(rand, seed, 6);
4543    mlkem_cbd_eta2(vec2 + 2 * MLKEM_N, rand + 0 * 25 * 8);
4544    mlkem_cbd_eta2(vec2 + 3 * MLKEM_N, rand + 1 * 25 * 8);
4545    if (poly != NULL) {
4546        mlkem_cbd_eta2(poly,               rand + 2 * 25 * 8);
4547    }
4548
4549    return ret;
4550}
4551#endif
4552#endif /* __aarch64__ && WOLFSSL_ARMASM */
4553
4554#if !(defined(__aarch64__) && defined(WOLFSSL_ARMASM))
4555
4556/* Get the noise/error by calculating random bytes and sampling to a binomial
4557 * distribution.
4558 *
4559 * @param  [in, out]  prf   Pseudo-random function object.
4560 * @param  [in]       k     Number of polynomials in vector.
4561 * @param  [out]      vec1  First Vector of polynomials.
4562 * @param  [in]       eta1  Size of noise/error integers with first vector.
4563 * @param  [out]      vec2  Second Vector of polynomials.
4564 * @param  [in]       eta2  Size of noise/error integers with second vector.
4565 * @param  [out]      poly  Polynomial.
4566 * @param  [in, out]  seed  Seed to use when calculating random.
4567 * @return  0 on success.
4568 */
4569static int mlkem_get_noise_c(MLKEM_PRF_T* prf, int k, sword16* vec1, int eta1,
4570    sword16* vec2, int eta2, sword16* poly, byte* seed)
4571{
4572    int ret = 0;
4573    int i;
4574
4575    /* First noise generation has a seed with 0x00 appended. */
4576    seed[WC_ML_KEM_SYM_SZ] = 0;
4577    /* Generate noise as private key. */
4578    for (i = 0; (ret == 0) && (i < k); i++) {
4579        /* Generate noise for each dimension of vector. */
4580        ret = mlkem_get_noise_eta1_c(prf, vec1 + i * MLKEM_N, seed, (byte)eta1);
4581        /* Increment value of appended byte. */
4582        seed[WC_ML_KEM_SYM_SZ]++;
4583    }
4584    if ((ret == 0) && (vec2 != NULL)) {
4585        /* Generate noise for error. */
4586        for (i = 0; (ret == 0) && (i < k); i++) {
4587            /* Generate noise for each dimension of vector. */
4588            ret = mlkem_get_noise_eta1_c(prf, vec2 + i * MLKEM_N, seed,
4589                (byte)eta2);
4590            /* Increment value of appended byte. */
4591            seed[WC_ML_KEM_SYM_SZ]++;
4592        }
4593    }
4594    else {
4595        seed[WC_ML_KEM_SYM_SZ] = (byte)(2 * k);
4596    }
4597    if ((ret == 0) && (poly != NULL)) {
4598        /* Generating random error polynomial. */
4599        ret = mlkem_get_noise_eta2_c(prf, poly, seed);
4600    }
4601
4602    return ret;
4603}
4604
4605#endif /* !(__aarch64__ && WOLFSSL_ARMASM) */
4606
4607/* Get the noise/error by calculating random bytes and sampling to a binomial
4608 * distribution.
4609 *
4610 * @param  [in, out]  prf   Pseudo-random function object.
4611 * @param  [in]       k     Number of polynomials in vector.
4612 * @param  [out]      vec1  First Vector of polynomials.
4613 * @param  [out]      vec2  Second Vector of polynomials.
4614 * @param  [out]      poly  Polynomial.
4615 * @param  [in, out]  seed  Seed to use when calculating random.
4616 * @return  0 on success.
4617 */
4618int mlkem_get_noise(MLKEM_PRF_T* prf, int k, sword16* vec1, sword16* vec2,
4619    sword16* poly, byte* seed)
4620{
4621    int ret;
4622
4623#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
4624    if (k == WC_ML_KEM_512_K) {
4625#if defined(WOLFSSL_ARMASM) && defined(__aarch64__)
4626        ret = mlkem_get_noise_k2_aarch64(vec1, vec2, poly, seed);
4627#else
4628    #if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
4629        if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
4630            ret = mlkem_get_noise_k2_avx2(prf, vec1, vec2, poly, seed);
4631            RESTORE_VECTOR_REGISTERS();
4632        }
4633        else
4634    #endif
4635        if (poly == NULL) {
4636            ret = mlkem_get_noise_c(prf, k, vec1, MLKEM_CBD_ETA3, vec2,
4637                MLKEM_CBD_ETA3, NULL, seed);
4638        }
4639        else {
4640            ret = mlkem_get_noise_c(prf, k, vec1, MLKEM_CBD_ETA3, vec2,
4641                MLKEM_CBD_ETA2, poly, seed);
4642        }
4643#endif
4644    }
4645    else
4646#endif
4647#if defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
4648    if (k == WC_ML_KEM_768_K) {
4649#if defined(WOLFSSL_ARMASM) && defined(__aarch64__)
4650        ret = mlkem_get_noise_k3_aarch64(vec1, vec2, poly, seed);
4651#else
4652    #if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
4653        if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
4654            ret = mlkem_get_noise_k3_avx2(vec1, vec2, poly, seed);
4655            RESTORE_VECTOR_REGISTERS();
4656        }
4657        else
4658    #endif
4659        {
4660            ret = mlkem_get_noise_c(prf, k, vec1, MLKEM_CBD_ETA2, vec2,
4661                MLKEM_CBD_ETA2, poly, seed);
4662        }
4663#endif
4664    }
4665    else
4666#endif
4667#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
4668    if (k == WC_ML_KEM_1024_K) {
4669#if defined(WOLFSSL_ARMASM) && defined(__aarch64__)
4670        ret = mlkem_get_noise_k4_aarch64(vec1, vec2, poly, seed);
4671#else
4672    #if defined(USE_INTEL_SPEEDUP) && !defined(WC_SHA3_NO_ASM)
4673        if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
4674            ret = mlkem_get_noise_k4_avx2(prf, vec1, vec2, poly, seed);
4675            RESTORE_VECTOR_REGISTERS();
4676        }
4677        else
4678    #endif
4679        {
4680            ret = mlkem_get_noise_c(prf, k, vec1, MLKEM_CBD_ETA2, vec2,
4681                MLKEM_CBD_ETA2, poly, seed);
4682        }
4683#endif
4684    }
4685    else
4686#endif
4687    {
4688        ret = BAD_STATE_E;
4689    }
4690
4691    (void)prf;
4692
4693    return ret;
4694}
4695
4696#if defined(WOLFSSL_MLKEM_MAKEKEY_SMALL_MEM) || \
4697    defined(WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM)
4698/* Get the noise/error by calculating random bytes and sampling to a binomial
4699 * distribution.
4700 *
4701 * @param  [in, out]  prf   Pseudo-random function object.
4702 * @param  [in]       k     Number of polynomials in vector.
4703 * @param  [out]      vec2  Second Vector of polynomials.
4704 * @param  [in, out]  seed  Seed to use when calculating random.
4705 * @param  [in]       i     Index of vector to generate.
4706 * @param  [in]       make  Indicates generation is for making a key.
4707 * @return  0 on success.
4708 */
4709static int mlkem_get_noise_i(MLKEM_PRF_T* prf, int k, sword16* vec2,
4710    byte* seed, int i, int make)
4711{
4712    int ret;
4713
4714    /* Initialize the PRF (generating matrix A leaves it in uninitialized
4715     * state). */
4716    mlkem_prf_init(prf);
4717
4718    /* Set index of polynomial of second vector into seed. */
4719    seed[WC_ML_KEM_SYM_SZ] = (byte)(k + i);
4720#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512)
4721    if ((k == WC_ML_KEM_512_K) && make) {
4722        ret = mlkem_get_noise_eta1_c(prf, vec2, seed, MLKEM_CBD_ETA3);
4723    }
4724    else
4725#endif
4726    {
4727        ret = mlkem_get_noise_eta1_c(prf, vec2, seed, MLKEM_CBD_ETA2);
4728    }
4729
4730    (void)make;
4731    return ret;
4732}
4733#endif
4734
4735/******************************************************************************/
4736
4737#if !(defined(__aarch64__) && defined(WOLFSSL_ARMASM))
4738/* Compare two byte arrays of equal size.
4739 *
4740 * @param [in]  a   First array to compare.
4741 * @param [in]  b   Second array to compare.
4742 * @param [in]  sz  Size of arrays in bytes.
4743 * @return  0 on success.
4744 * @return  -1 on failure.
4745 */
4746static int mlkem_cmp_c(const byte* a, const byte* b, int sz)
4747{
4748    int i;
4749    byte r = 0;
4750
4751    /* Constant time comparison of the encapsulated message and cipher text. */
4752    for (i = 0; i < sz; i++) {
4753        r |= a[i] ^ b[i];
4754    }
4755    return (int)(0 - ((-(word32)r) >> 31));
4756}
4757#endif
4758
4759/* Compare two byte arrays of equal size.
4760 *
4761 * @param [in]  a   First array to compare.
4762 * @param [in]  b   Second array to compare.
4763 * @param [in]  sz  Size of arrays in bytes.
4764 * @return  0 on success.
4765 * @return  -1 on failure.
4766 */
4767int mlkem_cmp(const byte* a, const byte* b, int sz)
4768{
4769#if defined(__aarch64__) && defined(WOLFSSL_ARMASM)
4770    return mlkem_cmp_neon(a, b, sz);
4771#else
4772    int fail;
4773
4774#ifdef USE_INTEL_SPEEDUP
4775    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
4776        fail = mlkem_cmp_avx2(a, b, sz);
4777        RESTORE_VECTOR_REGISTERS();
4778    }
4779    else
4780#endif
4781    {
4782        fail = mlkem_cmp_c(a, b, sz);
4783    }
4784
4785    return fail;
4786#endif
4787}
4788
4789/******************************************************************************/
4790
4791#if !defined(WOLFSSL_ARMASM)
4792
4793/* Conditional subtraction of q to each coefficient of a polynomial.
4794 *
4795 * FIPS 203, Section 4.2.1, Compression and decompression
4796 *
4797 * @param  [in, out]  p  Polynomial.
4798 */
4799static MLKEM_NOINLINE void mlkem_csubq_c(sword16* p)
4800{
4801    unsigned int i;
4802
4803    for (i = 0; i < MLKEM_N; ++i) {
4804        sword16 t = (sword16)(p[i] - MLKEM_Q);
4805        /* When top bit set, -ve number - need to add q back. */
4806        p[i] = (sword16)(((word16)(-((word16)t >> 15)) & MLKEM_Q) +
4807            (word16)t);
4808    }
4809}
4810
4811#elif defined(__aarch64__)
4812
4813/* Conditional subtraction of q to each coefficient of a polynomial.
4814 *
4815 * FIPS 203, Section 4.2.1, Compression and decompression
4816 *
4817 * @param  [in, out]  p  Polynomial.
4818 */
4819#define mlkem_csubq_c   mlkem_csubq_neon
4820
4821#elif defined(WOLFSSL_ARMASM_THUMB2)
4822
4823/* Conditional subtraction of q to each coefficient of a polynomial.
4824 *
4825 * FIPS 203, Section 4.2.1, Compression and decompression
4826 *
4827 * @param  [in, out]  p  Polynomial.
4828 */
4829#define mlkem_csubq_c   mlkem_thumb2_csubq
4830
4831#else
4832
4833/* Conditional subtraction of q to each coefficient of a polynomial.
4834 *
4835 * FIPS 203, Section 4.2.1, Compression and decompression
4836 *
4837 * @param  [in, out]  p  Polynomial.
4838 */
4839#define mlkem_csubq_c   mlkem_arm32_csubq
4840
4841#endif
4842
4843/******************************************************************************/
4844
4845#if defined(CONV_WITH_DIV) || !defined(WORD64_AVAILABLE)
4846
4847/* Compress value.
4848 *
4849 * Uses div operator that may be slow and not constant-time.
4850 *
4851 * FIPS 203, Section 4.2.1, Compression and decompression
4852 *
4853 * @param  [in]  v  Vector of polynomials.
4854 * @param  [in]  i  Index of polynomial in vector.
4855 * @param  [in]  j  Index into polynomial.
4856 * @param  [in]  k  Offset from indices.
4857 * @param  [in]  s  Shift amount to apply to value being compressed.
4858 * @param  [in]  m  Mask to apply get the required number of bits.
4859 * @return  Compressed value.
4860 */
4861#define TO_COMP_WORD_VEC(v, i, j, k, s, m) \
4862    ((((word32)v[i * MLKEM_N + j + k] << s) + MLKEM_Q_HALF) / MLKEM_Q) & m
4863
4864/* Compress value to 10 bits.
4865 *
4866 * Uses div operator that may be slow and not constant-time.
4867 *
4868 * FIPS 203, Section 4.2.1, Compression and decompression
4869 *
4870 * @param  [in]  v  Vector of polynomials.
4871 * @param  [in]  i  Index of polynomial in vector.
4872 * @param  [in]  j  Index into polynomial.
4873 * @param  [in]  k  Offset from indices.
4874 * @return  Compressed value.
4875 */
4876#define TO_COMP_WORD_10(v, i, j, k) \
4877    TO_COMP_WORD_VEC(v, i, j, k, 10, 0x3ff)
4878
4879/* Compress value to 11 bits.
4880 *
4881 * Uses div operator that may be slow and not constant-time.
4882 *
4883 * FIPS 203, Section 4.2.1, Compression and decompression
4884 *
4885 * @param  [in]  v  Vector of polynomials.
4886 * @param  [in]  i  Index of polynomial in vector.
4887 * @param  [in]  j  Index into polynomial.
4888 * @param  [in]  k  Offset from indices.
4889 * @return  Compressed value.
4890 */
4891#define TO_COMP_WORD_11(v, i, j, k) \
4892    TO_COMP_WORD_VEC(v, i, j, k, 11, 0x7ff)
4893
4894#else
4895
4896/* Multiplier that does div q.
4897 * ((1 << 53) + MLKEM_Q_HALF) / MLKEM_Q
4898 */
4899#define MLKEM_V53         0x275f6ed0176UL
4900/* Multiplier times half of q.
4901 * MLKEM_V53 * (MLKEM_Q_HALF + 1)
4902 */
4903#define MLKEM_V53_HALF    0x10013afb768076UL
4904
4905/* Multiplier that does div q.
4906 * ((1 << 54) + MLKEM_Q_HALF) / MLKEM_Q
4907 */
4908#define MLKEM_V54         0x4ebedda02ecUL
4909/* Multiplier times half of q.
4910 * MLKEM_V54 * (MLKEM_Q_HALF + 1)
4911 */
4912#define MLKEM_V54_HALF    0x200275f6ed00ecUL
4913
4914/* Compress value to 10 bits.
4915 *
4916 * Uses mul instead of div.
4917 *
4918 * FIPS 203, Section 4.2.1, Compression and decompression
4919 *
4920 * @param  [in]  v  Vector of polynomials.
4921 * @param  [in]  i  Index of polynomial in vector.
4922 * @param  [in]  j  Index into polynomial.
4923 * @param  [in]  k  Offset from indices.
4924 * @return  Compressed value.
4925 */
4926#define TO_COMP_WORD_10(v, i, j, k) \
4927    (sword16)((((MLKEM_V54 << 10) * (word64)(v)[(i) * MLKEM_N + (j) + (k)]) + \
4928               MLKEM_V54_HALF) >> 54)
4929
4930/* Compress value to 11 bits.
4931 *
4932 * Uses mul instead of div.
4933 * Only works for values in range: 0..3228
4934 *
4935 * FIPS 203, Section 4.2.1, Compression and decompression
4936 *
4937 * @param  [in]  v  Vector of polynomials.
4938 * @param  [in]  i  Index of polynomial in vector.
4939 * @param  [in]  j  Index into polynomial.
4940 * @param  [in]  k  Offset from indices.
4941 * @return  Compressed value.
4942 */
4943#define TO_COMP_WORD_11(v, i, j, k) \
4944    (sword16)((((MLKEM_V53 << 11) * (word64)(v)[(i) * MLKEM_N + (j) + (k)]) + \
4945               MLKEM_V53_HALF) >> 53)
4946
4947#endif /* CONV_WITH_DIV */
4948
4949#if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
4950    !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
4951#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512) || \
4952    defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
4953/* Compress the vector of polynomials into a byte array with 10 bits each.
4954 *
4955 * FIPS 203, Section 4.2.1, Compression and decompression
4956 *
4957 * @param  [out]      r  Array of bytes.
4958 * @param  [in, out]  v  Vector of polynomials.
4959 * @param  [in]       k  Number of polynomials in vector.
4960 */
4961static void mlkem_vec_compress_10_c(byte* r, sword16* v, unsigned int k)
4962{
4963    unsigned int i;
4964    unsigned int j;
4965
4966    for (i = 0; i < k; i++) {
4967        /* Reduce each coefficient to mod q. */
4968        mlkem_csubq_c(v + i * MLKEM_N);
4969        /* All values are now positive. */
4970    }
4971
4972    /* Each polynomial. */
4973    for (i = 0; i < k; i++) {
4974#if defined(WOLFSSL_SMALL_STACK) || defined(WOLFSSL_MLKEM_NO_LARGE_CODE) || \
4975    defined(BIG_ENDIAN_ORDER)
4976        /* Each 4 polynomial coefficients. */
4977        for (j = 0; j < MLKEM_N; j += 4) {
4978        #ifdef WOLFSSL_MLKEM_SMALL
4979            unsigned int l;
4980            sword16 t[4];
4981            /* Compress four polynomial values to 10 bits each. */
4982            for (l = 0; l < 4; l++) {
4983                t[l] = TO_COMP_WORD_10(v, i, j, l);
4984            }
4985
4986            /* Pack four 10-bit values into byte array. */
4987            r[ 0] = (t[0] >> 0);
4988            r[ 1] = (t[0] >> 8) | (t[1] << 2);
4989            r[ 2] = (t[1] >> 6) | (t[2] << 4);
4990            r[ 3] = (t[2] >> 4) | (t[3] << 6);
4991            r[ 4] = (t[3] >> 2);
4992        #else
4993            /* Compress four polynomial values to 10 bits each. */
4994            sword16 t0 = TO_COMP_WORD_10(v, i, j, 0);
4995            sword16 t1 = TO_COMP_WORD_10(v, i, j, 1);
4996            sword16 t2 = TO_COMP_WORD_10(v, i, j, 2);
4997            sword16 t3 = TO_COMP_WORD_10(v, i, j, 3);
4998
4999            /* Pack four 10-bit values into byte array. */
5000            r[ 0] = (byte)( t0 >> 0);
5001            r[ 1] = (byte)((t0 >> 8) | (t1 << 2));
5002            r[ 2] = (byte)((t1 >> 6) | (t2 << 4));
5003            r[ 3] = (byte)((t2 >> 4) | (t3 << 6));
5004            r[ 4] = (byte)( t3 >> 2);
5005        #endif
5006
5007            /* Move over set bytes. */
5008            r += 5;
5009        }
5010#else
5011        /* Each 16 polynomial coefficients. */
5012        for (j = 0; j < MLKEM_N; j += 16) {
5013            /* Compress four polynomial values to 10 bits each. */
5014            sword16 t0  = TO_COMP_WORD_10(v, i, j, 0);
5015            sword16 t1  = TO_COMP_WORD_10(v, i, j, 1);
5016            sword16 t2  = TO_COMP_WORD_10(v, i, j, 2);
5017            sword16 t3  = TO_COMP_WORD_10(v, i, j, 3);
5018            sword16 t4  = TO_COMP_WORD_10(v, i, j, 4);
5019            sword16 t5  = TO_COMP_WORD_10(v, i, j, 5);
5020            sword16 t6  = TO_COMP_WORD_10(v, i, j, 6);
5021            sword16 t7  = TO_COMP_WORD_10(v, i, j, 7);
5022            sword16 t8  = TO_COMP_WORD_10(v, i, j, 8);
5023            sword16 t9  = TO_COMP_WORD_10(v, i, j, 9);
5024            sword16 t10 = TO_COMP_WORD_10(v, i, j, 10);
5025            sword16 t11 = TO_COMP_WORD_10(v, i, j, 11);
5026            sword16 t12 = TO_COMP_WORD_10(v, i, j, 12);
5027            sword16 t13 = TO_COMP_WORD_10(v, i, j, 13);
5028            sword16 t14 = TO_COMP_WORD_10(v, i, j, 14);
5029            sword16 t15 = TO_COMP_WORD_10(v, i, j, 15);
5030
5031            word32* r32 = (word32*)r;
5032            /* Pack sixteen 10-bit values into byte array. */
5033            r32[0] =  (word32)t0         | ((word32)t1  << 10) |
5034                     ((word32)t2  << 20) | ((word32)t3  << 30);
5035            r32[1] = ((word32)t3  >>  2) | ((word32)t4  <<  8) |
5036                     ((word32)t5  << 18) | ((word32)t6  << 28);
5037            r32[2] = ((word32)t6  >>  4) | ((word32)t7  <<  6) |
5038                     ((word32)t8  << 16) | ((word32)t9  << 26);
5039            r32[3] = ((word32)t9  >>  6) | ((word32)t10 <<  4) |
5040                     ((word32)t11 << 14) | ((word32)t12 << 24);
5041            r32[4] = ((word32)t12 >>  8) | ((word32)t13 <<  2) |
5042                     ((word32)t14 << 12) | ((word32)t15 << 22);
5043
5044            /* Move over set bytes. */
5045            r += 20;
5046        }
5047#endif
5048    }
5049}
5050
5051/* Compress the vector of polynomials into a byte array with 10 bits each.
5052 *
5053 * FIPS 203, Section 4.2.1, Compression and decompression
5054 *
5055 * @param  [out]      r  Array of bytes.
5056 * @param  [in, out]  v  Vector of polynomials.
5057 * @param  [in]       k  Number of polynomials in vector.
5058 */
5059void mlkem_vec_compress_10(byte* r, sword16* v, unsigned int k)
5060{
5061#ifdef USE_INTEL_SPEEDUP
5062    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5063        mlkem_compress_10_avx2(r, v, (int)k);
5064        RESTORE_VECTOR_REGISTERS();
5065    }
5066    else
5067#endif
5068    {
5069        mlkem_vec_compress_10_c(r, v, k);
5070    }
5071}
5072#endif
5073
5074#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
5075/* Compress the vector of polynomials into a byte array with 11 bits each.
5076 *
5077 * FIPS 203, Section 4.2.1, Compression and decompression
5078 *
5079 * @param  [out]      r  Array of bytes.
5080 * @param  [in, out]  v  Vector of polynomials.
5081 */
5082static void mlkem_vec_compress_11_c(byte* r, sword16* v)
5083{
5084    unsigned int i;
5085    unsigned int j;
5086#ifdef WOLFSSL_MLKEM_SMALL
5087    unsigned int k;
5088#endif
5089
5090    for (i = 0; i < 4; i++) {
5091        /* Reduce each coefficient to mod q. */
5092        mlkem_csubq_c(v + i * MLKEM_N);
5093        /* All values are now positive. */
5094    }
5095
5096    /* Each polynomial. */
5097    for (i = 0; i < 4; i++) {
5098        /* Each 8 polynomial coefficients. */
5099        for (j = 0; j < MLKEM_N; j += 8) {
5100        #ifdef WOLFSSL_MLKEM_SMALL
5101            sword16 t[8];
5102            /* Compress eight polynomial values to 11 bits each. */
5103            for (k = 0; k < 8; k++) {
5104                t[k] = TO_COMP_WORD_11(v, i, j, k);
5105            }
5106
5107            /* Pack eight 11-bit values into byte array. */
5108            r[ 0] = (byte)( t[0] >>  0);
5109            r[ 1] = (byte)((t[0] >>  8) | (t[1] << 3));
5110            r[ 2] = (byte)((t[1] >>  5) | (t[2] << 6));
5111            r[ 3] = (byte)( t[2] >>  2);
5112            r[ 4] = (byte)((t[2] >> 10) | (t[3] << 1));
5113            r[ 5] = (byte)((t[3] >>  7) | (t[4] << 4));
5114            r[ 6] = (byte)((t[4] >>  4) | (t[5] << 7));
5115            r[ 7] = (byte)( t[5] >>  1);
5116            r[ 8] = (byte)((t[5] >>  9) | (t[6] << 2));
5117            r[ 9] = (byte)((t[6] >>  6) | (t[7] << 5));
5118            r[10] = (byte)( t[7] >>  3);
5119        #else
5120            /* Compress eight polynomial values to 11 bits each. */
5121            sword16 t0 = TO_COMP_WORD_11(v, i, j, 0);
5122            sword16 t1 = TO_COMP_WORD_11(v, i, j, 1);
5123            sword16 t2 = TO_COMP_WORD_11(v, i, j, 2);
5124            sword16 t3 = TO_COMP_WORD_11(v, i, j, 3);
5125            sword16 t4 = TO_COMP_WORD_11(v, i, j, 4);
5126            sword16 t5 = TO_COMP_WORD_11(v, i, j, 5);
5127            sword16 t6 = TO_COMP_WORD_11(v, i, j, 6);
5128            sword16 t7 = TO_COMP_WORD_11(v, i, j, 7);
5129
5130            /* Pack eight 11-bit values into byte array. */
5131            r[ 0] = (byte)( t0 >>  0);
5132            r[ 1] = (byte)((t0 >>  8) | (t1 << 3));
5133            r[ 2] = (byte)((t1 >>  5) | (t2 << 6));
5134            r[ 3] = (byte)( t2 >>  2);
5135            r[ 4] = (byte)((t2 >> 10) | (t3 << 1));
5136            r[ 5] = (byte)((t3 >>  7) | (t4 << 4));
5137            r[ 6] = (byte)((t4 >>  4) | (t5 << 7));
5138            r[ 7] = (byte)( t5 >>  1);
5139            r[ 8] = (byte)((t5 >>  9) | (t6 << 2));
5140            r[ 9] = (byte)((t6 >>  6) | (t7 << 5));
5141            r[10] = (byte)( t7 >>  3);
5142        #endif
5143
5144            /* Move over set bytes. */
5145            r += 11;
5146        }
5147    }
5148}
5149
5150/* Compress the vector of polynomials into a byte array with 11 bits each.
5151 *
5152 * FIPS 203, Section 4.2.1, Compression and decompression
5153 *
5154 * @param  [out]      r  Array of bytes.
5155 * @param  [in, out]  v  Vector of polynomials.
5156 */
5157void mlkem_vec_compress_11(byte* r, sword16* v)
5158{
5159#ifdef USE_INTEL_SPEEDUP
5160    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5161        mlkem_compress_11_avx2(r, v, 4);
5162        RESTORE_VECTOR_REGISTERS();
5163    }
5164    else
5165#endif
5166    {
5167        mlkem_vec_compress_11_c(r, v);
5168    }
5169}
5170#endif
5171#endif /* !WOLFSSL_MLKEM_NO_ENCAPSULATE || !WOLFSSL_MLKEM_NO_DECAPSULATE */
5172
5173#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE
5174/* Decompress a 10 bit value.
5175 *
5176 * FIPS 203, Section 4.2.1, Compression and decompression
5177 *
5178 * @param  [out]  v  Vector of polynomials.
5179 * @param  [in]   i  Index of polynomial in vector.
5180 * @param  [in]   j  Index into polynomial.
5181 * @param  [in]   k  Offset from indices.
5182 * @param  [in]   t  Value to decompress.
5183 */
5184#define DECOMP_10(v, i, j, k, t) \
5185    v[(i) * MLKEM_N + 4 * (j) + (k)] = \
5186        (sword16)((((word32)((t) & 0x3ff) * MLKEM_Q) + 512) >> 10)
5187
5188/* Decompress an 11 bit value.
5189 *
5190 * FIPS 203, Section 4.2.1, Compression and decompression
5191 *
5192 * @param  [out]  v  Vector of polynomials.
5193 * @param  [in]   i  Index of polynomial in vector.
5194 * @param  [in]   j  Index into polynomial.
5195 * @param  [in]   k  Offset from indices.
5196 * @param  [in]   t  Value to decompress.
5197 */
5198#define DECOMP_11(v, i, j, k, t) \
5199    v[(i) * MLKEM_N + 8 * (j) + (k)] = \
5200        (sword16)((((word32)((t) & 0x7ff) * MLKEM_Q) + 1024) >> 11)
5201
5202#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512) || \
5203    defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
5204/* Decompress the byte array of packed 10 bits into vector of polynomials.
5205 *
5206 * FIPS 203, Section 4.2.1, Compression and decompression
5207 *
5208 * @param  [out]  v  Vector of polynomials.
5209 * @param  [in]   b  Array of bytes.
5210 * @param  [in]   k  Number of polynomials in vector.
5211 */
5212static void mlkem_vec_decompress_10_c(sword16* v, const byte* b, unsigned int k)
5213{
5214    unsigned int i;
5215    unsigned int j;
5216#ifdef WOLFSSL_MLKEM_SMALL
5217    unsigned int l;
5218#endif
5219
5220    /* Each polynomial. */
5221    for (i = 0; i < k; i++) {
5222        /* Each 4 polynomial coefficients. */
5223        for (j = 0; j < MLKEM_N / 4; j++) {
5224        #ifdef WOLFSSL_MLKEM_SMALL
5225            word16 t[4];
5226            /* Extract out 4 values of 10 bits each. */
5227            t[0] = (word16)((b[0] >> 0) | ((word16)b[ 1] << 8));
5228            t[1] = (word16)((b[1] >> 2) | ((word16)b[ 2] << 6));
5229            t[2] = (word16)((b[2] >> 4) | ((word16)b[ 3] << 4));
5230            t[3] = (word16)((b[3] >> 6) | ((word16)b[ 4] << 2));
5231            b += 5;
5232
5233            /* Decompress 4 values. */
5234            for (l = 0; l < 4; l++) {
5235                DECOMP_10(v, i, j, l, t[l]);
5236            }
5237        #else
5238            /* Extract out 4 values of 10 bits each. */
5239            word16 t0 = (word16)((b[0] >> 0) | ((word16)b[ 1] << 8));
5240            word16 t1 = (word16)((b[1] >> 2) | ((word16)b[ 2] << 6));
5241            word16 t2 = (word16)((b[2] >> 4) | ((word16)b[ 3] << 4));
5242            word16 t3 = (word16)((b[3] >> 6) | ((word16)b[ 4] << 2));
5243            b += 5;
5244
5245            /* Decompress 4 values. */
5246            DECOMP_10(v, i, j, 0, t0);
5247            DECOMP_10(v, i, j, 1, t1);
5248            DECOMP_10(v, i, j, 2, t2);
5249            DECOMP_10(v, i, j, 3, t3);
5250        #endif
5251        }
5252    }
5253}
5254
5255/* Decompress the byte array of packed 10 bits into vector of polynomials.
5256 *
5257 * FIPS 203, Section 4.2.1, Compression and decompression
5258 *
5259 * @param  [out]  v  Vector of polynomials.
5260 * @param  [in]   b  Array of bytes.
5261 * @param  [in]   k  Number of polynomials in vector.
5262 */
5263void mlkem_vec_decompress_10(sword16* v, const byte* b, unsigned int k)
5264{
5265#ifdef USE_INTEL_SPEEDUP
5266    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5267        mlkem_decompress_10_avx2(v, b, (int)k);
5268        RESTORE_VECTOR_REGISTERS();
5269    }
5270    else
5271#endif
5272    {
5273        mlkem_vec_decompress_10_c(v, b, k);
5274    }
5275}
5276#endif
5277#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
5278/* Decompress the byte array of packed 11 bits into vector of polynomials.
5279 *
5280 * FIPS 203, Section 4.2.1, Compression and decompression
5281 *
5282 * @param  [out]  v       Vector of polynomials.
5283 * @param  [in]   b       Array of bytes.
5284 */
5285static void mlkem_vec_decompress_11_c(sword16* v, const byte* b)
5286{
5287    unsigned int i;
5288    unsigned int j;
5289#ifdef WOLFSSL_MLKEM_SMALL
5290    unsigned int l;
5291#endif
5292
5293    /* Each polynomial. */
5294    for (i = 0; i < 4; i++) {
5295        /* Each 8 polynomial coefficients. */
5296        for (j = 0; j < MLKEM_N / 8; j++) {
5297        #ifdef WOLFSSL_MLKEM_SMALL
5298            word16 t[8];
5299            /* Extract out 8 values of 11 bits each. */
5300            t[0] = (word16)((b[0] >> 0) | ((word16)b[ 1] << 8));
5301            t[1] = (word16)((b[1] >> 3) | ((word16)b[ 2] << 5));
5302            t[2] = (word16)((b[2] >> 6) | ((word16)b[ 3] << 2) |
5303                   ((word16)b[4] << 10));
5304            t[3] = (word16)((b[4] >> 1) | ((word16)b[ 5] << 7));
5305            t[4] = (word16)((b[5] >> 4) | ((word16)b[ 6] << 4));
5306            t[5] = (word16)((b[6] >> 7) | ((word16)b[ 7] << 1) |
5307                   ((word16)b[8] <<  9));
5308            t[6] = (word16)((b[8] >> 2) | ((word16)b[ 9] << 6));
5309            t[7] = (word16)((b[9] >> 5) | ((word16)b[10] << 3));
5310            b += 11;
5311
5312            /* Decompress 8 values. */
5313            for (l = 0; l < 8; l++) {
5314                DECOMP_11(v, i, j, l, t[l]);
5315            }
5316        #else
5317            /* Extract out 8 values of 11 bits each. */
5318            word16 t0 = (word16)((b[0] >> 0) | ((word16)b[ 1] << 8));
5319            word16 t1 = (word16)((b[1] >> 3) | ((word16)b[ 2] << 5));
5320            word16 t2 = (word16)((b[2] >> 6) | ((word16)b[ 3] << 2) |
5321                   ((word16)b[4] << 10));
5322            word16 t3 = (word16)((b[4] >> 1) | ((word16)b[ 5] << 7));
5323            word16 t4 = (word16)((b[5] >> 4) | ((word16)b[ 6] << 4));
5324            word16 t5 = (word16)((b[6] >> 7) | ((word16)b[ 7] << 1) |
5325                   ((word16)b[8] <<  9));
5326            word16 t6 = (word16)((b[8] >> 2) | ((word16)b[ 9] << 6));
5327            word16 t7 = (word16)((b[9] >> 5) | ((word16)b[10] << 3));
5328            b += 11;
5329
5330            /* Decompress 8 values. */
5331            DECOMP_11(v, i, j, 0, t0);
5332            DECOMP_11(v, i, j, 1, t1);
5333            DECOMP_11(v, i, j, 2, t2);
5334            DECOMP_11(v, i, j, 3, t3);
5335            DECOMP_11(v, i, j, 4, t4);
5336            DECOMP_11(v, i, j, 5, t5);
5337            DECOMP_11(v, i, j, 6, t6);
5338            DECOMP_11(v, i, j, 7, t7);
5339        #endif
5340        }
5341    }
5342}
5343
5344/* Decompress the byte array of packed 11 bits into vector of polynomials.
5345 *
5346 * FIPS 203, Section 4.2.1, Compression and decompression
5347 *
5348 * @param  [out]  v       Vector of polynomials.
5349 * @param  [in]   b       Array of bytes.
5350 */
5351void mlkem_vec_decompress_11(sword16* v, const byte* b)
5352{
5353#ifdef USE_INTEL_SPEEDUP
5354    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5355        mlkem_decompress_11_avx2(v, b, 4);
5356        RESTORE_VECTOR_REGISTERS();
5357    }
5358    else
5359#endif
5360    {
5361        mlkem_vec_decompress_11_c(v, b);
5362    }
5363}
5364#endif
5365#endif /* !WOLFSSL_MLKEM_NO_DECAPSULATE */
5366
5367#ifdef CONV_WITH_DIV
5368
5369/* Compress value.
5370 *
5371 * Uses div operator that may be slow and not constant-time.
5372 *
5373 * FIPS 203, Section 4.2.1, Compression and decompression
5374 *
5375 * @param  [in]  v  Vector of polynomials.
5376 * @param  [in]  i  Index into polynomial.
5377 * @param  [in]  j  Offset from indices.
5378 * @param  [in]  s  Shift amount to apply to value being compressed.
5379 * @param  [in]  m  Mask to apply to get the required number of bits.
5380 * @return  Compressed value.
5381 */
5382#define TO_COMP_WORD(v, i, j, s, m) \
5383    ((((word32)v[i + j] << s) + MLKEM_Q_HALF) / MLKEM_Q) & m
5384
5385/* Compress value to 4 bits.
5386 *
5387 * Uses div operator that may be slow and not constant-time.
5388 *
5389 * FIPS 203, Section 4.2.1, Compression and decompression
5390 *
5391 * @param  [in]  p  Polynomial.
5392 * @param  [in]  i  Index into polynomial.
5393 * @param  [in]  j  Offset from indices.
5394 * @return  Compressed value.
5395 */
5396#define TO_COMP_WORD_4(p, i, j) \
5397    TO_COMP_WORD(p, i, j, 4, 0xf)
5398
5399/* Compress value to 5 bits.
5400 *
5401 * Uses div operator that may be slow and not constant-time.
5402 *
5403 * FIPS 203, Section 4.2.1, Compression and decompression
5404 *
5405 * @param  [in]  p  Polynomial.
5406 * @param  [in]  i  Index into polynomial.
5407 * @param  [in]  j  Offset from indices.
5408 * @return  Compressed value.
5409 */
5410#define TO_COMP_WORD_5(p, i, j) \
5411    TO_COMP_WORD(p, i, j, 5, 0x1f)
5412
5413#else
5414
5415/* Multiplier that does div q. */
5416#define MLKEM_V28         ((word32)(((1U << 28) + MLKEM_Q_HALF)) / MLKEM_Q)
5417/* Multiplier times half of q plus one. */
5418#define MLKEM_V28_HALF    ((word32)(MLKEM_V28 * (MLKEM_Q_HALF + 1)))
5419
5420/* Multiplier that does div q. */
5421#define MLKEM_V27         ((word32)(((1U << 27) + MLKEM_Q_HALF)) / MLKEM_Q)
5422/* Multiplier times half of q. */
5423#define MLKEM_V27_HALF    ((word32)(MLKEM_V27 * MLKEM_Q_HALF))
5424
5425/* Compress value to 4 bits.
5426 *
5427 * Uses mul instead of div.
5428 *
5429 * FIPS 203, Section 4.2.1, Compression and decompression
5430 *
5431 * @param  [in]  p  Polynomial.
5432 * @param  [in]  i  Index into polynomial.
5433 * @param  [in]  j  Offset from indices.
5434 * @return  Compressed value.
5435 */
5436#define TO_COMP_WORD_4(p, i, j) \
5437    (byte)((((MLKEM_V28 << 4) * (word32)(p)[(i) + (j)]) + MLKEM_V28_HALF) >> 28)
5438
5439/* Compress value to 5 bits.
5440 *
5441 * Uses mul instead of div.
5442 *
5443 * FIPS 203, Section 4.2.1, Compression and decompression
5444 *
5445 * @param  [in]  p  Polynomial.
5446 * @param  [in]  i  Index into polynomial.
5447 * @param  [in]  j  Offset from indices.
5448 * @return  Compressed value.
5449 */
5450#define TO_COMP_WORD_5(p, i, j) \
5451    (byte)((((MLKEM_V27 << 5) * (word32)(p)[(i) + (j)]) + MLKEM_V27_HALF) >> 27)
5452
5453#endif /* CONV_WITH_DIV */
5454
5455#if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
5456    !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
5457#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512) || \
5458    defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
5459/* Compress a polynomial into byte array with coefficients of 4 bits.
5460 *
5461 * FIPS 203, Section 4.2.1, Compression and decompression
5462 *
5463 * @param  [out]      b  Array of bytes.
5464 * @param  [in, out]  p  Polynomial.
5465 */
5466static void mlkem_compress_4_c(byte* b, sword16* p)
5467{
5468    unsigned int i;
5469#ifdef WOLFSSL_MLKEM_SMALL
5470    unsigned int j;
5471    byte t[8];
5472#endif
5473
5474    /* Reduce each coefficient to mod q. */
5475    mlkem_csubq_c(p);
5476    /* All values are now positive. */
5477
5478    /* Each 8 polynomial coefficients. */
5479    for (i = 0; i < MLKEM_N; i += 8) {
5480    #ifdef WOLFSSL_MLKEM_SMALL
5481        /* Compress eight polynomial values to 4 bits each. */
5482        for (j = 0; j < 8; j++) {
5483            t[j] = TO_COMP_WORD_4(p, i, j);
5484        }
5485
5486        b[0] = (byte)(t[0] | (t[1] << 4));
5487        b[1] = (byte)(t[2] | (t[3] << 4));
5488        b[2] = (byte)(t[4] | (t[5] << 4));
5489        b[3] = (byte)(t[6] | (t[7] << 4));
5490    #else
5491        /* Compress eight polynomial values to 4 bits each. */
5492        byte t0 = TO_COMP_WORD_4(p, i, 0);
5493        byte t1 = TO_COMP_WORD_4(p, i, 1);
5494        byte t2 = TO_COMP_WORD_4(p, i, 2);
5495        byte t3 = TO_COMP_WORD_4(p, i, 3);
5496        byte t4 = TO_COMP_WORD_4(p, i, 4);
5497        byte t5 = TO_COMP_WORD_4(p, i, 5);
5498        byte t6 = TO_COMP_WORD_4(p, i, 6);
5499        byte t7 = TO_COMP_WORD_4(p, i, 7);
5500
5501        /* Pack eight 4-bit values into byte array. */
5502        b[0] = (byte)(t0 | (t1 << 4));
5503        b[1] = (byte)(t2 | (t3 << 4));
5504        b[2] = (byte)(t4 | (t5 << 4));
5505        b[3] = (byte)(t6 | (t7 << 4));
5506    #endif
5507
5508        /* Move over set bytes. */
5509        b += 4;
5510    }
5511}
5512
5513/* Compress a polynomial into byte array with coefficients of 4 bits.
5514 *
5515 * FIPS 203, Section 4.2.1, Compression and decompression
5516 *
5517 * @param  [out]      b  Array of bytes.
5518 * @param  [in, out]  p  Polynomial.
5519 */
5520void mlkem_compress_4(byte* b, sword16* p)
5521{
5522#ifdef USE_INTEL_SPEEDUP
5523    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5524        mlkem_compress_4_avx2(b, p);
5525        RESTORE_VECTOR_REGISTERS();
5526    }
5527    else
5528#endif
5529    {
5530        mlkem_compress_4_c(b, p);
5531    }
5532}
5533#endif
5534#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
5535/* Compress a polynomial into byte array with coefficients of 5 bits.
5536 *
5537 * FIPS 203, Section 4.2.1, Compression and decompression
5538 *
5539 * @param  [out]      b  Array of bytes.
5540 * @param  [in, out]  p  Polynomial.
5541 */
5542static void mlkem_compress_5_c(byte* b, sword16* p)
5543{
5544    unsigned int i;
5545#ifdef WOLFSSL_MLKEM_SMALL
5546    unsigned int j;
5547    byte t[8];
5548#endif
5549
5550    /* Reduce each coefficient to mod q. */
5551    mlkem_csubq_c(p);
5552    /* All values are now positive. */
5553
5554    for (i = 0; i < MLKEM_N; i += 8) {
5555    #ifdef WOLFSSL_MLKEM_SMALL
5556        /* Compress eight polynomial values to 5 bits each. */
5557        for (j = 0; j < 8; j++) {
5558            t[j] = TO_COMP_WORD_5(p, i, j);
5559        }
5560
5561        /* Pack 5 bits into byte array. */
5562        b[0] = (byte)((t[0] >> 0) | (t[1] << 5));
5563        b[1] = (byte)((t[1] >> 3) | (t[2] << 2) | (t[3] << 7));
5564        b[2] = (byte)((t[3] >> 1) | (t[4] << 4));
5565        b[3] = (byte)((t[4] >> 4) | (t[5] << 1) | (t[6] << 6));
5566        b[4] = (byte)((t[6] >> 2) | (t[7] << 3));
5567    #else
5568        /* Compress eight polynomial values to 5 bits each. */
5569        byte t0 = TO_COMP_WORD_5(p, i, 0);
5570        byte t1 = TO_COMP_WORD_5(p, i, 1);
5571        byte t2 = TO_COMP_WORD_5(p, i, 2);
5572        byte t3 = TO_COMP_WORD_5(p, i, 3);
5573        byte t4 = TO_COMP_WORD_5(p, i, 4);
5574        byte t5 = TO_COMP_WORD_5(p, i, 5);
5575        byte t6 = TO_COMP_WORD_5(p, i, 6);
5576        byte t7 = TO_COMP_WORD_5(p, i, 7);
5577
5578        /* Pack eight 5-bit values into byte array. */
5579        b[0] = (byte)((t0 >> 0) | (t1 << 5));
5580        b[1] = (byte)((t1 >> 3) | (t2 << 2) | (t3 << 7));
5581        b[2] = (byte)((t3 >> 1) | (t4 << 4));
5582        b[3] = (byte)((t4 >> 4) | (t5 << 1) | (t6 << 6));
5583        b[4] = (byte)((t6 >> 2) | (t7 << 3));
5584    #endif
5585
5586        /* Move over set bytes. */
5587        b += 5;
5588    }
5589}
5590
5591/* Compress a polynomial into byte array with coefficients of 5 bits.
5592 *
5593 * FIPS 203, Section 4.2.1, Compression and decompression
5594 *
5595 * @param  [out]      b  Array of bytes.
5596 * @param  [in, out]  p  Polynomial.
5597 */
5598void mlkem_compress_5(byte* b, sword16* p)
5599{
5600#ifdef USE_INTEL_SPEEDUP
5601    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5602        mlkem_compress_5_avx2(b, p);
5603        RESTORE_VECTOR_REGISTERS();
5604    }
5605    else
5606#endif
5607    {
5608        mlkem_compress_5_c(b, p);
5609    }
5610}
5611#endif
5612#endif /* !WOLFSSL_MLKEM_NO_ENCAPSULATE || !WOLFSSL_MLKEM_NO_DECAPSULATE */
5613
5614#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE
5615/* Decompress a 4 bit value.
5616 *
5617 * FIPS 203, Section 4.2.1, Compression and decompression
5618 *
5619 * @param  [out]  p  Polynomial.
5620 * @param  [in]   i  Index into polynomial.
5621 * @param  [in]   j  Offset from indices.
5622 * @param  [in]   t  Value to decompress.
5623 */
5624#define DECOMP_4(p, i, j, t) \
5625    p[(i) + (j)] = (sword16)(((word16)((t) * MLKEM_Q) + 8) >> 4)
5626
5627/* Decompress a 5 bit value.
5628 *
5629 * FIPS 203, Section 4.2.1, Compression and decompression
5630 *
5631 * @param  [out]  p  Polynomial.
5632 * @param  [in]   i  Index into polynomial.
5633 * @param  [in]   j  Offset from indices.
5634 * @param  [in]   t  Value to decompress.
5635 */
5636#define DECOMP_5(p, i, j, t) \
5637    p[(i) + (j)] = (sword16)((((word32)((t) & 0x1f) * MLKEM_Q) + 16) >> 5)
5638
5639#if defined(WOLFSSL_KYBER512) || defined(WOLFSSL_WC_ML_KEM_512) || \
5640    defined(WOLFSSL_KYBER768) || defined(WOLFSSL_WC_ML_KEM_768)
5641/* Decompress the byte array of packed 4 bits into polynomial.
5642 *
5643 * FIPS 203, Section 4.2.1, Compression and decompression
5644 *
5645 * @param  [out]  p       Polynomial.
5646 * @param  [in]   b       Array of bytes.
5647 */
5648static void mlkem_decompress_4_c(sword16* p, const byte* b)
5649{
5650    unsigned int i;
5651
5652    /* 2 coefficients at a time. */
5653    for (i = 0; i < MLKEM_N; i += 2) {
5654        /* 2 coefficients decompressed from one byte. */
5655        DECOMP_4(p, i, 0, b[0] & 0xf);
5656        DECOMP_4(p, i, 1, b[0] >>  4);
5657        b += 1;
5658    }
5659}
5660
5661/* Decompress the byte array of packed 4 bits into polynomial.
5662 *
5663 * FIPS 203, Section 4.2.1, Compression and decompression
5664 *
5665 * @param  [out]  p       Polynomial.
5666 * @param  [in]   b       Array of bytes.
5667 */
5668void mlkem_decompress_4(sword16* p, const byte* b)
5669{
5670#ifdef USE_INTEL_SPEEDUP
5671    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5672        mlkem_decompress_4_avx2(p, b);
5673        RESTORE_VECTOR_REGISTERS();
5674    }
5675    else
5676#endif
5677    {
5678        mlkem_decompress_4_c(p, b);
5679    }
5680}
5681#endif
5682#if defined(WOLFSSL_KYBER1024) || defined(WOLFSSL_WC_ML_KEM_1024)
5683/* Decompress the byte array of packed 5 bits into polynomial.
5684 *
5685 * FIPS 203, Section 4.2.1, Compression and decompression
5686 *
5687 * @param  [out]  p       Polynomial.
5688 * @param  [in]   b       Array of bytes.
5689 */
5690static void mlkem_decompress_5_c(sword16* p, const byte* b)
5691{
5692    unsigned int i;
5693
5694    /* Each 8 polynomial coefficients. */
5695    for (i = 0; i < MLKEM_N; i += 8) {
5696    #ifdef WOLFSSL_MLKEM_SMALL
5697        unsigned int j;
5698        byte t[8];
5699
5700        /* Extract out 8 values of 5 bits each. */
5701        t[0] = (b[0] >> 0);
5702        t[1] = (byte)((b[0] >> 5) | (b[1] << 3));
5703        t[2] = (b[1] >> 2);
5704        t[3] = (byte)((b[1] >> 7) | (b[2] << 1));
5705        t[4] = (byte)((b[2] >> 4) | (b[3] << 4));
5706        t[5] = (b[3] >> 1);
5707        t[6] = (byte)((b[3] >> 6) | (b[4] << 2));
5708        t[7] = (b[4] >> 3);
5709        b += 5;
5710
5711        /* Decompress 8 values. */
5712        for (j = 0; j < 8; j++) {
5713            DECOMP_5(p, i, j, t[j]);
5714        }
5715    #else
5716        /* Extract out 8 values of 5 bits each. */
5717        byte t0 = (b[0] >> 0);
5718        byte t1 = (byte)((b[0] >> 5) | (b[1] << 3));
5719        byte t2 = (b[1] >> 2);
5720        byte t3 = (byte)((b[1] >> 7) | (b[2] << 1));
5721        byte t4 = (byte)((b[2] >> 4) | (b[3] << 4));
5722        byte t5 = (b[3] >> 1);
5723        byte t6 = (byte)((b[3] >> 6) | (b[4] << 2));
5724        byte t7 = (b[4] >> 3);
5725        b += 5;
5726
5727        /* Decompress 8 values. */
5728        DECOMP_5(p, i, 0, t0);
5729        DECOMP_5(p, i, 1, t1);
5730        DECOMP_5(p, i, 2, t2);
5731        DECOMP_5(p, i, 3, t3);
5732        DECOMP_5(p, i, 4, t4);
5733        DECOMP_5(p, i, 5, t5);
5734        DECOMP_5(p, i, 6, t6);
5735        DECOMP_5(p, i, 7, t7);
5736    #endif
5737    }
5738}
5739
5740/* Decompress the byte array of packed 5 bits into polynomial.
5741 *
5742 * FIPS 203, Section 4.2.1, Compression and decompression
5743 *
5744 * @param  [out]  p       Polynomial.
5745 * @param  [in]   b       Array of bytes.
5746 */
5747void mlkem_decompress_5(sword16* p, const byte* b)
5748{
5749#ifdef USE_INTEL_SPEEDUP
5750    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5751        mlkem_decompress_5_avx2(p, b);
5752        RESTORE_VECTOR_REGISTERS();
5753    }
5754    else
5755#endif
5756    {
5757        mlkem_decompress_5_c(p, b);
5758    }
5759}
5760#endif
5761#endif /* !WOLFSSL_MLKEM_NO_DECAPSULATE */
5762
5763/******************************************************************************/
5764
5765#if !(defined(__aarch64__) && defined(WOLFSSL_ARMASM))
5766#if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
5767    !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
5768/* Convert bit from byte to 0 or (MLKEM_Q + 1) / 2.
5769 *
5770 * Constant time implementation.
5771 * XOR in wc_mlkem_opt_blocker() to ensure optimizer doesn't know what will be
5772 * ANDed with MLKEM_Q_1_HALF and can't optimize to non-constant time code.
5773 *
5774 * FIPS 203, Algorithm 6: ByteDecode_d(B)
5775 *
5776 * @param  [out]  p    Polynomial to hold converted value.
5777 * @param  [in]   msg  Message to get bit from byte.
5778 * @param  [in]   i    Index of byte from message.
5779 * @param  [in]   j    Index of bit in byte.
5780 */
5781#define FROM_MSG_BIT(p, msg, i, j) \
5782    ((p)[8 * (i) + (j)] = (((sword16)0 - (sword16)(((msg)[i] >> (j)) & 1)) ^ \
5783                          wc_mlkem_opt_blocker()) & MLKEM_Q_1_HALF)
5784
5785/* Convert message to polynomial.
5786 *
5787 * FIPS 203, Algorithm 6: ByteDecode_d(B)
5788 *
5789 * @param  [out]  p    Polynomial.
5790 * @param  [in]   msg  Message as a byte array.
5791 */
5792static void mlkem_from_msg_c(sword16* p, const byte* msg)
5793{
5794    unsigned int i;
5795
5796    /* For each byte of the message. */
5797    for (i = 0; i < MLKEM_N / 8; i++) {
5798    #ifdef WOLFSSL_MLKEM_SMALL
5799        unsigned int j;
5800        /* For each bit of the message. */
5801        for (j = 0; j < 8; j++) {
5802            FROM_MSG_BIT(p, msg, i, j);
5803        }
5804    #else
5805        FROM_MSG_BIT(p, msg, i, 0);
5806        FROM_MSG_BIT(p, msg, i, 1);
5807        FROM_MSG_BIT(p, msg, i, 2);
5808        FROM_MSG_BIT(p, msg, i, 3);
5809        FROM_MSG_BIT(p, msg, i, 4);
5810        FROM_MSG_BIT(p, msg, i, 5);
5811        FROM_MSG_BIT(p, msg, i, 6);
5812        FROM_MSG_BIT(p, msg, i, 7);
5813    #endif
5814    }
5815}
5816
5817/* Convert message to polynomial.
5818 *
5819 * FIPS 203, Algorithm 6: ByteDecode_d(B)
5820 *
5821 * @param  [out]  p    Polynomial.
5822 * @param  [in]   msg  Message as a byte array.
5823 */
5824void mlkem_from_msg(sword16* p, const byte* msg)
5825{
5826#ifdef USE_INTEL_SPEEDUP
5827    if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5828        mlkem_from_msg_avx2(p, msg);
5829        RESTORE_VECTOR_REGISTERS();
5830    }
5831    else
5832#endif
5833    {
5834        mlkem_from_msg_c(p, msg);
5835    }
5836}
5837#endif
5838
5839#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE
5840#ifdef CONV_WITH_DIV
5841
5842/* Convert value to bit.
5843 *
5844 * Uses div operator that may be slow.
5845 *
5846 * FIPS 203, Algorithm 5: ByteEncode_d(F)
5847 *
5848 * @param  [in, out]  m   Message.
5849 * @param  [in]       p   Polynomial.
5850 * @param  [in]       i   Index of byte in message.
5851 * @param  [in]       j   Index of bit in byte.
5852 */
5853#define TO_MSG_BIT(m, p, i, j) \
5854    m[i] |= (((((sword16)p[8 * i + j] << 1) + MLKEM_Q_HALF) / MLKEM_Q) & 1) << j
5855
5856#else
5857
5858/* Multiplier that does div q. */
5859#define MLKEM_V31       (((1U << 31) + (MLKEM_Q / 2)) / MLKEM_Q)
5860/* 2 * multiplier that does div q. Only need bit 32 of result. */
5861#define MLKEM_V31_2     ((word32)(MLKEM_V31 * 2))
5862/* Multiplier times half of q. */
5863#define MLKEM_V31_HALF    ((word32)(MLKEM_V31 * MLKEM_Q_HALF))
5864
5865/* Convert value to bit.
5866 *
5867 * Uses mul instead of div.
5868 *
5869 * FIPS 203, Algorithm 5: ByteEncode_d(F)
5870 *
5871 * @param  [in, out]  m   Message.
5872 * @param  [in]       p   Polynomial.
5873 * @param  [in]       i   Index of byte in message.
5874 * @param  [in]       j   Index of bit in byte.
5875 */
5876#define TO_MSG_BIT(m, p, i, j) \
5877    (m)[i] |= (byte)((((MLKEM_V31_2 * (word16)(p)[8 * (i) + (j)]) + \
5878                       MLKEM_V31_HALF) >> 31) << (j))
5879
5880#endif /* CONV_WITH_DIV */
5881
5882/* Convert polynomial to message.
5883 *
5884 * FIPS 203, Algorithm 5: ByteEncode_d(F)
5885 *
5886 * @param  [out]      msg  Message as a byte array.
5887 * @param  [in, out]  p    Polynomial.
5888 */
5889static void mlkem_to_msg_c(byte* msg, sword16* p)
5890{
5891    unsigned int i;
5892
5893    /* Reduce each coefficient to mod q. */
5894    mlkem_csubq_c(p);
5895    /* All values are now in range. */
5896
5897    for (i = 0; i < MLKEM_N / 8; i++) {
5898    #ifdef WOLFSSL_MLKEM_SMALL
5899        unsigned int j;
5900        msg[i] = 0;
5901        for (j = 0; j < 8; j++) {
5902            TO_MSG_BIT(msg, p, i, j);
5903        }
5904    #else
5905        msg[i] = 0;
5906        TO_MSG_BIT(msg, p, i, 0);
5907        TO_MSG_BIT(msg, p, i, 1);
5908        TO_MSG_BIT(msg, p, i, 2);
5909        TO_MSG_BIT(msg, p, i, 3);
5910        TO_MSG_BIT(msg, p, i, 4);
5911        TO_MSG_BIT(msg, p, i, 5);
5912        TO_MSG_BIT(msg, p, i, 6);
5913        TO_MSG_BIT(msg, p, i, 7);
5914    #endif
5915    }
5916}
5917
5918/* Convert polynomial to message.
5919 *
5920 * FIPS 203, Algorithm 5: ByteEncode_d(F)
5921 *
5922 * @param  [out]      msg  Message as a byte array.
5923 * @param  [in, out]  p    Polynomial.
5924 */
5925void mlkem_to_msg(byte* msg, sword16* p)
5926{
5927#ifdef USE_INTEL_SPEEDUP
5928     if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
5929        /* Convert the polynomial into an array of bytes (message). */
5930        mlkem_to_msg_avx2(msg, p);
5931        RESTORE_VECTOR_REGISTERS();
5932    }
5933    else
5934#endif
5935    {
5936        mlkem_to_msg_c(msg, p);
5937    }
5938}
5939#endif /* !WOLFSSL_MLKEM_NO_DECAPSULATE */
5940#else
5941#if !defined(WOLFSSL_MLKEM_NO_ENCAPSULATE) || \
5942    !defined(WOLFSSL_MLKEM_NO_DECAPSULATE)
5943/* Convert message to polynomial.
5944 *
5945 * FIPS 203, Algorithm 6: ByteDecode_d(B)
5946 *
5947 * @param  [out]  p    Polynomial.
5948 * @param  [in]   msg  Message as a byte array.
5949 */
5950void mlkem_from_msg(sword16* p, const byte* msg)
5951{
5952    mlkem_from_msg_neon(p, msg);
5953}
5954#endif /* !WOLFSSL_MLKEM_NO_ENCAPSULATE || !WOLFSSL_MLKEM_NO_DECAPSULATE */
5955
5956#ifndef WOLFSSL_MLKEM_NO_DECAPSULATE
5957/* Convert polynomial to message.
5958 *
5959 * FIPS 203, Algorithm 5: ByteEncode_d(F)
5960 *
5961 * @param  [out]      msg  Message as a byte array.
5962 * @param  [in, out]  p    Polynomial.
5963 */
5964void mlkem_to_msg(byte* msg, sword16* p)
5965{
5966    mlkem_to_msg_neon(msg, p);
5967}
5968#endif /* WOLFSSL_MLKEM_NO_DECAPSULATE */
5969#endif /* !(__aarch64__ && WOLFSSL_ARMASM) */
5970
5971/******************************************************************************/
5972
5973/* Convert bytes to polynomial.
5974 *
5975 * Consecutive 12 bits hold each coefficient of polynomial.
5976 * Used in decoding private and public keys.
5977 *
5978 * FIPS 203, Algorithm 6: ByteDecode_d(B)
5979 *
5980 * @param  [out]  p  Vector of polynomials.
5981 * @param  [in]   b  Array of bytes.
5982 * @param  [in]   k  Number of polynomials in vector.
5983 */
5984static void mlkem_from_bytes_c(sword16* p, const byte* b, int k)
5985{
5986    int i;
5987    int j;
5988
5989    for (j = 0; j < k; j++) {
5990        for (i = 0; i < MLKEM_N / 2; i++) {
5991            p[2 * i + 0] = ((b[3 * i + 0] >> 0) |
5992                            ((word16)b[3 * i + 1] << 8)) & 0xfff;
5993            p[2 * i + 1] = ((b[3 * i + 1] >> 4) |
5994                            ((word16)b[3 * i + 2] << 4)) & 0xfff;
5995        }
5996        p += MLKEM_N;
5997        b += WC_ML_KEM_POLY_SIZE;
5998    }
5999}
6000
6001/* Convert bytes to polynomial.
6002 *
6003 * Consecutive 12 bits hold each coefficient of polynomial.
6004 * Used in decoding private and public keys.
6005 *
6006 * FIPS 203, Algorithm 6: ByteDecode_d(B)
6007 *
6008 * @param  [out]  p  Vector of polynomials.
6009 * @param  [in]   b  Array of bytes.
6010 * @param  [in]   k  Number of polynomials in vector.
6011 */
6012void mlkem_from_bytes(sword16* p, const byte* b, int k)
6013{
6014#ifdef USE_INTEL_SPEEDUP
6015     if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
6016        int i;
6017
6018        for (i = 0; i < k; i++) {
6019            mlkem_from_bytes_avx2(p, b);
6020            p += MLKEM_N;
6021            b += WC_ML_KEM_POLY_SIZE;
6022        }
6023
6024        RESTORE_VECTOR_REGISTERS();
6025    }
6026    else
6027#endif
6028    {
6029        mlkem_from_bytes_c(p, b, k);
6030    }
6031}
6032
6033/* Convert polynomial to bytes.
6034 *
6035 * Consecutive 12 bits hold each coefficient of polynomial.
6036 * Used in encoding private and public keys.
6037 *
6038 * FIPS 203, Algorithm 5: ByteEncode_d(F)
6039 *
6040 * @param  [out]      b  Array of bytes.
6041 * @param  [in, out]  p  Polynomial.
6042 * @param  [in]       k  Number of polynomials in vector.
6043 */
6044static void mlkem_to_bytes_c(byte* b, sword16* p, int k)
6045{
6046    int i;
6047    int j;
6048
6049    for (j = 0; j < k; j++) {
6050        /* Reduce each coefficient to mod q. */
6051        mlkem_csubq_c(p);
6052        /* All values are now positive. */
6053
6054        for (i = 0; i < MLKEM_N / 2; i++) {
6055            word16 t0 = (word16)p[2 * i];
6056            word16 t1 = (word16)p[2 * i + 1];
6057            b[3 * i + 0] = (byte)(t0 >> 0);
6058            b[3 * i + 1] = (byte)((t0 >> 8) | (t1 << 4));
6059            b[3 * i + 2] = (byte)(t1 >> 4);
6060        }
6061        p += MLKEM_N;
6062        b += WC_ML_KEM_POLY_SIZE;
6063    }
6064}
6065
6066/* Convert polynomial to bytes.
6067 *
6068 * Consecutive 12 bits hold each coefficient of polynomial.
6069 * Used in encoding private and public keys.
6070 *
6071 * FIPS 203, Algorithm 5: ByteEncode_d(F)
6072 *
6073 * @param  [out]      b  Array of bytes.
6074 * @param  [in, out]  p  Polynomial.
6075 * @param  [in]       k  Number of polynomials in vector.
6076 */
6077void mlkem_to_bytes(byte* b, sword16* p, int k)
6078{
6079#ifdef USE_INTEL_SPEEDUP
6080     if (IS_INTEL_AVX2(cpuid_flags) && (SAVE_VECTOR_REGISTERS2() == 0)) {
6081        int i;
6082
6083        for (i = 0; i < k; i++) {
6084            mlkem_to_bytes_avx2(b, p);
6085            p += MLKEM_N;
6086            b += WC_ML_KEM_POLY_SIZE;
6087        }
6088
6089        RESTORE_VECTOR_REGISTERS();
6090    }
6091    else
6092#endif
6093    {
6094        mlkem_to_bytes_c(b, p, k);
6095    }
6096}
6097
6098/**
6099 * Check the public key values are smaller than the modulus.
6100 *
6101 * @param [in] p  Public key - vector.
6102 * @param [in] k  Number of polynomials in vector.
6103 * @return  0 when all values are in range.
6104 * @return  PUBLIC_KEY_E when at least one value is out of range.
6105 */
6106int mlkem_check_public(const sword16* p, int k)
6107{
6108    int ret = 0;
6109    int i;
6110
6111    for (i = 0; i < k * MLKEM_N; i++) {
6112        if (p[i] >= MLKEM_Q) {
6113            ret = PUBLIC_KEY_E;
6114            break;
6115        }
6116    }
6117
6118    return ret;
6119}
6120
6121#endif /* WOLFSSL_HAVE_MLKEM */