From b1d30129f8227cdd138903244689d370e1a93d06 Mon Sep 17 00:00:00 2001
From: Tiago Oliveira <tiago.oliveira@sandboxquantum.com>
Date: Thu, 18 Apr 2024 12:39:41 +0100
Subject: [PATCH] fix sct for mlkem/xwing
 (https://github.com/formosa-crypto/formosa-mlkem/commit/b8078109202057b58f1b352c3316bbf97dcc3327)

---
 src/crypto_kem/mlkem/mlkem768/amd64/avx2/indcpa.jinc |  2 +-
 src/crypto_kem/mlkem/mlkem768/amd64/avx2/kem.jazz    | 11 +++++++++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/src/crypto_kem/mlkem/mlkem768/amd64/avx2/indcpa.jinc b/src/crypto_kem/mlkem/mlkem768/amd64/avx2/indcpa.jinc
index 23c6b279..f2323f1b 100644
--- a/src/crypto_kem/mlkem/mlkem768/amd64/avx2/indcpa.jinc
+++ b/src/crypto_kem/mlkem/mlkem768/amd64/avx2/indcpa.jinc
@@ -4,7 +4,7 @@ require "polyvec.jinc"
 require "gen_matrix.jinc"
 
 inline
-fn __indcpa_keypair(reg u64 pkp, reg u64 skp, reg ptr u8[MLKEM_SYMBYTES] randomnessp)
+fn __indcpa_keypair(#spill_to_mmx reg u64 pkp skp, reg ptr u8[MLKEM_SYMBYTES] randomnessp)
 {
   stack u16[MLKEM_K*MLKEM_VECN] aa;
   stack u16[MLKEM_VECN] e pkpv skpv;
diff --git a/src/crypto_kem/mlkem/mlkem768/amd64/avx2/kem.jazz b/src/crypto_kem/mlkem/mlkem768/amd64/avx2/kem.jazz
index 13ba3cc9..14e4b580 100644
--- a/src/crypto_kem/mlkem/mlkem768/amd64/avx2/kem.jazz
+++ b/src/crypto_kem/mlkem/mlkem768/amd64/avx2/kem.jazz
@@ -7,6 +7,8 @@ export fn jade_kem_mlkem_mlkem768_amd64_avx2_keypair_derand(reg u64 public_key s
   reg ptr u8[MLKEM_SYMBYTES*2] randomnessp;
   inline int i;
 
+  _ = #init_msf();
+
   public_key = public_key;
   secret_key = secret_key;
 
@@ -29,6 +31,8 @@ export fn jade_kem_mlkem_mlkem768_amd64_avx2_enc_derand(reg u64 ciphertext share
   reg ptr u8[MLKEM_SYMBYTES] randomnessp;
   inline int i;
 
+  _ = #init_msf();
+
   ciphertext = ciphertext;
   shared_secret = shared_secret;
   public_key = public_key;
@@ -56,6 +60,8 @@ export fn jade_kem_mlkem_mlkem768_amd64_avx2_keypair(reg u64 public_key secret_k
 
   randomnessp = randomness;
   randomnessp = #randombytes(randomnessp);
+  _ = #init_msf();
+
   __crypto_kem_keypair_jazz(public_key, secret_key, randomnessp);
   ?{}, r = #set0();
   return r;
@@ -73,6 +79,8 @@ export fn jade_kem_mlkem_mlkem768_amd64_avx2_enc(reg u64 ciphertext shared_secre
 
   randomnessp = randomness;
   randomnessp = #randombytes(randomnessp);
+  _ = #init_msf();
+
   __crypto_kem_enc_jazz(ciphertext, shared_secret, public_key, randomnessp);
   ?{}, r = #set0();
   return r;
@@ -81,6 +89,9 @@ export fn jade_kem_mlkem_mlkem768_amd64_avx2_enc(reg u64 ciphertext shared_secre
 export fn jade_kem_mlkem_mlkem768_amd64_avx2_dec(reg u64 shared_secret ciphertext secret_key) -> reg u64
 {
   reg u64 r;
+
+  _ = #init_msf();
+
   __crypto_kem_dec_jazz(shared_secret, ciphertext, secret_key);
   ?{}, r = #set0();
   return r;