From ecabea6cd315eaf08571ef61ff267318140d147b Mon Sep 17 00:00:00 2001 From: Ben Perez Date: Wed, 14 May 2025 19:38:34 +0000 Subject: [PATCH] 8347608: Optimize Java implementation of ML-KEM Reviewed-by: weijun --- .../com/sun/crypto/provider/ML_KEM.java | 22 +++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/src/java.base/share/classes/com/sun/crypto/provider/ML_KEM.java b/src/java.base/share/classes/com/sun/crypto/provider/ML_KEM.java index b45b655e1f3..fa461e45d22 100644 --- a/src/java.base/share/classes/com/sun/crypto/provider/ML_KEM.java +++ b/src/java.base/share/classes/com/sun/crypto/provider/ML_KEM.java @@ -703,7 +703,7 @@ public final class ML_KEM { private K_PKE_CipherText kPkeEncrypt( K_PKE_EncryptionKey publicKey, byte[] message, byte[] sigma) { - short[][] zeroes = new short[mlKem_k][ML_KEM_N]; + short[][] zeroes = shortMatrixAlloc(mlKem_k, ML_KEM_N); byte[] pkBytes = publicKey.keyBytes; byte[] rho = Arrays.copyOfRange(pkBytes, pkBytes.length - 32, pkBytes.length); @@ -792,7 +792,7 @@ public final class ML_KEM { System.arraycopy(rho, 0, seedBuf, 0, rho.length); seedBuf[rhoLen + 2] = 0x1F; seedBuf[XOF_BLOCK_LEN - 1] = (byte)0x80; - byte[][] xofBufArr = new byte[nrPar][XOF_BLOCK_LEN + XOF_PAD]; + byte[][] xofBufArr = byteMatrixAlloc(nrPar, XOF_BLOCK_LEN + XOF_PAD); int[] iIndex = new int[nrPar]; int[] jIndex = new int[nrPar]; @@ -1550,4 +1550,22 @@ public final class ML_KEM { return (aHigh - ((m * MONT_Q) >> MONT_R_BITS)); // subtract signed high product } + + // For multidimensional array initialization, manually allocating each entry is + // faster than doing the entire initialization in one go + static short[][] shortMatrixAlloc(int first, int second) { + short[][] res = new short[first][]; + for (int i = 0; i < first; i++) { + res[i] = new short[second]; + } + return res; + } + + static byte[][] byteMatrixAlloc(int first, int second) { + byte[][] res = new byte[first][]; + for (int i = 0; i < first; i++) { + res[i] = new byte[second]; + } + return res; + } }