@@ -672,10 +672,8 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c)
672672 sword16 y [3 * MLKEM_MAX_K * MLKEM_N ];
673673#endif
674674#endif
675- #ifdef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM
676675 sword16 * u ;
677676 sword16 * v ;
678- #endif
679677
680678 /* Establish parameters based on key type. */
681679 switch (key -> type ) {
@@ -741,21 +739,15 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c)
741739 }
742740#endif
743741
744- if (ret == 0 ) {
745742#ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM
743+ if (ret == 0 ) {
746744 /* Assign allocated dynamic memory to pointers.
747745 * y (b) | a (m) | mu (p) | e1 (p) | e2 (v) | u (v) | v (p) */
748746 a = y + MLKEM_N * k ;
749747 mu = a + MLKEM_N * k * k ;
750748 e1 = mu + MLKEM_N ;
751749 e2 = e1 + MLKEM_N * k ;
752- #else
753- /* Assign allocated dynamic memory to pointers.
754- * y (v) | a (v) | u (v) */
755- a = y + MLKEM_N * k ;
756- #endif
757750
758- #ifndef WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM
759751 /* Convert msg to a polynomial.
760752 * Step 20: mu <- Decompress_1(ByteDecode_1(m)) */
761753 mlkem_from_msg (mu , m );
@@ -767,7 +759,7 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c)
767759 */
768760 ret = mlkem_get_noise (& key -> prf , k , y , e1 , e2 , r );
769761 }
770- #ifdef WOLFSSL_MLKEM_CACHE_A
762+ #ifdef WOLFSSL_MLKEM_CACHE_A
771763 if ((ret == 0 ) && ((key -> flags & MLKEM_FLAG_A_SET ) != 0 )) {
772764 unsigned int i ;
773765 /* Transpose matrix.
@@ -782,25 +774,28 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c)
782774 }
783775 }
784776 else
785- #endif
777+ #endif /* WOLFSSL_MLKEM_CACHE_A */
786778 if (ret == 0 ) {
787779 /* Generate the transposed matrix.
788780 * Step 4-8: generate matrix A_hat */
789781 ret = mlkem_gen_matrix (& key -> prf , a , k , key -> pubSeed , 1 );
790782 }
791783 if (ret == 0 ) {
792- sword16 * u ;
793- sword16 * v ;
794-
795784 /* Assign remaining allocated dynamic memory to pointers.
796785 * y (v) | a (m) | mu (p) | e1 (p) | r2 (v) | u (v) | v (p)*/
797786 u = e2 + MLKEM_N ;
798787 v = u + MLKEM_N * k ;
799788
800789 /* Perform encapsulation maths.
801790 * Steps 18-19, 21: calculate u and v */
802- mlkem_encapsulate (key -> pub , u , v , a , y , e1 , e2 , mu , k );
803- #else
791+ ret = mlkem_encapsulate (key -> pub , u , v , a , y , e1 , e2 , mu , k );
792+ }
793+ #else /* WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */
794+ if (ret == 0 ) {
795+ /* Assign allocated dynamic memory to pointers.
796+ * y (v) | a (v) | u (v) */
797+ a = y + MLKEM_N * k ;
798+
804799 /* Initialize the PRF for use in the noise generation. */
805800 mlkem_prf_init (& key -> prf );
806801 /* Generate noise using PRF.
@@ -819,41 +814,39 @@ static int mlkemkey_encapsulate(MlKemKey* key, const byte* m, byte* r, byte* c)
819814 ret = mlkem_encapsulate_seeds (key -> pub , & key -> prf , u , a , y , k , m ,
820815 key -> pubSeed , r );
821816 }
817+ #endif /* WOLFSSL_MLKEM_ENCAPSULATE_SMALL_MEM */
818+
822819 if (ret == 0 ) {
823- #endif
824- {
825- byte * c1 = c ;
826- byte * c2 = c + compVecSz ;
827-
828- #if defined(WOLFSSL_KYBER512 ) || defined(WOLFSSL_WC_ML_KEM_512 )
829- if (k == WC_ML_KEM_512_K ) {
830- /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */
831- mlkem_vec_compress_10 (c1 , u , k );
832- /* Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */
833- mlkem_compress_4 (c2 , v );
834- /* Step 24: return c <- (c_1||c_2) */
835- }
836- #endif
837- #if defined(WOLFSSL_KYBER768 ) || defined(WOLFSSL_WC_ML_KEM_768 )
838- if (k == WC_ML_KEM_768_K ) {
839- /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */
840- mlkem_vec_compress_10 (c1 , u , k );
841- /* Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */
842- mlkem_compress_4 (c2 , v );
843- /* Step 24: return c <- (c_1||c_2) */
844- }
845- #endif
846- #if defined(WOLFSSL_KYBER1024 ) || defined(WOLFSSL_WC_ML_KEM_1024 )
847- if (k == WC_ML_KEM_1024_K ) {
848- /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */
849- mlkem_vec_compress_11 (c1 , u );
850- /* Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */
851- mlkem_compress_5 (c2 , v );
852- /* Step 24: return c <- (c_1||c_2) */
853- }
854- #endif
820+ byte * c1 = c ;
821+ byte * c2 = c + compVecSz ;
855822
823+ #if defined(WOLFSSL_KYBER512 ) || defined(WOLFSSL_WC_ML_KEM_512 )
824+ if (k == WC_ML_KEM_512_K ) {
825+ /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */
826+ mlkem_vec_compress_10 (c1 , u , k );
827+ /* Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */
828+ mlkem_compress_4 (c2 , v );
829+ /* Step 24: return c <- (c_1||c_2) */
856830 }
831+ #endif
832+ #if defined(WOLFSSL_KYBER768 ) || defined(WOLFSSL_WC_ML_KEM_768 )
833+ if (k == WC_ML_KEM_768_K ) {
834+ /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */
835+ mlkem_vec_compress_10 (c1 , u , k );
836+ /* Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */
837+ mlkem_compress_4 (c2 , v );
838+ /* Step 24: return c <- (c_1||c_2) */
839+ }
840+ #endif
841+ #if defined(WOLFSSL_KYBER1024 ) || defined(WOLFSSL_WC_ML_KEM_1024 )
842+ if (k == WC_ML_KEM_1024_K ) {
843+ /* Step 22: c_1 <- ByteEncode_d_u(Compress_d_u(u)) */
844+ mlkem_vec_compress_11 (c1 , u );
845+ /* Step 23: c_2 <- ByteEncode_d_v(Compress_d_v(v)) */
846+ mlkem_compress_5 (c2 , v );
847+ /* Step 24: return c <- (c_1||c_2) */
848+ }
849+ #endif
857850 }
858851
859852#ifndef WOLFSSL_NO_MALLOC
0 commit comments