Define mlk_polyvec[_mulcache] as typedef for array

This commit changes the definition of `mlk_polyvec` from a struct
wrapper `struct { mlk_poly vec[MLKEM_K] }` to a typedef for
`mlk_poly[MLKEM_K]`. Similarly, it introduces `mlk_polymat`
as a typedef for `mlk_poly[MLKEM_K * MLKEM_K]`.

This ensures that when working with polynomial matrices
-- that is, arrays of `mlk_polyvec`s -- the polynomial
entries are consecutive in memory. This needs to be assumed
for the correctness of `gen_matrix()`.

From all we know, is a theoretical concern that without this
change, compilers would be free to introduce padding into
`mlk_polyvec`, breaking `gen_matrix()`. Yet, the behaviour of
`gen_matrix` is strictly speaking UB, so even if a compiler
does not use padding in `mlk_polyvec`, it is better to change it.

Defining `mlk_polymat` as a typedef for `mlk_poly[MLKEM_K * MLKEM_K]`
and not `mlk_polyvec[MLKEM_K]` is deliberate: In the latter case, one
would not be able to pass a `mlk_polymat` to a function expecting
`const polymat`, forcing cumbersome manual casts.

Signed-off-by: Hanno Becker <[email protected]>

Author: hanno-becker

examples/monolithic_build/mlkem_native_monobuild.c

@@ -181,6 +181,7 @@
 #undef mlk_poly_getnoise_eta1_4x
 #undef mlk_poly_getnoise_eta2
 #undef mlk_poly_getnoise_eta2_4x
+#undef mlk_polymat
 #undef mlk_polyvec
 #undef mlk_polyvec_add
 #undef mlk_polyvec_basemul_acc_montgomery_cached

mlkem/indcpa.c

@@ -36,15 +36,15 @@
  *              and the public seed used to generate the matrix A.
  *
  * Arguments:   uint8_t *r: pointer to the output serialized public key
- *              mlk_polyvec *pk: pointer to the input public-key mlk_polyvec.
+ *              mlk_polyvec pk: pointer to the input public-key mlk_polyvec.
  *                Must have coefficients within [0,..,q-1].
  *              const uint8_t *seed: pointer to the input public seed
  *
  * Specification:
  * Implements [FIPS 203, Algorithm 13 (K-PKE.KeyGen), L19]
  *
  **************************************************/
-static void mlk_pack_pk(uint8_t r[MLKEM_INDCPA_PUBLICKEYBYTES], mlk_polyvec *pk,
+static void mlk_pack_pk(uint8_t r[MLKEM_INDCPA_PUBLICKEYBYTES], mlk_polyvec pk,
                         const uint8_t seed[MLKEM_SYMBYTES])
 {
   mlk_assert_bound_2d(pk, MLKEM_K, MLKEM_N, 0, MLKEM_Q);
@@ -58,7 +58,7 @@ static void mlk_pack_pk(uint8_t r[MLKEM_INDCPA_PUBLICKEYBYTES], mlk_polyvec *pk,
  * Description: De-serialize public key from a byte array;
  *              approximate inverse of mlk_pack_pk
  *
- * Arguments:   - mlk_polyvec *pk: pointer to output public-key polynomial
+ * Arguments:   - mlk_polyvec pk: pointer to output public-key polynomial
  *                vector Coefficients will be normalized to [0,..,q-1].
  *              - uint8_t *seed: pointer to output seed to generate matrix A
  *              - const uint8_t *packedpk: pointer to input serialized public
@@ -68,7 +68,7 @@ static void mlk_pack_pk(uint8_t r[MLKEM_INDCPA_PUBLICKEYBYTES], mlk_polyvec *pk,
  * Implements [FIPS 203, Algorithm 14 (K-PKE.Encrypt), L2-3]
  *
  **************************************************/
-static void mlk_unpack_pk(mlk_polyvec *pk, uint8_t seed[MLKEM_SYMBYTES],
+static void mlk_unpack_pk(mlk_polyvec pk, uint8_t seed[MLKEM_SYMBYTES],
                           const uint8_t packedpk[MLKEM_INDCPA_PUBLICKEYBYTES])
 {
   mlk_polyvec_frombytes(pk, packedpk);
@@ -86,14 +86,14 @@ static void mlk_unpack_pk(mlk_polyvec *pk, uint8_t seed[MLKEM_SYMBYTES],
  * Description: Serialize the secret key
  *
  * Arguments:   - uint8_t *r: pointer to output serialized secret key
- *              - mlk_polyvec *sk: pointer to input vector of polynomials
+ *              - mlk_polyvec sk: pointer to input vector of polynomials
  *                (secret key)
  *
  * Specification:
  * Implements [FIPS 203, Algorithm 13 (K-PKE.KeyGen), L20]
  *
  **************************************************/
-static void mlk_pack_sk(uint8_t r[MLKEM_INDCPA_SECRETKEYBYTES], mlk_polyvec *sk)
+static void mlk_pack_sk(uint8_t r[MLKEM_INDCPA_SECRETKEYBYTES], mlk_polyvec sk)
 {
   mlk_assert_bound_2d(sk, MLKEM_K, MLKEM_N, 0, MLKEM_Q);
   mlk_polyvec_tobytes(r, sk);
@@ -104,7 +104,7 @@ static void mlk_pack_sk(uint8_t r[MLKEM_INDCPA_SECRETKEYBYTES], mlk_polyvec *sk)
  *
  * Description: De-serialize the secret key; inverse of mlk_pack_sk
  *
- * Arguments:   - mlk_polyvec *sk: pointer to output vector of polynomials
+ * Arguments:   - mlk_polyvec sk: pointer to output vector of polynomials
  *                (secret key)
  *              - const uint8_t *packedsk: pointer to input serialized secret
  *                key
@@ -113,7 +113,7 @@ static void mlk_pack_sk(uint8_t r[MLKEM_INDCPA_SECRETKEYBYTES], mlk_polyvec *sk)
  * Implements [FIPS 203, Algorithm 15 (K-PKE.Decrypt), L5]
  *
  **************************************************/
-static void mlk_unpack_sk(mlk_polyvec *sk,
+static void mlk_unpack_sk(mlk_polyvec sk,
                           const uint8_t packedsk[MLKEM_INDCPA_SECRETKEYBYTES])
 {
   mlk_polyvec_frombytes(sk, packedsk);
@@ -134,7 +134,7 @@ static void mlk_unpack_sk(mlk_polyvec *sk,
  * Implements [FIPS 203, Algorithm 14 (K-PKE.Encrypt), L22-23]
  *
  **************************************************/
-static void mlk_pack_ciphertext(uint8_t r[MLKEM_INDCPA_BYTES], mlk_polyvec *b,
+static void mlk_pack_ciphertext(uint8_t r[MLKEM_INDCPA_BYTES], mlk_polyvec b,
                                 mlk_poly *v)
 {
   mlk_polyvec_compress_du(r, b);
@@ -147,15 +147,15 @@ static void mlk_pack_ciphertext(uint8_t r[MLKEM_INDCPA_BYTES], mlk_polyvec *b,
  * Description: De-serialize and decompress ciphertext from a byte array;
  *              approximate inverse of mlk_pack_ciphertext
  *
- * Arguments:   - mlk_polyvec *b: pointer to the output vector of polynomials b
+ * Arguments:   - mlk_polyvec b: pointer to the output vector of polynomials b
  *              - mlk_poly *v: pointer to the output polynomial v
  *              - const uint8_t *c: pointer to the input serialized ciphertext
  *
  * Specification:
  * Implements [FIPS 203, Algorithm 15 (K-PKE.Decrypt), L1-4]
  *
  **************************************************/
-static void mlk_unpack_ciphertext(mlk_polyvec *b, mlk_poly *v,
+static void mlk_unpack_ciphertext(mlk_polyvec b, mlk_poly *v,
                                   const uint8_t c[MLKEM_INDCPA_BYTES])
 {
   mlk_polyvec_decompress_du(b, c);
@@ -186,7 +186,7 @@ __contract__(
  *
  * Not static for benchmarking */
 MLK_INTERNAL_API
-void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
+void mlk_gen_matrix(mlk_polymat a, const uint8_t seed[MLKEM_SYMBYTES],
                     int transposed)
 {
   unsigned i, j;
@@ -228,7 +228,7 @@ void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
      * This call writes across mlk_polyvec boundaries for K=2 and K=3.
      * This is intentional and safe.
      */
-    mlk_poly_rej_uniform_x4(&a[0].vec[0] + i, seed_ext);
+    mlk_poly_rej_uniform_x4(&a[i], seed_ext);
   }
 
   /* For MLKEM_K == 3, sample the last entry individually. */
@@ -249,7 +249,7 @@ void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
       seed_ext[0][MLKEM_SYMBYTES + 1] = x;
     }
 
-    mlk_poly_rej_uniform(&a[0].vec[0] + i, seed_ext[0]);
+    mlk_poly_rej_uniform(&a[i], seed_ext[0]);
     i++;
   }
 
@@ -259,12 +259,9 @@ void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
    * The public matrix is generated in NTT domain. If the native backend
    * uses a custom order in NTT domain, permute A accordingly.
    */
-  for (i = 0; i < MLKEM_K; i++)
+  for (i = 0; i < MLKEM_K * MLKEM_K; i++)
   {
-    for (j = 0; j < MLKEM_K; j++)
-    {
-      mlk_poly_permute_bitrev_to_custom(a[i].vec[j].coeffs);
-    }
+    mlk_poly_permute_bitrev_to_custom(a[i].coeffs);
   }
 
   /* Specification: Partially implements
@@ -278,27 +275,26 @@ void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
  * Description: Computes matrix-vector product in NTT domain,
  *              via Montgomery multiplication.
  *
- * Arguments:   - mlk_polyvec *out: Pointer to output polynomial vector
- *              - mlk_polyvec a[MLKEM_K]: Input matrix. Must be in NTT domain
+ * Arguments:   - mlk_polyvec out: Pointer to output polynomial vector
+ *              - mlk_polymat a: Input matrix. Must be in NTT domain
  *                  and have coefficients of absolute value < 4096.
- *              - mlk_polyvec *v: Input polynomial vector. Must be in NTT
+ *              - mlk_polyvec v: Input polynomial vector. Must be in NTT
  *                  domain.
- *              - mlk_polyvec *vc: Mulcache for v, computed via
+ *              - mlk_polyvec vc: Mulcache for v, computed via
  *                  mlk_polyvec_mulcache_compute().
  *
  * Specification: Implements [FIPS 203, Section 2.4.7, Eq (2.12), (2.13)]
  *
  **************************************************/
-static void mlk_matvec_mul(mlk_polyvec *out, const mlk_polyvec a[MLKEM_K],
-                           const mlk_polyvec *v, const mlk_polyvec_mulcache *vc)
+static void mlk_matvec_mul(mlk_polyvec out, const mlk_polymat a,
+                           const mlk_polyvec v, const mlk_polyvec_mulcache vc)
 __contract__(
   requires(memory_no_alias(out, sizeof(mlk_polyvec)))
-  requires(memory_no_alias(a, sizeof(mlk_polyvec) * MLKEM_K))
+  requires(memory_no_alias(a, sizeof(mlk_polymat)))
   requires(memory_no_alias(v, sizeof(mlk_polyvec)))
   requires(memory_no_alias(vc, sizeof(mlk_polyvec_mulcache)))
-  requires(forall(k0, 0, MLKEM_K,
-    forall(k1, 0, MLKEM_K,
-      array_bound(a[k0].vec[k1].coeffs, 0, MLKEM_N, 0, MLKEM_UINT12_LIMIT))))
+  requires(forall(k0, 0, MLKEM_K * MLKEM_K,
+    array_bound(a[k0].coeffs, 0, MLKEM_N, 0, MLKEM_UINT12_LIMIT)))
   assigns(object_whole(out)))
 {
   unsigned i;
@@ -307,7 +303,7 @@ __contract__(
     assigns(i, object_whole(out))
     invariant(i <= MLKEM_K))
   {
-    mlk_polyvec_basemul_acc_montgomery_cached(&out->vec[i], &a[i], v, vc);
+    mlk_polyvec_basemul_acc_montgomery_cached(&out[i], &a[MLKEM_K * i], v, vc);
   }
 }
 
@@ -327,7 +323,8 @@ void mlk_indcpa_keypair_derand(uint8_t pk[MLKEM_INDCPA_PUBLICKEYBYTES],
   MLK_ALIGN uint8_t buf[2 * MLKEM_SYMBYTES];
   const uint8_t *publicseed = buf;
   const uint8_t *noiseseed = buf + MLKEM_SYMBYTES;
-  mlk_polyvec a[MLKEM_K], e, pkpv, skpv;
+  mlk_polymat a;
+  mlk_polyvec e, pkpv, skpv;
   mlk_polyvec_mulcache skpv_cache;
 
   MLK_ALIGN uint8_t coins_with_domain_separator[MLKEM_SYMBYTES + 1];
@@ -348,40 +345,38 @@ void mlk_indcpa_keypair_derand(uint8_t pk[MLKEM_INDCPA_PUBLICKEYBYTES],
   mlk_gen_matrix(a, publicseed, 0 /* no transpose */);
 
 #if MLKEM_K == 2
-  mlk_poly_getnoise_eta1_4x(skpv.vec + 0, skpv.vec + 1, e.vec + 0, e.vec + 1,
-                            noiseseed, 0, 1, 2, 3);
+  mlk_poly_getnoise_eta1_4x(&skpv[0], &skpv[1], &e[0], &e[1], noiseseed, 0, 1,
+                            2, 3);
 #elif MLKEM_K == 3
   /*
    * Only the first three output buffers are needed.
    * The laster parameter is a dummy that's overwritten later.
    */
-  mlk_poly_getnoise_eta1_4x(skpv.vec + 0, skpv.vec + 1, skpv.vec + 2,
-                            pkpv.vec + 0 /* irrelevant */, noiseseed, 0, 1, 2,
+  mlk_poly_getnoise_eta1_4x(&skpv[0], &skpv[1], &skpv[2],
+                            &pkpv[0] /* irrelevant */, noiseseed, 0, 1, 2,
                             0xFF /* irrelevant */);
   /* Same here */
-  mlk_poly_getnoise_eta1_4x(e.vec + 0, e.vec + 1, e.vec + 2,
-                            pkpv.vec + 0 /* irrelevant */, noiseseed, 3, 4, 5,
-                            0xFF /* irrelevant */);
+  mlk_poly_getnoise_eta1_4x(&e[0], &e[1], &e[2], &pkpv[0] /* irrelevant */,
+                            noiseseed, 3, 4, 5, 0xFF /* irrelevant */);
 #elif MLKEM_K == 4
-  mlk_poly_getnoise_eta1_4x(skpv.vec + 0, skpv.vec + 1, skpv.vec + 2,
-                            skpv.vec + 3, noiseseed, 0, 1, 2, 3);
-  mlk_poly_getnoise_eta1_4x(e.vec + 0, e.vec + 1, e.vec + 2, e.vec + 3,
-                            noiseseed, 4, 5, 6, 7);
+  mlk_poly_getnoise_eta1_4x(&skpv[0], &skpv[1], &skpv[2], &skpv[3], noiseseed,
+                            0, 1, 2, 3);
+  mlk_poly_getnoise_eta1_4x(&e[0], &e[1], &e[2], &e[3], noiseseed, 4, 5, 6, 7);
 #endif
 
-  mlk_polyvec_ntt(&skpv);
-  mlk_polyvec_ntt(&e);
+  mlk_polyvec_ntt(skpv);
+  mlk_polyvec_ntt(e);
 
-  mlk_polyvec_mulcache_compute(&skpv_cache, &skpv);
-  mlk_matvec_mul(&pkpv, a, &skpv, &skpv_cache);
-  mlk_polyvec_tomont(&pkpv);
+  mlk_polyvec_mulcache_compute(skpv_cache, skpv);
+  mlk_matvec_mul(pkpv, a, skpv, skpv_cache);
+  mlk_polyvec_tomont(pkpv);
 
-  mlk_polyvec_add(&pkpv, &e);
-  mlk_polyvec_reduce(&pkpv);
-  mlk_polyvec_reduce(&skpv);
+  mlk_polyvec_add(pkpv, e);
+  mlk_polyvec_reduce(pkpv);
+  mlk_polyvec_reduce(skpv);
 
-  mlk_pack_sk(sk, &skpv);
-  mlk_pack_pk(pk, &pkpv, publicseed);
+  mlk_pack_sk(sk, skpv);
+  mlk_pack_pk(pk, pkpv, publicseed);
 
   /* Specification: Partially implements
    * [FIPS 203, Section 3.3, Destruction of intermediate values] */
@@ -408,11 +403,12 @@ void mlk_indcpa_enc(uint8_t c[MLKEM_INDCPA_BYTES],
                     const uint8_t coins[MLKEM_SYMBYTES])
 {
   MLK_ALIGN uint8_t seed[MLKEM_SYMBYTES];
-  mlk_polyvec sp, pkpv, ep, at[MLKEM_K], b;
+  mlk_polymat at;
+  mlk_polyvec sp, pkpv, ep, b;
   mlk_poly v, k, epp;
   mlk_polyvec_mulcache sp_cache;
 
-  mlk_unpack_pk(&pkpv, seed, pk);
+  mlk_unpack_pk(pkpv, seed, pk);
   mlk_poly_frommsg(&k, m);
 
   /*
@@ -426,44 +422,41 @@ void mlk_indcpa_enc(uint8_t c[MLKEM_INDCPA_BYTES],
   mlk_gen_matrix(at, seed, 1 /* transpose */);
 
 #if MLKEM_K == 2
-  mlk_poly_getnoise_eta1122_4x(sp.vec + 0, sp.vec + 1, ep.vec + 0, ep.vec + 1,
-                               coins, 0, 1, 2, 3);
+  mlk_poly_getnoise_eta1122_4x(&sp[0], &sp[1], &ep[0], &ep[1], coins, 0, 1, 2,
+                               3);
   mlk_poly_getnoise_eta2(&epp, coins, 4);
 #elif MLKEM_K == 3
   /*
    * In this call, only the first three output buffers are needed.
    * The last parameter is a dummy that's overwritten later.
    */
-  mlk_poly_getnoise_eta1_4x(sp.vec + 0, sp.vec + 1, sp.vec + 2, &b.vec[0],
-                            coins, 0, 1, 2, 0xFF);
+  mlk_poly_getnoise_eta1_4x(&sp[0], &sp[1], &sp[2], &b[0], coins, 0, 1, 2,
+                            0xFF);
   /* The fourth output buffer in this call _is_ used. */
-  mlk_poly_getnoise_eta2_4x(ep.vec + 0, ep.vec + 1, ep.vec + 2, &epp, coins, 3,
-                            4, 5, 6);
+  mlk_poly_getnoise_eta2_4x(&ep[0], &ep[1], &ep[2], &epp, coins, 3, 4, 5, 6);
 #elif MLKEM_K == 4
-  mlk_poly_getnoise_eta1_4x(sp.vec + 0, sp.vec + 1, sp.vec + 2, sp.vec + 3,
-                            coins, 0, 1, 2, 3);
-  mlk_poly_getnoise_eta2_4x(ep.vec + 0, ep.vec + 1, ep.vec + 2, ep.vec + 3,
-                            coins, 4, 5, 6, 7);
+  mlk_poly_getnoise_eta1_4x(&sp[0], &sp[1], &sp[2], &sp[3], coins, 0, 1, 2, 3);
+  mlk_poly_getnoise_eta2_4x(&ep[0], &ep[1], &ep[2], &ep[3], coins, 4, 5, 6, 7);
   mlk_poly_getnoise_eta2(&epp, coins, 8);
 #endif
 
-  mlk_polyvec_ntt(&sp);
+  mlk_polyvec_ntt(sp);
 
-  mlk_polyvec_mulcache_compute(&sp_cache, &sp);
-  mlk_matvec_mul(&b, at, &sp, &sp_cache);
-  mlk_polyvec_basemul_acc_montgomery_cached(&v, &pkpv, &sp, &sp_cache);
+  mlk_polyvec_mulcache_compute(sp_cache, sp);
+  mlk_matvec_mul(b, at, sp, sp_cache);
+  mlk_polyvec_basemul_acc_montgomery_cached(&v, pkpv, sp, sp_cache);
 
-  mlk_polyvec_invntt_tomont(&b);
+  mlk_polyvec_invntt_tomont(b);
   mlk_poly_invntt_tomont(&v);
 
-  mlk_polyvec_add(&b, &ep);
+  mlk_polyvec_add(b, ep);
   mlk_poly_add(&v, &epp);
   mlk_poly_add(&v, &k);
 
-  mlk_polyvec_reduce(&b);
+  mlk_polyvec_reduce(b);
   mlk_poly_reduce(&v);
 
-  mlk_pack_ciphertext(c, &b, &v);
+  mlk_pack_ciphertext(c, b, &v);
 
   /* Specification: Partially implements
    * [FIPS 203, Section 3.3, Destruction of intermediate values] */
@@ -490,12 +483,12 @@ void mlk_indcpa_dec(uint8_t m[MLKEM_INDCPA_MSGBYTES],
   mlk_poly v, sb;
   mlk_polyvec_mulcache b_cache;
 
-  mlk_unpack_ciphertext(&b, &v, c);
-  mlk_unpack_sk(&skpv, sk);
+  mlk_unpack_ciphertext(b, &v, c);
+  mlk_unpack_sk(skpv, sk);
 
-  mlk_polyvec_ntt(&b);
-  mlk_polyvec_mulcache_compute(&b_cache, &b);
-  mlk_polyvec_basemul_acc_montgomery_cached(&sb, &skpv, &b, &b_cache);
+  mlk_polyvec_ntt(b);
+  mlk_polyvec_mulcache_compute(b_cache, b);
+  mlk_polyvec_basemul_acc_montgomery_cached(&sb, skpv, b, b_cache);
   mlk_poly_invntt_tomont(&sb);
 
   mlk_poly_sub(&v, &sb);

mlkem/indcpa.h

@@ -19,7 +19,7 @@
  *              uniformly random. Performs rejection sampling on output of
  *              a XOF
  *
- * Arguments:   - mlk_polyvec *a: pointer to output matrix A
+ * Arguments:   - mlk_polymat a: pointer to output matrix A
  *              - const uint8_t *seed: pointer to input seed
  *              - int transposed: boolean deciding whether A or A^T is generated
  *
@@ -29,15 +29,15 @@
  *
  **************************************************/
 MLK_INTERNAL_API
-void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
+void mlk_gen_matrix(mlk_polymat a, const uint8_t seed[MLKEM_SYMBYTES],
                     int transposed)
 __contract__(
-  requires(memory_no_alias(a, sizeof(mlk_polyvec) * MLKEM_K))
+  requires(memory_no_alias(a, sizeof(mlk_polymat)))
   requires(memory_no_alias(seed, MLKEM_SYMBYTES))
   requires(transposed == 0 || transposed == 1)
   assigns(object_whole(a))
-  ensures(forall(x, 0, MLKEM_K, forall(y, 0, MLKEM_K,
-  array_bound(a[x].vec[y].coeffs, 0, MLKEM_N, 0, MLKEM_Q))));
+  ensures(forall(x, 0, MLKEM_K * MLKEM_K,
+    array_bound(a[x].coeffs, 0, MLKEM_N, 0, MLKEM_Q)))
 );
 
 #define mlk_indcpa_keypair_derand MLK_NAMESPACE_K(indcpa_keypair_derand)

mlkem/kem.c

@@ -55,9 +55,9 @@ static int mlk_check_pk(const uint8_t pk[MLKEM_INDCCA_PUBLICKEYBYTES])
   mlk_polyvec p;
   uint8_t p_reencoded[MLKEM_POLYVECBYTES];
 
-  mlk_polyvec_frombytes(&p, pk);
-  mlk_polyvec_reduce(&p);
-  mlk_polyvec_tobytes(p_reencoded, &p);
+  mlk_polyvec_frombytes(p, pk);
+  mlk_polyvec_reduce(p);
+  mlk_polyvec_tobytes(p_reencoded, p);
 
   /* We use a constant-time memcmp here to avoid having to
    * declassify the PK before the PCT has succeeded. */