mlkem_native.h: Add PCT-aware MLK_TOTAL_ALLOC constants

Split KEYPAIR allocation constants into _NO_PCT and _PCT variants to
accurately reflect memory usage with and without pairwise consistency
testing (PCT). MLK_TOTAL_ALLOC_*_KEYPAIR now automatically selects
the appropriate value based on MLK_CONFIG_KEYGEN_PCT.

For legacy configs, assume PCT is enabled (conservative).

The default alloc test config no longer enables PCT; PCT is tested
via config-variations.

 - Ported from pq-code-package/mldsa-native#869

Signed-off-by: Matthias J. Kannwischer <[email protected]>

Author: mkannwischer

mlkem/mlkem_native.S

@@ -177,14 +177,20 @@
 #undef MLK_TOTAL_ALLOC_1024_DECAPS
 #undef MLK_TOTAL_ALLOC_1024_ENCAPS
 #undef MLK_TOTAL_ALLOC_1024_KEYPAIR
+#undef MLK_TOTAL_ALLOC_1024_KEYPAIR_NO_PCT
+#undef MLK_TOTAL_ALLOC_1024_KEYPAIR_PCT
 #undef MLK_TOTAL_ALLOC_512
 #undef MLK_TOTAL_ALLOC_512_DECAPS
 #undef MLK_TOTAL_ALLOC_512_ENCAPS
 #undef MLK_TOTAL_ALLOC_512_KEYPAIR
+#undef MLK_TOTAL_ALLOC_512_KEYPAIR_NO_PCT
+#undef MLK_TOTAL_ALLOC_512_KEYPAIR_PCT
 #undef MLK_TOTAL_ALLOC_768
 #undef MLK_TOTAL_ALLOC_768_DECAPS
 #undef MLK_TOTAL_ALLOC_768_ENCAPS
 #undef MLK_TOTAL_ALLOC_768_KEYPAIR
+#undef MLK_TOTAL_ALLOC_768_KEYPAIR_NO_PCT
+#undef MLK_TOTAL_ALLOC_768_KEYPAIR_PCT
 #undef crypto_kem_check_pk
 #undef crypto_kem_check_sk
 #undef crypto_kem_dec

mlkem/mlkem_native.c

@@ -166,14 +166,20 @@
 #undef MLK_TOTAL_ALLOC_1024_DECAPS
 #undef MLK_TOTAL_ALLOC_1024_ENCAPS
 #undef MLK_TOTAL_ALLOC_1024_KEYPAIR
+#undef MLK_TOTAL_ALLOC_1024_KEYPAIR_NO_PCT
+#undef MLK_TOTAL_ALLOC_1024_KEYPAIR_PCT
 #undef MLK_TOTAL_ALLOC_512
 #undef MLK_TOTAL_ALLOC_512_DECAPS
 #undef MLK_TOTAL_ALLOC_512_ENCAPS
 #undef MLK_TOTAL_ALLOC_512_KEYPAIR
+#undef MLK_TOTAL_ALLOC_512_KEYPAIR_NO_PCT
+#undef MLK_TOTAL_ALLOC_512_KEYPAIR_PCT
 #undef MLK_TOTAL_ALLOC_768
 #undef MLK_TOTAL_ALLOC_768_DECAPS
 #undef MLK_TOTAL_ALLOC_768_ENCAPS
 #undef MLK_TOTAL_ALLOC_768_KEYPAIR
+#undef MLK_TOTAL_ALLOC_768_KEYPAIR_NO_PCT
+#undef MLK_TOTAL_ALLOC_768_KEYPAIR_PCT
 #undef crypto_kem_check_pk
 #undef crypto_kem_check_sk
 #undef crypto_kem_dec

mlkem/mlkem_native.h

@@ -447,22 +447,38 @@ int MLK_API_NAMESPACE(check_sk)(
  * fixed-sized buffer and a simple allocator (e.g., bump allocator).
  */
 /* check-magic: off */
-#define MLK_TOTAL_ALLOC_512_KEYPAIR 10048
+#define MLK_TOTAL_ALLOC_512_KEYPAIR_NO_PCT 5824
+#define MLK_TOTAL_ALLOC_512_KEYPAIR_PCT 10048
 #define MLK_TOTAL_ALLOC_512_ENCAPS 8384
 #define MLK_TOTAL_ALLOC_512_DECAPS 9152
-#define MLK_TOTAL_ALLOC_768_KEYPAIR 15552
+#define MLK_TOTAL_ALLOC_768_KEYPAIR_NO_PCT 10176
+#define MLK_TOTAL_ALLOC_768_KEYPAIR_PCT 15552
 #define MLK_TOTAL_ALLOC_768_ENCAPS 13248
 #define MLK_TOTAL_ALLOC_768_DECAPS 14336
-#define MLK_TOTAL_ALLOC_1024_KEYPAIR 22400
+#define MLK_TOTAL_ALLOC_1024_KEYPAIR_NO_PCT 15552
+#define MLK_TOTAL_ALLOC_1024_KEYPAIR_PCT 22400
 #define MLK_TOTAL_ALLOC_1024_ENCAPS 19136
 #define MLK_TOTAL_ALLOC_1024_DECAPS 20704
 /* check-magic: on */
 
+/*
+ * MLK_TOTAL_ALLOC_*_KEYPAIR adapts based on MLK_CONFIG_KEYGEN_PCT.
+ * For legacy config, we don't know which options are used, so assume
+ * the worst case (PCT enabled).
+ */
+#if defined(MLK_API_LEGACY_CONFIG) || defined(MLK_CONFIG_KEYGEN_PCT)
+#define MLK_TOTAL_ALLOC_512_KEYPAIR MLK_TOTAL_ALLOC_512_KEYPAIR_PCT
+#define MLK_TOTAL_ALLOC_768_KEYPAIR MLK_TOTAL_ALLOC_768_KEYPAIR_PCT
+#define MLK_TOTAL_ALLOC_1024_KEYPAIR MLK_TOTAL_ALLOC_1024_KEYPAIR_PCT
+#else
+#define MLK_TOTAL_ALLOC_512_KEYPAIR MLK_TOTAL_ALLOC_512_KEYPAIR_NO_PCT
+#define MLK_TOTAL_ALLOC_768_KEYPAIR MLK_TOTAL_ALLOC_768_KEYPAIR_NO_PCT
+#define MLK_TOTAL_ALLOC_1024_KEYPAIR MLK_TOTAL_ALLOC_1024_KEYPAIR_NO_PCT
+#endif
+
 /*
  * `MLK_MAX_TOTAL_ALLOC_{KEYPAIR,ENCAPS,DECAPS}` is the maximum across all
  * parameter sets for each operation.
- * `MLK_MAX_TOTAL_ALLOC` is the maximum across all parameter sets and
- * operations.
  */
 #define MLK_MAX_TOTAL_ALLOC_KEYPAIR MLK_TOTAL_ALLOC_1024_KEYPAIR
 #define MLK_MAX_TOTAL_ALLOC_ENCAPS MLK_TOTAL_ALLOC_1024_ENCAPS

test/configs/configs.yml

@@ -434,11 +434,6 @@ configs:
     description: "Using custom allocation that can be made fail at specific invocation"
     defines:
       MLK_CONFIG_NAMESPACE_PREFIX: mlk
-      MLK_CONFIG_KEYGEN_PCT:
-        content: |
-          #if !defined(MLK_CONFIG_KEYGEN_PCT)
-          #define MLK_CONFIG_KEYGEN_PCT
-          #endif
       MLK_CONFIG_CUSTOM_ALLOC_FREE:
         content: |
           #define MLK_CONFIG_CUSTOM_ALLOC_FREE

test/configs/test_alloc_config.h

@@ -31,7 +31,6 @@
  * This configuration differs from the default mlkem/mlkem_native_config.h in
  * the following places:
  *   - MLK_CONFIG_NAMESPACE_PREFIX
- *   - MLK_CONFIG_KEYGEN_PCT
  *   - MLK_CONFIG_CUSTOM_ALLOC_FREE
  */
 
@@ -632,10 +631,7 @@ void custom_free(void *p, size_t sz, const char *file, int line,
  *              key generation.
  *
  *****************************************************************************/
-#if !defined(MLK_CONFIG_KEYGEN_PCT)
-#define MLK_CONFIG_KEYGEN_PCT
-#endif
-
+/* #define MLK_CONFIG_KEYGEN_PCT */
 
 /******************************************************************************
  * Name:        MLK_CONFIG_KEYGEN_PCT_BREAKAGE_TEST

test/src/test_alloc.c

@@ -16,6 +16,9 @@
  * Level-dependent allocation limit macros.
  * These expand to the right MLK_TOTAL_ALLOC_{512,768,1024}_* constant
  * based on MLK_CONFIG_API_PARAMETER_SET.
+ *
+ * Note: MLK_TOTAL_ALLOC_*_KEYPAIR in the header automatically adapts
+ * based on MLK_CONFIG_KEYGEN_PCT.
  */
 #define MLK_TOTAL_ALLOC_KEYPAIR__(LVL) MLK_TOTAL_ALLOC_##LVL##_KEYPAIR
 #define MLK_TOTAL_ALLOC_KEYPAIR_(LVL) MLK_TOTAL_ALLOC_KEYPAIR__(LVL)
@@ -470,6 +473,7 @@ int main(void)
   /*
    * For parameter set 1024, also check that the high watermarks match
    * the MLK_MAX_TOTAL_ALLOC_* constants (which are defined as the 1024 values).
+   * MLK_MAX_TOTAL_ALLOC_KEYPAIR adapts based on MLK_CONFIG_KEYGEN_PCT.
    */
 #if MLK_CONFIG_API_PARAMETER_SET == 1024
   CHECK_ALLOC_MATCH(global_bump_high_mark_keypair, MLK_MAX_TOTAL_ALLOC_KEYPAIR);