Undo splitting of 2D arrays into separate objects

Previously, we had to split various 2D arrays into separate objects
to accommodate the strict semantics of is_fresh/memory_no_alias in
function contract preconditions.

With the new relaxed semantics of is_fresh/memory_no_alias, this is
no longer necessary.

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

Author: hanno-becker

examples/monolithic_build/mlkem_native_monobuild.c

@@ -198,6 +198,7 @@
 #undef mlk_polyvec_tomont
 /* mlkem/sys.h */
 #undef MLK_ALIGN
+#undef MLK_ALIGN_UP
 #undef MLK_ALWAYS_INLINE
 #undef MLK_CET_ENDBR
 #undef MLK_CT_TESTING_DECLASSIFY

mlkem/indcpa.c

@@ -197,19 +197,11 @@ void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
    * of the same parent object.
    */
 
-  MLK_ALIGN uint8_t seed0[MLKEM_SYMBYTES + 2];
-  MLK_ALIGN uint8_t seed1[MLKEM_SYMBYTES + 2];
-  MLK_ALIGN uint8_t seed2[MLKEM_SYMBYTES + 2];
-  MLK_ALIGN uint8_t seed3[MLKEM_SYMBYTES + 2];
-  uint8_t *seedxy[4];
-  seedxy[0] = seed0;
-  seedxy[1] = seed1;
-  seedxy[2] = seed2;
-  seedxy[3] = seed3;
+  MLK_ALIGN uint8_t seed_ext[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)];
 
   for (j = 0; j < 4; j++)
   {
-    memcpy(seedxy[j], seed, MLKEM_SYMBYTES);
+    memcpy(seed_ext[j], seed, MLKEM_SYMBYTES);
   }
 
   /* Sample 4 matrix entries a time. */
@@ -223,21 +215,21 @@ void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
       y = (i + j) % MLKEM_K;
       if (transposed)
       {
-        seedxy[j][MLKEM_SYMBYTES + 0] = x;
-        seedxy[j][MLKEM_SYMBYTES + 1] = y;
+        seed_ext[j][MLKEM_SYMBYTES + 0] = x;
+        seed_ext[j][MLKEM_SYMBYTES + 1] = y;
       }
       else
       {
-        seedxy[j][MLKEM_SYMBYTES + 0] = y;
-        seedxy[j][MLKEM_SYMBYTES + 1] = x;
+        seed_ext[j][MLKEM_SYMBYTES + 0] = y;
+        seed_ext[j][MLKEM_SYMBYTES + 1] = x;
       }
     }
 
     /*
      * 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, seedxy);
+    mlk_poly_rej_uniform_x4(&a[0].vec[0] + i, seed_ext);
   }
 
   /* For MLKEM_K == 3, sample the last entry individually. */
@@ -249,16 +241,16 @@ void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
 
     if (transposed)
     {
-      seed0[MLKEM_SYMBYTES + 0] = x;
-      seed0[MLKEM_SYMBYTES + 1] = y;
+      seed_ext[0][MLKEM_SYMBYTES + 0] = x;
+      seed_ext[0][MLKEM_SYMBYTES + 1] = y;
     }
     else
     {
-      seed0[MLKEM_SYMBYTES + 0] = y;
-      seed0[MLKEM_SYMBYTES + 1] = x;
+      seed_ext[0][MLKEM_SYMBYTES + 0] = y;
+      seed_ext[0][MLKEM_SYMBYTES + 1] = x;
     }
 
-    mlk_poly_rej_uniform(&a[0].vec[0] + i, seed0);
+    mlk_poly_rej_uniform(&a[0].vec[0] + i, seed_ext[0]);
     i++;
   }
 
@@ -278,10 +270,7 @@ void mlk_gen_matrix(mlk_polyvec *a, const uint8_t seed[MLKEM_SYMBYTES],
 
   /* Specification: Partially implements
    * [FIPS 203, Section 3.3, Destruction of intermediate values] */
-  mlk_zeroize(seed0, sizeof(seed0));
-  mlk_zeroize(seed1, sizeof(seed1));
-  mlk_zeroize(seed2, sizeof(seed2));
-  mlk_zeroize(seed3, sizeof(seed3));
+  mlk_zeroize(seed_ext, sizeof(seed_ext));
 }
 
 /*************************************************

mlkem/poly_k.c

@@ -286,27 +286,21 @@ void mlk_poly_getnoise_eta1_4x(mlk_poly *r0, mlk_poly *r1, mlk_poly *r2,
                                uint8_t nonce0, uint8_t nonce1, uint8_t nonce2,
                                uint8_t nonce3)
 {
-  MLK_ALIGN uint8_t buf0[MLKEM_ETA1 * MLKEM_N / 4];
-  MLK_ALIGN uint8_t buf1[MLKEM_ETA1 * MLKEM_N / 4];
-  MLK_ALIGN uint8_t buf2[MLKEM_ETA1 * MLKEM_N / 4];
-  MLK_ALIGN uint8_t buf3[MLKEM_ETA1 * MLKEM_N / 4];
-  MLK_ALIGN uint8_t extkey0[MLKEM_SYMBYTES + 1];
-  MLK_ALIGN uint8_t extkey1[MLKEM_SYMBYTES + 1];
-  MLK_ALIGN uint8_t extkey2[MLKEM_SYMBYTES + 1];
-  MLK_ALIGN uint8_t extkey3[MLKEM_SYMBYTES + 1];
-  memcpy(extkey0, seed, MLKEM_SYMBYTES);
-  memcpy(extkey1, seed, MLKEM_SYMBYTES);
-  memcpy(extkey2, seed, MLKEM_SYMBYTES);
-  memcpy(extkey3, seed, MLKEM_SYMBYTES);
-  extkey0[MLKEM_SYMBYTES] = nonce0;
-  extkey1[MLKEM_SYMBYTES] = nonce1;
-  extkey2[MLKEM_SYMBYTES] = nonce2;
-  extkey3[MLKEM_SYMBYTES] = nonce3;
-  mlk_prf_eta1_x4(buf0, buf1, buf2, buf3, extkey0, extkey1, extkey2, extkey3);
-  mlk_poly_cbd_eta1(r0, buf0);
-  mlk_poly_cbd_eta1(r1, buf1);
-  mlk_poly_cbd_eta1(r2, buf2);
-  mlk_poly_cbd_eta1(r3, buf3);
+  MLK_ALIGN uint8_t buf[4][MLK_ALIGN_UP(MLKEM_ETA1 * MLKEM_N / 4)];
+  MLK_ALIGN uint8_t extkey[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 1)];
+  memcpy(extkey[0], seed, MLKEM_SYMBYTES);
+  memcpy(extkey[1], seed, MLKEM_SYMBYTES);
+  memcpy(extkey[2], seed, MLKEM_SYMBYTES);
+  memcpy(extkey[3], seed, MLKEM_SYMBYTES);
+  extkey[0][MLKEM_SYMBYTES] = nonce0;
+  extkey[1][MLKEM_SYMBYTES] = nonce1;
+  extkey[2][MLKEM_SYMBYTES] = nonce2;
+  extkey[3][MLKEM_SYMBYTES] = nonce3;
+  mlk_prf_eta1_x4(buf, extkey);
+  mlk_poly_cbd_eta1(r0, buf[0]);
+  mlk_poly_cbd_eta1(r1, buf[1]);
+  mlk_poly_cbd_eta1(r2, buf[2]);
+  mlk_poly_cbd_eta1(r3, buf[3]);
 
   mlk_assert_abs_bound(r0, MLKEM_N, MLKEM_ETA1 + 1);
   mlk_assert_abs_bound(r1, MLKEM_N, MLKEM_ETA1 + 1);
@@ -315,14 +309,8 @@ void mlk_poly_getnoise_eta1_4x(mlk_poly *r0, mlk_poly *r1, mlk_poly *r2,
 
   /* Specification: Partially implements
    * [FIPS 203, Section 3.3, Destruction of intermediate values] */
-  mlk_zeroize(buf0, sizeof(buf0));
-  mlk_zeroize(buf1, sizeof(buf1));
-  mlk_zeroize(buf2, sizeof(buf2));
-  mlk_zeroize(buf3, sizeof(buf3));
-  mlk_zeroize(extkey0, sizeof(extkey0));
-  mlk_zeroize(extkey1, sizeof(extkey1));
-  mlk_zeroize(extkey2, sizeof(extkey2));
-  mlk_zeroize(extkey3, sizeof(extkey3));
+  mlk_zeroize(buf, sizeof(buf));
+  mlk_zeroize(extkey, sizeof(extkey));
 }
 
 #if MLKEM_K == 2 || MLKEM_K == 4
@@ -400,42 +388,34 @@ void mlk_poly_getnoise_eta1122_4x(mlk_poly *r0, mlk_poly *r1, mlk_poly *r2,
 #if MLKEM_ETA2 >= MLKEM_ETA1
 #error mlk_poly_getnoise_eta1122_4x assumes MLKEM_ETA1 > MLKEM_ETA2
 #endif
-  MLK_ALIGN uint8_t buf0[MLKEM_ETA1 * MLKEM_N / 4];
-  MLK_ALIGN uint8_t buf1[MLKEM_ETA1 * MLKEM_N / 4];
-  /* Pad to larger buffer */
-  MLK_ALIGN uint8_t buf2[MLKEM_ETA1 * MLKEM_N / 4];
-  MLK_ALIGN uint8_t buf3[MLKEM_ETA1 * MLKEM_N / 4];
-
-  MLK_ALIGN uint8_t extkey0[MLKEM_SYMBYTES + 1];
-  MLK_ALIGN uint8_t extkey1[MLKEM_SYMBYTES + 1];
-  MLK_ALIGN uint8_t extkey2[MLKEM_SYMBYTES + 1];
-  MLK_ALIGN uint8_t extkey3[MLKEM_SYMBYTES + 1];
-
-  memcpy(extkey0, seed, MLKEM_SYMBYTES);
-  memcpy(extkey1, seed, MLKEM_SYMBYTES);
-  memcpy(extkey2, seed, MLKEM_SYMBYTES);
-  memcpy(extkey3, seed, MLKEM_SYMBYTES);
-  extkey0[MLKEM_SYMBYTES] = nonce0;
-  extkey1[MLKEM_SYMBYTES] = nonce1;
-  extkey2[MLKEM_SYMBYTES] = nonce2;
-  extkey3[MLKEM_SYMBYTES] = nonce3;
+  MLK_ALIGN uint8_t buf[4][MLK_ALIGN_UP(MLKEM_ETA1 * MLKEM_N / 4)];
+  MLK_ALIGN uint8_t extkey[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 1)];
+
+  memcpy(extkey[0], seed, MLKEM_SYMBYTES);
+  memcpy(extkey[1], seed, MLKEM_SYMBYTES);
+  memcpy(extkey[2], seed, MLKEM_SYMBYTES);
+  memcpy(extkey[3], seed, MLKEM_SYMBYTES);
+  extkey[0][MLKEM_SYMBYTES] = nonce0;
+  extkey[1][MLKEM_SYMBYTES] = nonce1;
+  extkey[2][MLKEM_SYMBYTES] = nonce2;
+  extkey[3][MLKEM_SYMBYTES] = nonce3;
 
   /* On systems with fast batched Keccak, we use 4-fold batched PRF,
-   * even though that means generating more random data in buf2 and buf3
+   * even though that means generating more random data in buf[2] and buf[3]
    * than necessary. */
 #if !defined(FIPS202_X4_DEFAULT_IMPLEMENTATION)
-  mlk_prf_eta1_x4(buf0, buf1, buf2, buf3, extkey0, extkey1, extkey2, extkey3);
+  mlk_prf_eta1_x4(buf, extkey);
 #else  /* FIPS202_X4_DEFAULT_IMPLEMENTATION */
-  mlk_prf_eta1(buf0, extkey0);
-  mlk_prf_eta1(buf1, extkey1);
-  mlk_prf_eta2(buf2, extkey2);
-  mlk_prf_eta2(buf3, extkey3);
+  mlk_prf_eta1(buf[0], extkey[0]);
+  mlk_prf_eta1(buf[1], extkey[1]);
+  mlk_prf_eta2(buf[2], extkey[2]);
+  mlk_prf_eta2(buf[3], extkey[3]);
 #endif /* FIPS202_X4_DEFAULT_IMPLEMENTATION */
 
-  mlk_poly_cbd_eta1(r0, buf0);
-  mlk_poly_cbd_eta1(r1, buf1);
-  mlk_poly_cbd_eta2(r2, buf2);
-  mlk_poly_cbd_eta2(r3, buf3);
+  mlk_poly_cbd_eta1(r0, buf[0]);
+  mlk_poly_cbd_eta1(r1, buf[1]);
+  mlk_poly_cbd_eta2(r2, buf[2]);
+  mlk_poly_cbd_eta2(r3, buf[3]);
 
   mlk_assert_abs_bound(r0, MLKEM_N, MLKEM_ETA1 + 1);
   mlk_assert_abs_bound(r1, MLKEM_N, MLKEM_ETA1 + 1);
@@ -444,14 +424,8 @@ void mlk_poly_getnoise_eta1122_4x(mlk_poly *r0, mlk_poly *r1, mlk_poly *r2,
 
   /* Specification: Partially implements
    * [FIPS 203, Section 3.3, Destruction of intermediate values] */
-  mlk_zeroize(buf0, sizeof(buf0));
-  mlk_zeroize(buf1, sizeof(buf1));
-  mlk_zeroize(buf2, sizeof(buf2));
-  mlk_zeroize(buf3, sizeof(buf3));
-  mlk_zeroize(extkey0, sizeof(extkey0));
-  mlk_zeroize(extkey1, sizeof(extkey1));
-  mlk_zeroize(extkey2, sizeof(extkey2));
-  mlk_zeroize(extkey3, sizeof(extkey3));
+  mlk_zeroize(buf, sizeof(buf));
+  mlk_zeroize(extkey, sizeof(extkey));
 }
 #endif /* MLKEM_K == 2 */
 

mlkem/sampling.c

@@ -132,35 +132,31 @@ __contract__(
  *            - x4-batched version of `rej_uniform()` from the
  *              reference implementation, leveraging x4-batched Keccak-f1600. */
 MLK_INTERNAL_API
-void mlk_poly_rej_uniform_x4(mlk_poly *vec, uint8_t *seed[4])
+void mlk_poly_rej_uniform_x4(mlk_poly *vec,
+                             uint8_t seed[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)])
 {
   /* Temporary buffers for XOF output before rejection sampling */
-  MLK_ALIGN uint8_t buf0[MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE];
-  MLK_ALIGN uint8_t buf1[MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE];
-  MLK_ALIGN uint8_t buf2[MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE];
-  MLK_ALIGN uint8_t buf3[MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE];
+  MLK_ALIGN uint8_t
+      buf[4][MLK_ALIGN_UP(MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE)];
 
   /* Tracks the number of coefficients we have already sampled */
   unsigned ctr[4];
   mlk_xof_x4_ctx statex;
   unsigned buflen;
 
-  /* seed is MLKEM_SYMBYTES + 2 bytes long, but padded to MLKEM_SYMBYTES + 16 */
   mlk_xof_x4_init(&statex);
-  mlk_xof_x4_absorb(&statex, seed[0], seed[1], seed[2], seed[3],
-                    MLKEM_SYMBYTES + 2);
+  mlk_xof_x4_absorb(&statex, seed, MLKEM_SYMBYTES + 2);
 
   /*
    * Initially, squeeze heuristic number of MLKEM_GEN_MATRIX_NBLOCKS.
    * This should generate the matrix entries with high probability.
    */
-  mlk_xof_x4_squeezeblocks(buf0, buf1, buf2, buf3, MLKEM_GEN_MATRIX_NBLOCKS,
-                           &statex);
+  mlk_xof_x4_squeezeblocks(buf, MLKEM_GEN_MATRIX_NBLOCKS, &statex);
   buflen = MLKEM_GEN_MATRIX_NBLOCKS * MLK_XOF_RATE;
-  ctr[0] = mlk_rej_uniform(vec[0].coeffs, MLKEM_N, 0, buf0, buflen);
-  ctr[1] = mlk_rej_uniform(vec[1].coeffs, MLKEM_N, 0, buf1, buflen);
-  ctr[2] = mlk_rej_uniform(vec[2].coeffs, MLKEM_N, 0, buf2, buflen);
-  ctr[3] = mlk_rej_uniform(vec[3].coeffs, MLKEM_N, 0, buf3, buflen);
+  ctr[0] = mlk_rej_uniform(vec[0].coeffs, MLKEM_N, 0, buf[0], buflen);
+  ctr[1] = mlk_rej_uniform(vec[1].coeffs, MLKEM_N, 0, buf[1], buflen);
+  ctr[2] = mlk_rej_uniform(vec[2].coeffs, MLKEM_N, 0, buf[2], buflen);
+  ctr[3] = mlk_rej_uniform(vec[3].coeffs, MLKEM_N, 0, buf[3], buflen);
 
   /*
    * So long as not all matrix entries have been generated, squeeze
@@ -170,30 +166,27 @@ void mlk_poly_rej_uniform_x4(mlk_poly *vec, uint8_t *seed[4])
   while (ctr[0] < MLKEM_N || ctr[1] < MLKEM_N || ctr[2] < MLKEM_N ||
          ctr[3] < MLKEM_N)
   __loop__(
-    assigns(ctr, statex, memory_slice(vec, sizeof(mlk_poly) * 4), object_whole(buf0),
-       object_whole(buf1), object_whole(buf2), object_whole(buf3))
+    assigns(ctr, statex, memory_slice(vec, sizeof(mlk_poly) * 4), object_whole(buf[0]),
+       object_whole(buf[1]), object_whole(buf[2]), object_whole(buf[3]))
     invariant(ctr[0] <= MLKEM_N && ctr[1] <= MLKEM_N)
     invariant(ctr[2] <= MLKEM_N && ctr[3] <= MLKEM_N)
     invariant(array_bound(vec[0].coeffs, 0, ctr[0], 0, MLKEM_Q))
     invariant(array_bound(vec[1].coeffs, 0, ctr[1], 0, MLKEM_Q))
     invariant(array_bound(vec[2].coeffs, 0, ctr[2], 0, MLKEM_Q))
     invariant(array_bound(vec[3].coeffs, 0, ctr[3], 0, MLKEM_Q)))
   {
-    mlk_xof_x4_squeezeblocks(buf0, buf1, buf2, buf3, 1, &statex);
-    ctr[0] = mlk_rej_uniform(vec[0].coeffs, MLKEM_N, ctr[0], buf0, buflen);
-    ctr[1] = mlk_rej_uniform(vec[1].coeffs, MLKEM_N, ctr[1], buf1, buflen);
-    ctr[2] = mlk_rej_uniform(vec[2].coeffs, MLKEM_N, ctr[2], buf2, buflen);
-    ctr[3] = mlk_rej_uniform(vec[3].coeffs, MLKEM_N, ctr[3], buf3, buflen);
+    mlk_xof_x4_squeezeblocks(buf, 1, &statex);
+    ctr[0] = mlk_rej_uniform(vec[0].coeffs, MLKEM_N, ctr[0], buf[0], buflen);
+    ctr[1] = mlk_rej_uniform(vec[1].coeffs, MLKEM_N, ctr[1], buf[1], buflen);
+    ctr[2] = mlk_rej_uniform(vec[2].coeffs, MLKEM_N, ctr[2], buf[2], buflen);
+    ctr[3] = mlk_rej_uniform(vec[3].coeffs, MLKEM_N, ctr[3], buf[3], buflen);
   }
 
   mlk_xof_x4_release(&statex);
 
   /* Specification: Partially implements
    * [FIPS 203, Section 3.3, Destruction of intermediate values] */
-  mlk_zeroize(buf0, sizeof(buf0));
-  mlk_zeroize(buf1, sizeof(buf1));
-  mlk_zeroize(buf2, sizeof(buf2));
-  mlk_zeroize(buf3, sizeof(buf3));
+  mlk_zeroize(buf, sizeof(buf));
 }
 
 MLK_INTERNAL_API

mlkem/sampling.h

@@ -55,24 +55,21 @@ void mlk_poly_cbd3(mlk_poly *r, const uint8_t buf[3 * MLKEM_N / 4]);
  * Description: Generate four polynomials using rejection sampling
  *              on (pseudo-)uniformly random bytes sampled from a seed.
  *
- * Arguments:   - mlk_poly *vec:           Pointer to an array of 4 polynomials
- *                                     to be sampled.
- *              - uint8_t *seed[4]:    Pointer to array of four pointers
- *                                     pointing to the seed buffers of size
- *                                     MLKEM_SYMBYTES + 2 each.
+ * Arguments:   - mlk_poly *vec:
+ *                Pointer to an array of 4 polynomials to be sampled.
+ *              - uint8_t seed[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)]:
+ *                Pointer consecutive array of seed buffers of size
+ *                MLKEM_SYMBYTES + 2 each, plus padding for alignment.
  *
  * Specification: Implements [FIPS 203, Algorithm 7, SampleNTT]
  *
  **************************************************/
 MLK_INTERNAL_API
-void mlk_poly_rej_uniform_x4(mlk_poly *vec, uint8_t *seed[4])
+void mlk_poly_rej_uniform_x4(mlk_poly *vec,
+                             uint8_t seed[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)])
 __contract__(
   requires(memory_no_alias(vec, sizeof(mlk_poly) * 4))
-  requires(memory_no_alias(seed, sizeof(uint8_t*) * 4))
-  requires(memory_no_alias(seed[0], MLKEM_SYMBYTES + 2))
-  requires(memory_no_alias(seed[1], MLKEM_SYMBYTES + 2))
-  requires(memory_no_alias(seed[2], MLKEM_SYMBYTES + 2))
-  requires(memory_no_alias(seed[3], MLKEM_SYMBYTES + 2))
+  requires(memory_no_alias(seed, 4 * MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)))
   assigns(memory_slice(vec, sizeof(mlk_poly) * 4))
   ensures(array_bound(vec[0].coeffs, 0, MLKEM_N, 0, MLKEM_Q))
   ensures(array_bound(vec[1].coeffs, 0, MLKEM_N, 0, MLKEM_Q))

mlkem/symmetric.h

@@ -30,9 +30,10 @@
   mlk_shake256(OUT, (ETA) * MLKEM_N / 4, IN, MLKEM_SYMBYTES + 1)
 #define mlk_prf_eta1(OUT, IN) mlk_prf_eta(MLKEM_ETA1, OUT, IN)
 #define mlk_prf_eta2(OUT, IN) mlk_prf_eta(MLKEM_ETA2, OUT, IN)
-#define mlk_prf_eta1_x4(OUT0, OUT1, OUT2, OUT3, IN0, IN1, IN2, IN3)            \
-  mlk_shake256x4(OUT0, OUT1, OUT2, OUT3, (MLKEM_ETA1 * MLKEM_N / 4), IN0, IN1, \
-                 IN2, IN3, MLKEM_SYMBYTES + 1)
+#define mlk_prf_eta1_x4(OUT, IN)                                        \
+  mlk_shake256x4((OUT)[0], (OUT)[1], (OUT)[2], (OUT)[3],                \
+                 (MLKEM_ETA1 * MLKEM_N / 4), (IN)[0], (IN)[1], (IN)[2], \
+                 (IN)[3], MLKEM_SYMBYTES + 1)
 
 /* XOF function, FIPS 203 4.1 */
 #define mlk_xof_ctx mlk_shake128ctx
@@ -45,10 +46,12 @@
 #define mlk_xof_release(CTX) mlk_shake128_release((CTX))
 
 #define mlk_xof_x4_init(CTX) mlk_shake128x4_init((CTX))
-#define mlk_xof_x4_absorb(CTX, IN0, IN1, IN2, IN3, INBYTES) \
-  mlk_shake128x4_absorb_once((CTX), (IN0), (IN1), (IN2), (IN3), (INBYTES))
-#define mlk_xof_x4_squeezeblocks(BUF0, BUF1, BUF2, BUF3, NBLOCKS, CTX) \
-  mlk_shake128x4_squeezeblocks((BUF0), (BUF1), (BUF2), (BUF3), (NBLOCKS), (CTX))
+#define mlk_xof_x4_absorb(CTX, IN, INBYTES)                             \
+  mlk_shake128x4_absorb_once((CTX), (IN)[0], (IN)[1], (IN)[2], (IN)[3], \
+                             (INBYTES))
+#define mlk_xof_x4_squeezeblocks(BUF, NBLOCKS, CTX)                    \
+  mlk_shake128x4_squeezeblocks((BUF)[0], (BUF)[1], (BUF)[2], (BUF)[3], \
+                               (NBLOCKS), (CTX))
 #define mlk_xof_x4_release(CTX) mlk_shake128x4_release((CTX))
 
 #define MLK_XOF_RATE SHAKE128_RATE

mlkem/sys.h

@@ -109,6 +109,8 @@
 #endif
 
 #define MLK_DEFAULT_ALIGN 32
+#define MLK_ALIGN_UP(N) \
+  ((((N) + (MLK_DEFAULT_ALIGN - 1)) / MLK_DEFAULT_ALIGN) * MLK_DEFAULT_ALIGN)
 #if defined(__GNUC__)
 #define MLK_ALIGN __attribute__((aligned(MLK_DEFAULT_ALIGN)))
 #elif defined(_MSC_VER)

proofs/cbmc/poly_rej_uniform_x4/poly_rej_uniform_x4_harness.c

@@ -8,6 +8,6 @@
 void harness(void)
 {
   mlk_poly out[4];
-  uint8_t *seed[4];
+  uint8_t seed[4][MLK_ALIGN_UP(MLKEM_SYMBYTES + 2)];
   mlk_poly_rej_uniform_x4(out, seed);
 }