From 0c1c0372c7cc899d07347b1cb7b6b4378e14188f Mon Sep 17 00:00:00 2001
From: Yawning Angel <yawning@schwanenlied.me>
Date: Sat, 16 May 2026 06:49:21 +0900
Subject: [PATCH] core/crypto/mlkem: Minor cleanups

---
 core/crypto/_mlkem/poly.odin    |  6 +++---
 core/crypto/_mlkem/polyvec.odin |  4 ++--
 core/crypto/mlkem/api.odin      | 36 +++++++++++++++++++++++++++++++++
 3 files changed, 41 insertions(+), 5 deletions(-)

diff --git a/core/crypto/_mlkem/poly.odin b/core/crypto/_mlkem/poly.odin
index 4a79e9770..1982f9102 100644
--- a/core/crypto/_mlkem/poly.odin
+++ b/core/crypto/_mlkem/poly.odin
@@ -58,7 +58,7 @@ poly_compress :: proc "contextless" (r: []byte, a: ^Poly) #no_bounds_check {
 			r = r[5:]
 		}
 	case:
-		panic_contextless("crypto/mlkem: invalid POLYCOMPRESSEDBYTES")
+		unreachable()
 	}
 }
 
@@ -91,7 +91,7 @@ poly_decompress :: proc "contextless" (r: ^Poly, a: []byte) {
 			}
 		}
 	case:
-		panic_contextless("crypto/mlkem: invalid POLYCOMPRESSEDBYTES")
+		unreachable()
 	}
 }
 
@@ -236,6 +236,6 @@ poly_compressed_bytes :: #force_inline proc "contextless" (k: int) -> int {
 	case K_1024:
 		return POLYCOMPRESSEDBYTES_1024
 	case:
-		panic_contextless("crypto/mlkem: invalid k")
+		unreachable()
 	}
 }
diff --git a/core/crypto/_mlkem/polyvec.odin b/core/crypto/_mlkem/polyvec.odin
index 387acade8..2b971d2df 100644
--- a/core/crypto/_mlkem/polyvec.odin
+++ b/core/crypto/_mlkem/polyvec.odin
@@ -75,7 +75,7 @@ polyvec_compress :: proc "contextless" (r: []byte, a: ^Polyvec, kay: int) #no_bo
 			}
 		}
 	case:
-		panic_contextless("crypto/mlkem: invalid POLYVECCOMPRESSEDBYTES")
+		unreachable()
 	}
 }
 
@@ -123,7 +123,7 @@ polyvec_decompress :: proc "contextless" (r: ^Polyvec, a: []byte, kay: int) #no_
 			}
 		}
 	case:
-		panic_contextless("crypto/mlkem: invalid POLYVECCOMPRESSEDBYTES")
+		unreachable()
 	}
 }
 
diff --git a/core/crypto/mlkem/api.odin b/core/crypto/mlkem/api.odin
index 9c8fd6a2b..3e3956ba8 100644
--- a/core/crypto/mlkem/api.odin
+++ b/core/crypto/mlkem/api.odin
@@ -253,6 +253,28 @@ decaps :: proc(dk: ^Decapsulation_Key, ciphertext, shared_secret: []byte) -> boo
 	return true
 }
 
+// params returns the Parameters used by a Decapsulation_Key or
+// Encapsulation_Key instance.
+@(require_results)
+params :: proc(k: ^$T) -> Parameters where (T == Encapsulation_Key || T == Decapsulation_Key) {
+	when T == Encapsulation_Key {
+		return k_to_params(k.pke_ek.k)
+	} else {
+		return k_to_params(k.pke_dk.k)
+	}
+}
+
+// key_size returns the key size of a Decapsulation_Key or Encapsulation_Key
+// in bytes.
+@(require_results)
+key_size :: proc(k: ^$T) -> int where (T == Encapsulation_Key || T == Decapsulation_Key) {
+	when T == Encapsulation_Key {
+		return ENCAPSULATION_KEY_SIZES[k.pke_ek.k]
+	} else {
+		return DECAPSULATION_KEY_SEED_SIZE
+	}
+}
+
 @(private="file")
 params_to_k :: #force_inline proc "contextless" (params: Parameters) -> int {
 	#partial switch params {
@@ -266,3 +288,17 @@ params_to_k :: #force_inline proc "contextless" (params: Parameters) -> int {
 
 	return 0
 }
+
+@(private="file")
+k_to_params :: #force_inline proc "contextless" (k: int) -> Parameters {
+	switch k {
+	case _mlkem.K_512:
+		return .ML_KEM_512
+	case _mlkem.K_768:
+		return .ML_KEM_768
+	case _mlkem.K_1024:
+		return .ML_KEM_1024
+	}
+
+	return .Invalid
+}