encode ciphertext into bytes instead of polynomial matrices

kyber/decrypt.py

@@ -1,9 +1,9 @@
 import numpy as np
 from numpy.polynomial.polynomial import Polynomial
-from kyber.utils.compression import compress
+from kyber.utils.compression import compress, decompress
 from kyber.utils.encoding import encode, decode
 from kyber.utils.modulo import polmod
-from kyber.constants import k
+from kyber.constants import n, k, du, dv
 
 class Decrypt:
     def __init__(self, private_key, ciphertext) -> None:
@@ -22,7 +22,16 @@ def decrypt(self) -> bytes:
         s = np.array([
             decode(self._sk[32*12*i : 32*12*(i+1)], 12) for i in range(len(self._sk)//(32*12))
         ])
-        u, v = self._c
+
+        u, v = self._c[:du*k*n//8], self._c[du*k*n//8:]
+
+        u = np.array([
+            decode(u[32*du*i : 32*du*(i+1)], du) for i in range(len(u)//(32*du))
+        ])
+        v = decode(v, dv)
+
+        u = np.array([decompress(pol, du) for pol in u])
+        v = decompress(v, dv)
 
         m: Polynomial = v - np.matmul(s.T, u)
         m = polmod(m)

kyber/encrypt.py

@@ -3,10 +3,10 @@
 from numpy.polynomial.polynomial import Polynomial
 from kyber.utils.cbd import cbd
 from kyber.utils.pseudo_random import prf
-from kyber.utils.modulo import polmod
-from kyber.utils.compression import decompress
-from kyber.utils.encoding import decode
-from kyber.constants import k, eta1, eta2, n
+from kyber.utils.modulo import matmod, polmod
+from kyber.utils.compression import compress, decompress
+from kyber.utils.encoding import encode, decode
+from kyber.constants import k, eta1, eta2, n, du, dv
 from kyber.utils.byte_conversion import int_to_bytes
 from kyber.utils.parse import parse
 from kyber.utils.pseudo_random import xof
@@ -62,7 +62,16 @@ def encrypt(self):
         u = np.matmul(A.T, r) + e1
         v = np.matmul(t.T, r) + e2 + decompress(decode(m, 1), 1)
 
-        u = [polmod(item) for item in u]
+        u = matmod(u)
         v = polmod(v)
+
+        u = compress(u, du)
+        v = compress([v], dv)
+
+        u = encode(u, du)
+        v = encode(v, dv)
 
-        return (u, v)
+        assert len(u) == du * k * n//8
+        assert len(v) == dv * n//8
+
+        return u + v