From d6ca6e617ef183836ad07bb2e7358eb24c333769 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Tue, 21 Nov 2023 15:48:24 +0200
Subject: [PATCH 01/15] added polynomial ring data structure

---
 kyber/encryption/decrypt.py |  2 +-
 kyber/encryption/encrypt.py | 12 +++----
 kyber/encryption/keygen.py  | 10 +++---
 kyber/entities/polring.py   | 65 +++++++++++++++++++++++++++++++++++++
 kyber/utils/cbd.py          |  5 +--
 kyber/utils/compression.py  |  9 ++---
 kyber/utils/encoding.py     |  7 ++--
 kyber/utils/parse.py        |  5 +--
 tests/test_polring.py       | 39 ++++++++++++++++++++++
 9 files changed, 128 insertions(+), 26 deletions(-)
 create mode 100644 kyber/entities/polring.py
 create mode 100644 tests/test_polring.py

diff --git a/kyber/encryption/decrypt.py b/kyber/encryption/decrypt.py
index cdca206..b297fe5 100644
--- a/kyber/encryption/decrypt.py
+++ b/kyber/encryption/decrypt.py
@@ -4,6 +4,7 @@
 from kyber.utils.encoding import encode, decode
 from kyber.utils.modulo import polmod
 from kyber.constants import n, k, du, dv
+from kyber.entities.polring import PolynomialRing
 
 class Decrypt:
     def __init__(self, private_key, ciphertext) -> None:
@@ -31,7 +32,6 @@ def decrypt(self) -> bytes:
         v = decompress(v, dv)
 
         m: Polynomial = v - np.matmul(s.T, u)
-        m = polmod(m)
         m: bytes = encode(compress([m], 1), 1)
 
         assert len(m) == 32
diff --git a/kyber/encryption/encrypt.py b/kyber/encryption/encrypt.py
index fbcb215..46944ab 100644
--- a/kyber/encryption/encrypt.py
+++ b/kyber/encryption/encrypt.py
@@ -10,6 +10,7 @@
 from kyber.utils.byte_conversion import int_to_bytes
 from kyber.utils.parse import parse
 from kyber.utils.pseudo_random import xof
+from kyber.entities.polring import PolynomialRing
 
 class Encrypt:
     def __init__(self, public_key: bytes, m: bytes = None, r: bytes = None) -> None:
@@ -38,22 +39,20 @@ def encrypt(self):
         t, rho = self._pk[:-32], self._pk[-32:]
         t = np.array(decode(t, 12))
 
-        A = np.empty((k, k), Polynomial)
+        A = np.empty((k, k), PolynomialRing)
         for i in range(k):
             for j in range(k):
                 A[i][j] = parse(xof(rho, int_to_bytes(i), int_to_bytes(j)))
 
         N = 0
-        r = np.empty((k, ), Polynomial)
+        r = np.empty((k, ), PolynomialRing)
         for i in range(k):
             r[i] = cbd(prf(rb, int_to_bytes(N)), eta1)
-            r[i] = polmod(r[i])
             N += 1
 
-        e1 = np.empty((k, ), Polynomial)
+        e1 = np.empty((k, ), PolynomialRing)
         for i in range(k):
             e1[i] = cbd(prf(rb, int_to_bytes(N)), eta2)
-            e1[i] = polmod(e1[i])
             N += 1
 
         e2 = cbd(prf(rb, int_to_bytes(N)), eta2)
@@ -62,9 +61,6 @@ def encrypt(self):
         u = np.matmul(A.T, r) + e1
         v = np.matmul(t.T, r) + e2 + decompress(decode(m, 1)[0], 1)
 
-        u = matmod(u)
-        v = polmod(v)
-
         u = compress(u, du)
         v = compress([v], dv)
 
diff --git a/kyber/encryption/keygen.py b/kyber/encryption/keygen.py
index 3e3c99b..646c77f 100644
--- a/kyber/encryption/keygen.py
+++ b/kyber/encryption/keygen.py
@@ -8,6 +8,7 @@
 from kyber.utils.byte_conversion import int_to_bytes
 from kyber.utils.encoding import encode
 from kyber.utils.parse import parse
+from kyber.entities.polring import PolynomialRing
 
 def generate_keys() -> tuple:
     """
@@ -18,26 +19,23 @@ def generate_keys() -> tuple:
     d = randbytes(32)
     rho, sigma = G(d)[:32], G(d)[32:]
 
-    A = np.empty((k, k), Polynomial)
+    A = np.empty((k, k), PolynomialRing)
     for i in range(k):
         for j in range(k):
             A[i][j] = parse(xof(rho, int_to_bytes(i), int_to_bytes(j)))
 
     N = 0
-    s = np.empty((k, ), Polynomial)
+    s = np.empty((k, ), PolynomialRing)
     for i in range(k):
         s[i] = cbd(prf(sigma, int_to_bytes(N)), eta1)
-        s[i] = polmod(s[i])
         N += 1
 
-    e = np.empty((k, ), Polynomial)
+    e = np.empty((k, ), PolynomialRing)
     for i in range(k):
         e[i] = cbd(prf(sigma, int_to_bytes(N)), eta1)
-        e[i] = polmod(e[i])
         N += 1
 
     t = np.matmul(A, s) + e     # t is a polynomial matrix with shape (k, )
-    t = matmod(t)
 
     s: bytes = encode(s, 12)
     t: bytes = encode(t, 12)
diff --git a/kyber/entities/polring.py b/kyber/entities/polring.py
new file mode 100644
index 0000000..8f12117
--- /dev/null
+++ b/kyber/entities/polring.py
@@ -0,0 +1,65 @@
+from kyber.constants import q, n
+from numpy.polynomial.polynomial import Polynomial
+
+class PolynomialRing:
+    def __init__(self, coefs: list[int]) -> None:
+        """Input `(1, 2, 3)` represents `1+2x+3x^2`."""
+        self._coefs = coefs
+        self._coef_limit = q
+        self._degree_limit = n-1
+        self._apply_limits()
+
+    @property
+    def coefs(self) -> list[int]:
+        return self._coefs
+
+    def _apply_limits(self) -> None:
+        # apply degree limit
+        divisor = Polynomial([1] + [0 for _ in range(self._degree_limit)] + [1])    # x^n + 1
+        self._coefs = [int(c) for c in (Polynomial(self.coefs) % divisor).coef]
+
+        # apply coef limit
+        for i in range(len(self._coefs)):
+            self._coefs[i] %= self._coef_limit
+
+        # remove trailing zero coefficients
+        while len(self._coefs) > 0 and self._coefs[-1] == 0:
+            self._coefs.pop()
+
+    def __add__(self, other: "PolynomialRing") -> "PolynomialRing":
+        result = []
+        for i in range(max(len(self.coefs), len(other.coefs))):
+            self_coef = self.coefs[i] if i < len(self.coefs) else 0
+            other_coef = other.coefs[i] if i < len(other.coefs) else 0
+            result.append(self_coef + other_coef)
+        return PolynomialRing(result)
+
+    def __sub__(self, other: "PolynomialRing") -> "PolynomialRing":
+        result = []
+        for i in range(max(len(self.coefs), len(other.coefs))):
+            self_coef = self.coefs[i] if i < len(self.coefs) else 0
+            other_coef = other.coefs[i] if i < len(other.coefs) else 0
+            result.append(self_coef - other_coef)
+        return PolynomialRing(result)
+
+    def __mul__(self, other: "PolynomialRing") -> "PolynomialRing":
+        result = [0 for _ in range(256)]
+        for a in range(len(self.coefs)):
+            for b in range(len(other.coefs)):
+                # check if the term of this degree would be too high
+                if a+b > 255:
+                    continue
+                result[a+b] += self.coefs[a] * other.coefs[b]
+        return PolynomialRing(result)
+
+    def __eq__(self, other: "PolynomialRing") -> bool:
+        return self.coefs == other.coefs
+
+    # modulo with another polring (is needed?)
+
+    # subtraction
+
+    # equal with another (at least for debug)
+
+    def __repr__(self) -> str:
+        return "PolRing(" + ", ".join([str(c) for c in self.coefs]) + ")"
diff --git a/kyber/utils/cbd.py b/kyber/utils/cbd.py
index 60abbee..6294702 100644
--- a/kyber/utils/cbd.py
+++ b/kyber/utils/cbd.py
@@ -1,7 +1,8 @@
 from numpy.polynomial.polynomial import Polynomial
 from kyber.utils.byte_conversion import bytes_to_bits
+from kyber.entities.polring import PolynomialRing
 
-def cbd(b: bytes, eta: int) -> Polynomial:
+def cbd(b: bytes, eta: int) -> PolynomialRing:
     """
     Deterministically creates and returns a polynomial (degree 255)
     from the given byte array (length 64*eta).
@@ -20,4 +21,4 @@ def cbd(b: bytes, eta: int) -> Polynomial:
             b += bits[2 * i * eta + eta + j]
         f.append(a-b)
     assert len(f) == 256
-    return Polynomial(f)
+    return PolynomialRing(f)
diff --git a/kyber/utils/compression.py b/kyber/utils/compression.py
index 6dc47c3..3a8e8e7 100644
--- a/kyber/utils/compression.py
+++ b/kyber/utils/compression.py
@@ -2,8 +2,9 @@
 from numpy.polynomial.polynomial import Polynomial
 from kyber.constants import q
 from kyber.utils.round import normal_round
+from kyber.entities.polring import PolynomialRing
 
-def compress(pols: list[Polynomial], d: int) -> list[Polynomial]:
+def compress(pols: list[PolynomialRing], d: int) -> list[PolynomialRing]:
     """
     Reduces every coefficient of every polynomial in the given list
     to range `0...2**d-1` (inclusive).
@@ -11,15 +12,15 @@ def compress(pols: list[Polynomial], d: int) -> list[Polynomial]:
     result = []
     for pol in pols:
         f = [compress_int(c, d) for c in pol.coef]
-        result.append(Polynomial(f))
+        result.append(PolynomialRing(f))
     return result
 
-def decompress(pol: Polynomial, d: int) -> Polynomial:
+def decompress(pol: PolynomialRing, d: int) -> PolynomialRing:
     """
     Multiplies each coefficient of the given polynomial by `q/(2**d)`.
     Each coefficient of the given polynomial must be in range `0...2^d-1` (inclusive).
     """
-    return Polynomial([decompress_int(c, d) for c in pol.coef])
+    return PolynomialRing([decompress_int(c, d) for c in pol.coef])
 
 def compress_int(x: int, d: int) -> int:
     """
diff --git a/kyber/utils/encoding.py b/kyber/utils/encoding.py
index 2aaa533..653e2f7 100644
--- a/kyber/utils/encoding.py
+++ b/kyber/utils/encoding.py
@@ -1,8 +1,9 @@
 import numpy as np
 from numpy.polynomial.polynomial import Polynomial
 from kyber.utils.byte_conversion import bytes_to_bits, bits_to_bytes
+from kyber.entities.polring import PolynomialRing
 
-def encode(pols: list[Polynomial], l: int) -> bytes:
+def encode(pols: list[PolynomialRing], l: int) -> bytes:
     """
     Converts the given polynomial (degree 255, each coefficient in range `0...2**l-1` inclusive)
     into a byte array of lenght `32*l`.
@@ -32,7 +33,7 @@ def encode(pols: list[Polynomial], l: int) -> bytes:
     assert len(result) == 32*l*len(pols)
     return bytes(result)
 
-def decode(b: bytes, l: int) -> list[Polynomial]:
+def decode(b: bytes, l: int) -> list[PolynomialRing]:
     """
     Converts the given byte array (length `32*l*x` for some integer x) into
     a list of polynomials (length x, each degree 255)
@@ -48,5 +49,5 @@ def decode(b: bytes, l: int) -> list[Polynomial]:
         for i in range(256):
             f[i] = sum(bits[i*l+j]*2**j for j in range(l))      # accesses each bit exactly once
             assert 0 <= f[i] and f[i] <= 2**l-1
-        result.append(Polynomial(f))
+        result.append(PolynomialRing(f))
     return result
diff --git a/kyber/utils/parse.py b/kyber/utils/parse.py
index 1f320bb..2744b17 100644
--- a/kyber/utils/parse.py
+++ b/kyber/utils/parse.py
@@ -3,12 +3,13 @@
 from numpy.polynomial.polynomial import Polynomial
 import numpy as np
 from kyber.constants import n, q
+from kyber.entities.polring import PolynomialRing
 
 def byte_to_int(b: bytes) -> int:
     """Returns the unsigned integer that the given big-endian byte array represents."""
     return int.from_bytes(b)
 
-def parse(stream: Generator[bytes, None, None]) -> Polynomial:
+def parse(stream: Generator[bytes, None, None]) -> PolynomialRing:
     """
     Deterministically creates a polynomial (degree n-1, each coefficient in
     range `0...4095` inclusive) from the given bytestream.
@@ -27,4 +28,4 @@ def parse(stream: Generator[bytes, None, None]) -> Polynomial:
             a[j] = d2
             j += 1
         i += 3
-    return Polynomial(a)
+    return PolynomialRing(a)
diff --git a/tests/test_polring.py b/tests/test_polring.py
new file mode 100644
index 0000000..bec908c
--- /dev/null
+++ b/tests/test_polring.py
@@ -0,0 +1,39 @@
+import unittest
+from random import seed, randint
+from numpy.polynomial.polynomial import Polynomial
+from kyber.entities.polring import PolynomialRing
+from kyber.constants import q, n
+
+# q = 3329
+
+class TestPolynomialRing(unittest.TestCase):
+    def test_initialization(self):
+        pol = PolynomialRing([71, -5, 0, 1, 3329, 3330, 3328, 15])
+        self.assertListEqual(pol.coefs, [71, 3324, 0, 1, 0, 1, 3328, 15])
+
+    def test_init_with_random_inputs(self):
+        # test with 1000 samples that randomily initialized polring matches expected
+        seed(42)
+        for _ in range(1000):
+            coef_count = randint(1, 300)
+            coefs = [
+                (1 if randint(1, 2) == 1 else -1) * randint(1, 10000)
+                for _ in range(coef_count)
+            ]
+            pol = PolynomialRing(coefs)
+            self.assertEqual(len(pol.coefs), min(coef_count, 256))
+            expected_result = Polynomial(coefs) % Polynomial([1] + [0 for _ in range(n-1)] + [1])
+            for i in range(len(pol.coefs)):
+                self.assertEqual(pol.coefs[i], expected_result.coef[i] % q)
+
+    def test_sum(self):
+        pol1 = PolynomialRing([581,  -50,  100, 31, -4500, 4567, 11, 12])
+        pol2 = PolynomialRing([1986, -150, -99, 34, 500,   0])
+        # sum before mod [2567, -200, 1, 65, -4000, 4567, 11, 12]
+        self.assertListEqual((pol1 + pol2).coefs, [2567, 3129, 1, 65, 2658, 1238, 11, 12])
+
+    def test_multiplication(self):
+        pol1 = PolynomialRing([15, -13, 0, 7,   472, -88,  112, 5])
+        pol2 = PolynomialRing([2,  17,  8, 590, -11, -101, 91])
+        # product before mod [30, 229, -101, 8760, -6772, 6532, 9312, 278430, -56838, 20053, 53558, -19375, 9687, 455]
+        self.assertListEqual((pol1 * pol2).coefs, [30, 229, 3228, 2102, 3215, 3203, 2654, 2123, 3084, 79, 294, 599, 3029, 455])

From 1650eb65c1cbcdf1fe298b90b09b8943cce11cc6 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Tue, 21 Nov 2023 15:56:50 +0200
Subject: [PATCH 02/15] fixing broken tests and bugs

---
 kyber/encryption/encrypt.py |  1 -
 kyber/utils/compression.py  |  4 ++--
 kyber/utils/encoding.py     |  4 +---
 tests/test_cbd.py           | 11 ++++++-----
 tests/test_compression.py   | 21 +++++++++++----------
 tests/test_encoding.py      | 18 ++++++------------
 tests/test_parse.py         |  4 ++--
 7 files changed, 28 insertions(+), 35 deletions(-)

diff --git a/kyber/encryption/encrypt.py b/kyber/encryption/encrypt.py
index 46944ab..e7520fc 100644
--- a/kyber/encryption/encrypt.py
+++ b/kyber/encryption/encrypt.py
@@ -56,7 +56,6 @@ def encrypt(self):
             N += 1
 
         e2 = cbd(prf(rb, int_to_bytes(N)), eta2)
-        e2 = polmod(e2)
 
         u = np.matmul(A.T, r) + e1
         v = np.matmul(t.T, r) + e2 + decompress(decode(m, 1)[0], 1)
diff --git a/kyber/utils/compression.py b/kyber/utils/compression.py
index 3a8e8e7..3b73585 100644
--- a/kyber/utils/compression.py
+++ b/kyber/utils/compression.py
@@ -11,7 +11,7 @@ def compress(pols: list[PolynomialRing], d: int) -> list[PolynomialRing]:
     """
     result = []
     for pol in pols:
-        f = [compress_int(c, d) for c in pol.coef]
+        f = [compress_int(c, d) for c in pol.coefs]
         result.append(PolynomialRing(f))
     return result
 
@@ -20,7 +20,7 @@ def decompress(pol: PolynomialRing, d: int) -> PolynomialRing:
     Multiplies each coefficient of the given polynomial by `q/(2**d)`.
     Each coefficient of the given polynomial must be in range `0...2^d-1` (inclusive).
     """
-    return PolynomialRing([decompress_int(c, d) for c in pol.coef])
+    return PolynomialRing([decompress_int(c, d) for c in pol.coefs])
 
 def compress_int(x: int, d: int) -> int:
     """
diff --git a/kyber/utils/encoding.py b/kyber/utils/encoding.py
index 653e2f7..73bb06c 100644
--- a/kyber/utils/encoding.py
+++ b/kyber/utils/encoding.py
@@ -13,9 +13,7 @@ def encode(pols: list[PolynomialRing], l: int) -> bytes:
 
     result = bytearray()
     for pol in pols:
-        if len(pol.coef) > 256:
-            raise ValueError("too high polynomial degree")
-        f = list(pol.coef) + [0 for _ in range(256-len(pol.coef))]
+        f = list(pol.coefs) + [0 for _ in range(256-len(pol.coefs))]
         bits = np.empty((256*l, ))
         for i in range(256):
             f_item = f[i]
diff --git a/tests/test_cbd.py b/tests/test_cbd.py
index c946f2f..bb4ff08 100644
--- a/tests/test_cbd.py
+++ b/tests/test_cbd.py
@@ -25,8 +25,9 @@ def test_cbd_throws_with_incorrect_argument_length(self):
         with self.assertRaises(ValueError):
             cbd(argument, eta)
 
-    def test_cbd_result_polynomial_degree(self):
-        eta = 5
-        argument = randbytes(320)
-        result = cbd(argument, eta)
-        self.assertEqual(len(result.coef), 256)
+    # TODO: test with random iterations
+    # def test_cbd_result_polynomial_degree(self):
+    #     eta = 5
+    #     argument = randbytes(320)
+    #     result = cbd(argument, eta)
+    #     self.assertEqual(len(result.coefs), 256)
diff --git a/tests/test_compression.py b/tests/test_compression.py
index cac6d40..30c14ce 100644
--- a/tests/test_compression.py
+++ b/tests/test_compression.py
@@ -2,42 +2,43 @@
 from numpy.polynomial.polynomial import Polynomial
 from random import seed, randint
 from kyber.utils.compression import compress, decompress
+from kyber.entities.polring import PolynomialRing
 
 class TestCompression(unittest.TestCase):
     def test_compression_symmetry(self):
         seed(42)
-        polynomial = Polynomial([randint(0, 2047) for _ in range(256)])
+        polynomial = PolynomialRing([randint(0, 2047) for _ in range(256)])
         decompressed = decompress(polynomial, 11)
         compressed = compress([decompressed], 11)[0]
-        self.assertListEqual(list(polynomial.coef), list(compressed.coef))
+        self.assertListEqual(list(polynomial.coefs), list(compressed.coefs))
 
     def test_compression(self):
-        polynomial = Polynomial([416, 2913, 0, 1248])
-        expected_result = Polynomial([1, 7, 0, 3])
+        polynomial = PolynomialRing([416, 2913, 0, 1248])
+        expected_result = PolynomialRing([1, 7, 0, 3])
         self.assertEqual(compress([polynomial], 3)[0], expected_result)
 
     def test_compression_result_coefficients_in_range(self):
         # each coefficient in the result should be in range 0...2**d-1 (inclusive)
         d = 11
-        polynomial = Polynomial([4000, -700, 0, 32])
+        polynomial = PolynomialRing([4000, -700, 0, 32])
         result = compress([polynomial], d)[0]
-        for c in result.coef:
+        for c in result.coefs:
             self.assertTrue(0 <= c and c <= 2**d-1)
 
 
     def test_decompression(self):
-        polynomial = Polynomial([1, 7, 0, 3])
-        expected_result = Polynomial([416, 2913, 0, 1248])
+        polynomial = PolynomialRing([1, 7, 0, 3])
+        expected_result = PolynomialRing([416, 2913, 0, 1248])
         self.assertEqual(decompress(polynomial, 3), expected_result)
 
     def test_decompression_raises_with_negative_coefficient(self):
         # coefficient should not be negative
-        polynomial = Polynomial([2, -1, 3])
+        polynomial = PolynomialRing([2, -1, 3])
         with self.assertRaises(ValueError):
             decompress(polynomial, 3)
 
     def test_decompression_raises_with_too_large_coefficient(self):
         # coefficient should not be greather than 2**d-1 = 7
-        polynomial = Polynomial([2, 8, 3])
+        polynomial = PolynomialRing([2, 8, 3])
         with self.assertRaises(ValueError):
             decompress(polynomial, 3)
diff --git a/tests/test_encoding.py b/tests/test_encoding.py
index 0199d5a..9a8a1db 100644
--- a/tests/test_encoding.py
+++ b/tests/test_encoding.py
@@ -2,14 +2,15 @@
 from numpy.polynomial.polynomial import Polynomial
 from random import seed, randbytes, randint
 from kyber.utils.encoding import encode, decode
+from kyber.entities.polring import PolynomialRing
 
 class TestEncoding(unittest.TestCase):
     def setUp(self):
         seed(42)
         self.l = 5
         self.data = randbytes(32 * self.l)
-        self.polynomial = Polynomial([randint(0, 1) for _ in range(256)])
-        self.polynomial2 = Polynomial([randint(0, 1) for _ in range(256)])
+        self.polynomial = PolynomialRing([randint(0, 1) for _ in range(256)])
+        self.polynomial2 = PolynomialRing([randint(0, 1) for _ in range(256)])
 
     def test_encoding_symmetry(self):
         polynomials = decode(self.data, self.l)
@@ -18,12 +19,12 @@ def test_encoding_symmetry(self):
 
     def test_decode_coefficients(self):
         polynomial = decode(self.data, self.l)[0]
-        for c in polynomial.coef:
+        for c in polynomial.coefs:
             self.assertTrue(0 <= int(c) or int(c) <= 2**self.l-1)
 
     def test_decode_degree(self):
         polynomial = decode(self.data, self.l)[0]
-        self.assertEqual(len(polynomial.coef), 256)
+        self.assertEqual(len(polynomial.coefs), 256)
 
     def test_decode_raises_with_invalid_argument_length(self):
         data = self.data + bytes([42])
@@ -42,13 +43,6 @@ def test_encode_raises_with_too_large_coefficient(self):
         # each coefficient should be in range 0...2**l-1 (inclusive)
         coefs = [i % (2**self.l-1) for i in range(1, 257)]
         coefs[4] = 2**self.l
-        polynomial = Polynomial(coefs)
-        with self.assertRaises(ValueError):
-            encode([polynomial], self.l)
-
-    def test_encode_raises_with_too_high_degree(self):
-        # encode should be given a polynomial with degree of 255, not 256
-        coefs = [1 for _ in range(257)]
-        polynomial = Polynomial(coefs)
+        polynomial = PolynomialRing(coefs)
         with self.assertRaises(ValueError):
             encode([polynomial], self.l)
diff --git a/tests/test_parse.py b/tests/test_parse.py
index 631a593..23a9288 100644
--- a/tests/test_parse.py
+++ b/tests/test_parse.py
@@ -12,8 +12,8 @@ def sample_generator(seed: int = 0) -> Generator[bytes, None, None]:
 class TestParse(unittest.TestCase):
     def test_parse_outputted_polynomial_characteristics(self):
         pol = parse(sample_generator())
-        self.assertEqual(len(pol.coef), 256)
-        for c in pol.coef:
+        self.assertEqual(len(pol.coefs), 256)
+        for c in pol.coefs:
             self.assertTrue(0 <= c and c <= 4095)
 
     def test_parse_outputs_same_polynomial_with_same_input(self):

From 647e2c0d87c4a94be13a70c7c993931704fb3525 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Tue, 21 Nov 2023 16:01:24 +0200
Subject: [PATCH 03/15] added one more test for CBD

---
 tests/test_cbd.py | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/tests/test_cbd.py b/tests/test_cbd.py
index bb4ff08..7006af3 100644
--- a/tests/test_cbd.py
+++ b/tests/test_cbd.py
@@ -25,9 +25,15 @@ def test_cbd_throws_with_incorrect_argument_length(self):
         with self.assertRaises(ValueError):
             cbd(argument, eta)
 
-    # TODO: test with random iterations
-    # def test_cbd_result_polynomial_degree(self):
-    #     eta = 5
-    #     argument = randbytes(320)
-    #     result = cbd(argument, eta)
-    #     self.assertEqual(len(result.coefs), 256)
+    def test_cbd_result_polynomial_degree(self):
+        # test that CBD always outputs a polynomial with degree <= 255
+        # also make sure that the degree is not always less than 255
+        eta = 5
+        some_had_degree_255 = False
+        for _ in range(100):
+            argument = randbytes(320)
+            result = cbd(argument, eta)
+            self.assertLessEqual(len(result.coefs), 256)
+            if len(result.coefs) == 256:
+                some_had_degree_255 = True
+        self.assertTrue(some_had_degree_255)

From b6f4a95485791bbd048a7dbc808378f2dc1f88b4 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Tue, 21 Nov 2023 16:04:43 +0200
Subject: [PATCH 04/15] decrease test iterations for performance

---
 tests/test_polring.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_polring.py b/tests/test_polring.py
index bec908c..846d78c 100644
--- a/tests/test_polring.py
+++ b/tests/test_polring.py
@@ -14,7 +14,7 @@ def test_initialization(self):
     def test_init_with_random_inputs(self):
         # test with 1000 samples that randomily initialized polring matches expected
         seed(42)
-        for _ in range(1000):
+        for _ in range(100):
             coef_count = randint(1, 300)
             coefs = [
                 (1 if randint(1, 2) == 1 else -1) * randint(1, 10000)

From de8e269230b892bea8d4c13d0ad244bf1e85b90b Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Wed, 22 Nov 2023 15:51:20 +0200
Subject: [PATCH 05/15] improved polynomial ring modulo calculation performance

---
 kyber/entities/polring.py | 28 ++++++++++++++++++++++++----
 tests/test_polring.py     |  2 +-
 2 files changed, 25 insertions(+), 5 deletions(-)

diff --git a/kyber/entities/polring.py b/kyber/entities/polring.py
index 8f12117..d2d563c 100644
--- a/kyber/entities/polring.py
+++ b/kyber/entities/polring.py
@@ -4,7 +4,7 @@
 class PolynomialRing:
     def __init__(self, coefs: list[int]) -> None:
         """Input `(1, 2, 3)` represents `1+2x+3x^2`."""
-        self._coefs = coefs
+        self._coefs = [int(c) for c in coefs]
         self._coef_limit = q
         self._degree_limit = n-1
         self._apply_limits()
@@ -14,9 +14,8 @@ def coefs(self) -> list[int]:
         return self._coefs
 
     def _apply_limits(self) -> None:
-        # apply degree limit
-        divisor = Polynomial([1] + [0 for _ in range(self._degree_limit)] + [1])    # x^n + 1
-        self._coefs = [int(c) for c in (Polynomial(self.coefs) % divisor).coef]
+        # apply degree limit by dividing self by x^n+1
+        self._apply_polynomial_modulo_limit()
 
         # apply coef limit
         for i in range(len(self._coefs)):
@@ -26,6 +25,15 @@ def _apply_limits(self) -> None:
         while len(self._coefs) > 0 and self._coefs[-1] == 0:
             self._coefs.pop()
 
+    def _apply_polynomial_modulo_limit(self) -> None:
+        """Replaces `self._coefs` with the remainder of division `self._coefs / (x^n+1)`."""
+        # this is an optimal version of polynomial long division
+        while len(self._coefs) >= n+1:
+            self._coefs[-n-1] -= self._coefs[-1]
+            self._coefs[-1] = 0
+            while self._coefs[-1] == 0:
+                self._coefs.pop()
+
     def __add__(self, other: "PolynomialRing") -> "PolynomialRing":
         result = []
         for i in range(max(len(self.coefs), len(other.coefs))):
@@ -52,6 +60,18 @@ def __mul__(self, other: "PolynomialRing") -> "PolynomialRing":
                 result[a+b] += self.coefs[a] * other.coefs[b]
         return PolynomialRing(result)
 
+    # def modded(self) -> list[int]:
+    #     """Returns the coefs of the remainder of division self.coefs / (x^n+1)."""
+    #     r = self._coefs[:]
+
+    #     while len(r) >= n+1:
+    #         r[-n-1] -= r[-1]
+    #         r[-1] = 0
+    #         while r[-1] == 0:
+    #             r.pop()
+
+    #     return r
+
     def __eq__(self, other: "PolynomialRing") -> bool:
         return self.coefs == other.coefs
 
diff --git a/tests/test_polring.py b/tests/test_polring.py
index 846d78c..308938c 100644
--- a/tests/test_polring.py
+++ b/tests/test_polring.py
@@ -15,7 +15,7 @@ def test_init_with_random_inputs(self):
         # test with 1000 samples that randomily initialized polring matches expected
         seed(42)
         for _ in range(100):
-            coef_count = randint(1, 300)
+            coef_count = randint(1, 500)
             coefs = [
                 (1 if randint(1, 2) == 1 else -1) * randint(1, 10000)
                 for _ in range(coef_count)

From a0ac290e14208f978fbb7cda78f4c42ebfc874fe Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Wed, 22 Nov 2023 15:55:45 +0200
Subject: [PATCH 06/15] removing unnecessary imports

---
 kyber/encryption/decrypt.py | 4 +---
 kyber/encryption/encrypt.py | 2 --
 kyber/encryption/keygen.py  | 2 --
 kyber/entities/polring.py   | 1 -
 kyber/utils/cbd.py          | 1 -
 kyber/utils/compression.py  | 1 -
 kyber/utils/encoding.py     | 1 -
 kyber/utils/parse.py        | 1 -
 tests/test_compression.py   | 1 -
 tests/test_encoding.py      | 1 -
 10 files changed, 1 insertion(+), 14 deletions(-)

diff --git a/kyber/encryption/decrypt.py b/kyber/encryption/decrypt.py
index b297fe5..25b19f3 100644
--- a/kyber/encryption/decrypt.py
+++ b/kyber/encryption/decrypt.py
@@ -1,8 +1,6 @@
 import numpy as np
-from numpy.polynomial.polynomial import Polynomial
 from kyber.utils.compression import compress, decompress
 from kyber.utils.encoding import encode, decode
-from kyber.utils.modulo import polmod
 from kyber.constants import n, k, du, dv
 from kyber.entities.polring import PolynomialRing
 
@@ -31,7 +29,7 @@ def decrypt(self) -> bytes:
         u = np.array([decompress(pol, du) for pol in u])
         v = decompress(v, dv)
 
-        m: Polynomial = v - np.matmul(s.T, u)
+        m: PolynomialRing = v - np.matmul(s.T, u)
         m: bytes = encode(compress([m], 1), 1)
 
         assert len(m) == 32
diff --git a/kyber/encryption/encrypt.py b/kyber/encryption/encrypt.py
index e7520fc..90ade12 100644
--- a/kyber/encryption/encrypt.py
+++ b/kyber/encryption/encrypt.py
@@ -1,9 +1,7 @@
 from random import randbytes
 import numpy as np
-from numpy.polynomial.polynomial import Polynomial
 from kyber.utils.cbd import cbd
 from kyber.utils.pseudo_random import prf
-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
diff --git a/kyber/encryption/keygen.py b/kyber/encryption/keygen.py
index 646c77f..1951d23 100644
--- a/kyber/encryption/keygen.py
+++ b/kyber/encryption/keygen.py
@@ -1,10 +1,8 @@
 from random import randbytes
 import numpy as np
-from numpy.polynomial.polynomial import Polynomial
 from kyber.constants import k, eta1
 from kyber.utils.pseudo_random import prf, G, xof
 from kyber.utils.cbd import cbd
-from kyber.utils.modulo import matmod, polmod
 from kyber.utils.byte_conversion import int_to_bytes
 from kyber.utils.encoding import encode
 from kyber.utils.parse import parse
diff --git a/kyber/entities/polring.py b/kyber/entities/polring.py
index d2d563c..47934ee 100644
--- a/kyber/entities/polring.py
+++ b/kyber/entities/polring.py
@@ -1,5 +1,4 @@
 from kyber.constants import q, n
-from numpy.polynomial.polynomial import Polynomial
 
 class PolynomialRing:
     def __init__(self, coefs: list[int]) -> None:
diff --git a/kyber/utils/cbd.py b/kyber/utils/cbd.py
index 6294702..d87878b 100644
--- a/kyber/utils/cbd.py
+++ b/kyber/utils/cbd.py
@@ -1,4 +1,3 @@
-from numpy.polynomial.polynomial import Polynomial
 from kyber.utils.byte_conversion import bytes_to_bits
 from kyber.entities.polring import PolynomialRing
 
diff --git a/kyber/utils/compression.py b/kyber/utils/compression.py
index 3b73585..5f290d3 100644
--- a/kyber/utils/compression.py
+++ b/kyber/utils/compression.py
@@ -1,5 +1,4 @@
 from math import log2, ceil
-from numpy.polynomial.polynomial import Polynomial
 from kyber.constants import q
 from kyber.utils.round import normal_round
 from kyber.entities.polring import PolynomialRing
diff --git a/kyber/utils/encoding.py b/kyber/utils/encoding.py
index 73bb06c..cbd2d05 100644
--- a/kyber/utils/encoding.py
+++ b/kyber/utils/encoding.py
@@ -1,5 +1,4 @@
 import numpy as np
-from numpy.polynomial.polynomial import Polynomial
 from kyber.utils.byte_conversion import bytes_to_bits, bits_to_bytes
 from kyber.entities.polring import PolynomialRing
 
diff --git a/kyber/utils/parse.py b/kyber/utils/parse.py
index 2744b17..2e96e59 100644
--- a/kyber/utils/parse.py
+++ b/kyber/utils/parse.py
@@ -1,6 +1,5 @@
 from math import floor
 from typing import Generator
-from numpy.polynomial.polynomial import Polynomial
 import numpy as np
 from kyber.constants import n, q
 from kyber.entities.polring import PolynomialRing
diff --git a/tests/test_compression.py b/tests/test_compression.py
index 30c14ce..6c8d33f 100644
--- a/tests/test_compression.py
+++ b/tests/test_compression.py
@@ -1,5 +1,4 @@
 import unittest
-from numpy.polynomial.polynomial import Polynomial
 from random import seed, randint
 from kyber.utils.compression import compress, decompress
 from kyber.entities.polring import PolynomialRing
diff --git a/tests/test_encoding.py b/tests/test_encoding.py
index 9a8a1db..051870c 100644
--- a/tests/test_encoding.py
+++ b/tests/test_encoding.py
@@ -1,5 +1,4 @@
 import unittest
-from numpy.polynomial.polynomial import Polynomial
 from random import seed, randbytes, randint
 from kyber.utils.encoding import encode, decode
 from kyber.entities.polring import PolynomialRing

From c989f1a329b668da6471d2d258a727c0875969e1 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Wed, 22 Nov 2023 15:56:19 +0200
Subject: [PATCH 07/15] removed unnecessary modulo utility

---
 kyber/utils/modulo.py | 44 -------------------------------------------
 tests/test_modulo.py  | 26 -------------------------
 2 files changed, 70 deletions(-)
 delete mode 100644 kyber/utils/modulo.py
 delete mode 100644 tests/test_modulo.py

diff --git a/kyber/utils/modulo.py b/kyber/utils/modulo.py
deleted file mode 100644
index 8bd5624..0000000
--- a/kyber/utils/modulo.py
+++ /dev/null
@@ -1,44 +0,0 @@
-import numpy as np
-from numpy.polynomial.polynomial import Polynomial
-from kyber.constants import n, q
-
-def matmod(
-        matrix: np.ndarray,
-        pol_divisor: Polynomial = None,
-        int_divisor: int = None
-    ) -> np.ndarray:
-    """
-    Applies `polmod` to each element of the given matrix of polynomials.
-    :param matrix A matrix of polynomials
-    :param pol_divisor A polynomial divisor. Only for testing purposes. (default: x^n+1)
-    :param int_divisor An integer divisor. Only for testing purposes. (default: constant q)
-    :returns A new matrix of polynomials.
-    """
-
-    result = matrix.copy()
-    for index, item in np.ndenumerate(matrix):
-        result[index] = polmod(item, pol_divisor, int_divisor)
-    return result
-
-def polmod(pol: Polynomial, pol_divisor: Polynomial = None, int_divisor: int = None) -> Polynomial:
-    """
-    Finds given polynomial modulo another polynomial and coefficients modulo an integer.
-    :param pol The polynomial to be modded
-    :param pol_divisor A polynomial divisor. Only for testing purposes. (default: x^n+1)
-    :param int_divisor An integer divisor. Only for testing purposes. (default: constant q)
-    :returns The remainder of polynomial division whose coefficients are taken modulo int_divisor.
-    """
-
-    # divide the given polynomial by another polynomial and take the remainder
-    if pol_divisor is not None:
-        divisor = pol_divisor
-    else:
-        divisor = Polynomial([1] + [0 for _ in range(n-1)] + [1])       # x^n + 1
-    pol = pol % divisor
-
-    # divide each coefficient of the polynomial by an integer and take the remainder
-    divisor = int_divisor if int_divisor is not None else q
-    for i in range(len(pol.coef)):
-        pol.coef[i] %= divisor
-
-    return pol
diff --git a/tests/test_modulo.py b/tests/test_modulo.py
deleted file mode 100644
index b23a780..0000000
--- a/tests/test_modulo.py
+++ /dev/null
@@ -1,26 +0,0 @@
-import unittest
-import numpy as np
-from numpy.polynomial.polynomial import Polynomial
-from kyber.utils.modulo import matmod, polmod
-
-class TestModulo(unittest.TestCase):
-    def setUp(self):
-        self.polynomial = Polynomial([72, -2, 0, 52, -3, 17])
-        self.pol_divisor = Polynomial([1, 0, 0, 1])
-        self.int_divisor = 7
-
-    def test_polmod(self):
-        # remainder of division polynomial/pol_divisor is -17x^2+x+20
-        # then each coefficient is taken modulo int_divisor to get the expected_result
-        result = polmod(self.polynomial, self.pol_divisor, self.int_divisor)
-        expected_result = Polynomial([6, 1, 4])
-        self.assertEqual(result, expected_result)
-
-    def test_matmod(self):
-        # basically the same test than `test_polmod` but for a matrix of two identical polynomials
-        matrix = np.array([self.polynomial, self.polynomial])
-        expected_result = Polynomial([6, 1, 4])
-        result = matmod(matrix, self.pol_divisor, self.int_divisor)
-        self.assertEqual(result.shape, (2, ))
-        self.assertEqual(result[0], expected_result)
-        self.assertEqual(result[1], expected_result)

From 310b36665b71fce3a9a36a366d70cba502b4b7c6 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Wed, 22 Nov 2023 15:59:01 +0200
Subject: [PATCH 08/15] remove todo comments

---
 kyber/entities/polring.py | 18 ------------------
 1 file changed, 18 deletions(-)

diff --git a/kyber/entities/polring.py b/kyber/entities/polring.py
index 47934ee..2827823 100644
--- a/kyber/entities/polring.py
+++ b/kyber/entities/polring.py
@@ -59,26 +59,8 @@ def __mul__(self, other: "PolynomialRing") -> "PolynomialRing":
                 result[a+b] += self.coefs[a] * other.coefs[b]
         return PolynomialRing(result)
 
-    # def modded(self) -> list[int]:
-    #     """Returns the coefs of the remainder of division self.coefs / (x^n+1)."""
-    #     r = self._coefs[:]
-
-    #     while len(r) >= n+1:
-    #         r[-n-1] -= r[-1]
-    #         r[-1] = 0
-    #         while r[-1] == 0:
-    #             r.pop()
-
-    #     return r
-
     def __eq__(self, other: "PolynomialRing") -> bool:
         return self.coefs == other.coefs
 
-    # modulo with another polring (is needed?)
-
-    # subtraction
-
-    # equal with another (at least for debug)
-
     def __repr__(self) -> str:
         return "PolRing(" + ", ".join([str(c) for c in self.coefs]) + ")"

From c2e5c918ef58ca183925188207caf5327955885a Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Thu, 23 Nov 2023 08:54:58 +0200
Subject: [PATCH 09/15] fixed a bug in polynomial ring multiplication

---
 kyber/entities/polring.py | 10 ++++------
 tests/test_polring.py     | 12 ++++++++++++
 2 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/kyber/entities/polring.py b/kyber/entities/polring.py
index 2827823..5c266e1 100644
--- a/kyber/entities/polring.py
+++ b/kyber/entities/polring.py
@@ -1,12 +1,13 @@
 from kyber.constants import q, n
 
 class PolynomialRing:
-    def __init__(self, coefs: list[int]) -> None:
+    def __init__(self, coefs: list[int], check_limits: bool = True) -> None:
         """Input `(1, 2, 3)` represents `1+2x+3x^2`."""
         self._coefs = [int(c) for c in coefs]
         self._coef_limit = q
         self._degree_limit = n-1
-        self._apply_limits()
+        if check_limits:
+            self._apply_limits()
 
     @property
     def coefs(self) -> list[int]:
@@ -50,12 +51,9 @@ def __sub__(self, other: "PolynomialRing") -> "PolynomialRing":
         return PolynomialRing(result)
 
     def __mul__(self, other: "PolynomialRing") -> "PolynomialRing":
-        result = [0 for _ in range(256)]
+        result = [0 for _ in range(len(self.coefs) + len(other.coefs) - 1)]
         for a in range(len(self.coefs)):
             for b in range(len(other.coefs)):
-                # check if the term of this degree would be too high
-                if a+b > 255:
-                    continue
                 result[a+b] += self.coefs[a] * other.coefs[b]
         return PolynomialRing(result)
 
diff --git a/tests/test_polring.py b/tests/test_polring.py
index 308938c..06052d2 100644
--- a/tests/test_polring.py
+++ b/tests/test_polring.py
@@ -37,3 +37,15 @@ def test_multiplication(self):
         pol2 = PolynomialRing([2,  17,  8, 590, -11, -101, 91])
         # product before mod [30, 229, -101, 8760, -6772, 6532, 9312, 278430, -56838, 20053, 53558, -19375, 9687, 455]
         self.assertListEqual((pol1 * pol2).coefs, [30, 229, 3228, 2102, 3215, 3203, 2654, 2123, 3084, 79, 294, 599, 3029, 455])
+
+    def test_multiplication_with_random_inputs(self):
+        seed(42)
+        for _ in range(100):
+            pol1_degree, pol2_degree = randint(1, 500), randint(1, 500)
+            coefs_1 = [randint(-4000, 4000) for _ in range(pol1_degree)]
+            coefs_2 = [randint(-4000, 4000) for _ in range(pol2_degree)]
+            result = PolynomialRing(coefs_1) * PolynomialRing(coefs_2)
+            expected = (Polynomial(coefs_1) * Polynomial(coefs_2)) % Polynomial([1] + [0 for _ in range(n-1)] + [1])
+            self.assertEqual(len(result.coefs), len(expected.coef))
+            for i in range(len(result.coefs)):
+                self.assertEqual(result.coefs[i], expected.coef[i] % q)

From 76020e51f6965e0c77a07d1c516e94421806699e Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Fri, 24 Nov 2023 17:52:55 +0200
Subject: [PATCH 10/15] added possibility to customize ccakem shared secret
 length

---
 kyber/ccakem.py      | 10 +++++-----
 tests/test_ccakem.py |  8 ++++++++
 2 files changed, 13 insertions(+), 5 deletions(-)

diff --git a/kyber/ccakem.py b/kyber/ccakem.py
index 2752f01..1e7b84c 100644
--- a/kyber/ccakem.py
+++ b/kyber/ccakem.py
@@ -21,7 +21,7 @@ def ccakem_generate_keys() -> tuple[bytes, bytes]:
         pk      # public key
     )
 
-def ccakem_encrypt(public_key: bytes) -> tuple[bytes, bytes]:
+def ccakem_encrypt(public_key: bytes, shared_secret_length: int = 32) -> tuple[bytes, bytes]:
     """
     Takes public key as input and returns (ciphertext, shered_secret) as a tuple.
     Shared secret is 32 bytes in length.
@@ -33,14 +33,14 @@ def ccakem_encrypt(public_key: bytes) -> tuple[bytes, bytes]:
     Kr = G(m + H(public_key))
     K, r = Kr[:32], Kr[32:]
     c = Encrypt(public_key, m, r).encrypt()
-    K = kdf(K + H(c), 32)
+    K = kdf(K + H(c), shared_secret_length)
 
     return (
         c,      # ciphertext
         K       # shared secret
     )
 
-def ccakem_decrypt(ciphertext: bytes, private_key: bytes) -> bytes:
+def ccakem_decrypt(ciphertext: bytes, private_key: bytes, shared_secret_length: int = 32) -> bytes:
     """
     Decrypts the given ciphertext with the private key.
     :returns Decrypted 32-byte shared secret.
@@ -62,5 +62,5 @@ def ccakem_decrypt(ciphertext: bytes, private_key: bytes) -> bytes:
     c = Encrypt(pk, m, r).encrypt()
 
     if c == ciphertext:
-        return kdf(K + H(c), 32)
-    return kdf(z + H(c), 32)
+        return kdf(K + H(c), shared_secret_length)
+    return kdf(z + H(c), shared_secret_length)
diff --git a/tests/test_ccakem.py b/tests/test_ccakem.py
index e8c0a5d..eaf2a3c 100644
--- a/tests/test_ccakem.py
+++ b/tests/test_ccakem.py
@@ -39,3 +39,11 @@ def test_ccakem_integration(self):
         ciphertext, shared_secret1 = ccakem_encrypt(public_key)
         shared_secret2 = ccakem_decrypt(ciphertext, private_key)
         self.assertEqual(shared_secret1, shared_secret2)
+
+    def test_ccakem_symmetry_with_long_shared_secret(self):
+        # generate 1000-byte shared secret
+        private_key, public_key = ccakem_generate_keys()
+        ciphertext, shared_secret1 = ccakem_encrypt(public_key, 1000)
+        shared_secret2 = ccakem_decrypt(ciphertext, private_key, 1000)
+        self.assertEqual(len(shared_secret1), 1000)
+        self.assertEqual(shared_secret1, shared_secret2)

From 173a3440c32ccd7b9a1746da3c1df0ddfd9574cc Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Fri, 24 Nov 2023 18:02:59 +0200
Subject: [PATCH 11/15] added performance test

---
 perf_tests/__main__.py        |  3 +++
 perf_tests/test_encryption.py | 41 +++++++++++++++++++++++++++++++++++
 tasks.py                      |  5 +++++
 3 files changed, 49 insertions(+)
 create mode 100644 perf_tests/__main__.py
 create mode 100644 perf_tests/test_encryption.py

diff --git a/perf_tests/__main__.py b/perf_tests/__main__.py
new file mode 100644
index 0000000..a6b752e
--- /dev/null
+++ b/perf_tests/__main__.py
@@ -0,0 +1,3 @@
+from perf_tests.test_encryption import runner as encryption_test_runner
+
+encryption_test_runner()
diff --git a/perf_tests/test_encryption.py b/perf_tests/test_encryption.py
new file mode 100644
index 0000000..0f7e48b
--- /dev/null
+++ b/perf_tests/test_encryption.py
@@ -0,0 +1,41 @@
+from random import randbytes
+from time import time
+from kyber.encryption import Encrypt, Decrypt, generate_keys
+from kyber.constants import k, n, du, dv
+
+def run(payload: bytes) -> tuple[float, float, float]:
+    t0 = time()
+
+    private_key, public_key = generate_keys()
+
+    t1 = time()
+
+    ciphertext = bytearray()
+    for i in range(0, len(payload), 32):
+        ciphertext += Encrypt(public_key, payload[i:i+32]).encrypt()
+
+    t2 = time()
+
+    ciphertext_chunk_size = du*k*n//8 + dv*n//8
+    restored_payload = bytearray()
+    for i in range(0, len(ciphertext), ciphertext_chunk_size):
+        restored_payload += Decrypt(private_key, ciphertext[i:i+ciphertext_chunk_size]).decrypt()
+    
+    t3 = time()
+
+    assert payload == restored_payload
+
+    return (t1-t0, t2-t1, t3-t2)
+
+def runner():
+    payload = randbytes(10_048)   # about 10 kiB, multipla of 32
+    print("Starting encryption performance test (about 1 min)")
+    durations = run(payload)
+    print("Results:")
+    print(f"Keypair generation: {durations[0]:.2f} sec")
+    print(f"Encryption: {durations[1]:.2f} sec")
+    print(f"Decryption: {durations[2]:.2f} sec")
+    print(f"Total: {sum(durations):.2f} sec")
+
+if __name__ == "__main__":
+    runner()
diff --git a/tasks.py b/tasks.py
index acc40fc..29e232e 100644
--- a/tasks.py
+++ b/tasks.py
@@ -1,9 +1,14 @@
+import os
 from invoke import task
 
 @task
 def test(ctx):
     ctx.run("pytest tests/", pty=True)
 
+@task
+def performance(ctx):
+    ctx.run("python perf_tests/", pty=True)
+
 @task
 def coverage(ctx):
     ctx.run("coverage run -m pytest --ignore tests/test_integration.py tests/", pty=True)

From d802c9d4d7d8aece6e072d3ba2ec50b86bc29959 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Fri, 24 Nov 2023 18:20:56 +0200
Subject: [PATCH 12/15] added performance tests to docs

---
 docs/tests.md | 6 ++++++
 docs/usage.md | 8 ++++++++
 2 files changed, 14 insertions(+)

diff --git a/docs/tests.md b/docs/tests.md
index 4adb789..da178cc 100644
--- a/docs/tests.md
+++ b/docs/tests.md
@@ -53,3 +53,9 @@ TOTAL                              331      0     92      0   100%
 ```
 
 For more detailed report, run `poetry run invoke coverage-report` and then open `htmlcov/index.html`.
+
+## Performance tests
+
+The asymmetric encryption part of Kyber (CPAPKE in the specification document) only works with fixed-lengthed input, but we can split larger payload into 32-byte chunks and encrypt them separately. Ciphertexts can be concatenated and the whole process can be reversed during decryption. Using this method, encryption is tested with about 10 kibibytes of random payload.
+
+Performance tests can be run with `poetry run invoke performance`.
diff --git a/docs/usage.md b/docs/usage.md
index 052ef2e..a827f18 100644
--- a/docs/usage.md
+++ b/docs/usage.md
@@ -50,6 +50,14 @@ poetry run invoke coverage-report
 
 after which the report will appear at `htmlcov/index.html`.
 
+### Performance tests
+
+Run performance tests with
+
+```
+poetry run invoke performance
+```
+
 ### Lint
 
 Run static style cheking with

From 9db608a1228cc391751659ce2abcc93273957a28 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Sat, 25 Nov 2023 08:56:13 +0200
Subject: [PATCH 13/15] added week 4 report

---
 README.md      | 1 +
 docs/week-4.md | 7 +++++++
 2 files changed, 8 insertions(+)
 create mode 100644 docs/week-4.md

diff --git a/README.md b/README.md
index 2974723..ee8758a 100644
--- a/README.md
+++ b/README.md
@@ -18,3 +18,4 @@ Implementation of [CRYSTALS-Kyber](https://pq-crystals.org/kyber/index.shtml) en
 * [Week 1](docs/week-1.md)
 * [Week 2](docs/week-2.md)
 * [Week 3](docs/week-3.md)
+* [Week 4](docs/week-4.md)
diff --git a/docs/week-4.md b/docs/week-4.md
new file mode 100644
index 0000000..d0eda2b
--- /dev/null
+++ b/docs/week-4.md
@@ -0,0 +1,7 @@
+# Week 4
+
+_20. – 26.11.2023_
+
+This week I started implementing own data structures. First, I added `PolynomialRing` which, as the name tells, represents polynomial ring structure and handles its core calculations. My original plan was also to implement matrix structure, but it turned out to be much slower than Numpy array, and therefore will not be added. I also added more unit tests and started adding performance tests.
+
+Total working time: 5 hours

From 96c8513fe2e686e3cd2610a1e3c566399c4301a8 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Sat, 25 Nov 2023 09:13:17 +0200
Subject: [PATCH 14/15] improve docstrings

---
 kyber/entities/polring.py | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/kyber/entities/polring.py b/kyber/entities/polring.py
index 5c266e1..7d8bf01 100644
--- a/kyber/entities/polring.py
+++ b/kyber/entities/polring.py
@@ -2,7 +2,10 @@
 
 class PolynomialRing:
     def __init__(self, coefs: list[int], check_limits: bool = True) -> None:
-        """Input `(1, 2, 3)` represents `1+2x+3x^2`."""
+        """
+        :param coefs: Coefficients of the polynomial. E.g. `[1, 2, 3]` represents `1+2x+3x^2`.
+        :param check_limits: Set to `False` if coefs is already taken to modulo.
+        """
         self._coefs = [int(c) for c in coefs]
         self._coef_limit = q
         self._degree_limit = n-1
@@ -11,9 +14,11 @@ def __init__(self, coefs: list[int], check_limits: bool = True) -> None:
 
     @property
     def coefs(self) -> list[int]:
+        """Coefficients of the polynomial. E.g. `[1, 2, 3]` represents `1+2x+3x^2`."""
         return self._coefs
 
     def _apply_limits(self) -> None:
+        """Take this polynomial to modulo `x^n+1` and coefs to modulo `q`."""
         # apply degree limit by dividing self by x^n+1
         self._apply_polynomial_modulo_limit()
 

From d659b3d3f5d7087497c5dfcf952ba23128fb28c8 Mon Sep 17 00:00:00 2001
From: Pyry Lahtinen <pyry@pyry.info>
Date: Sat, 25 Nov 2023 09:46:58 +0200
Subject: [PATCH 15/15] improved polring sum performance

---
 kyber/entities/polring.py | 25 ++++++++++++++-----------
 1 file changed, 14 insertions(+), 11 deletions(-)

diff --git a/kyber/entities/polring.py b/kyber/entities/polring.py
index 7d8bf01..a1c69f2 100644
--- a/kyber/entities/polring.py
+++ b/kyber/entities/polring.py
@@ -7,8 +7,6 @@ def __init__(self, coefs: list[int], check_limits: bool = True) -> None:
         :param check_limits: Set to `False` if coefs is already taken to modulo.
         """
         self._coefs = [int(c) for c in coefs]
-        self._coef_limit = q
-        self._degree_limit = n-1
         if check_limits:
             self._apply_limits()
 
@@ -23,8 +21,7 @@ def _apply_limits(self) -> None:
         self._apply_polynomial_modulo_limit()
 
         # apply coef limit
-        for i in range(len(self._coefs)):
-            self._coefs[i] %= self._coef_limit
+        self._coefs = [c % q for c in self._coefs]
 
         # remove trailing zero coefficients
         while len(self._coefs) > 0 and self._coefs[-1] == 0:
@@ -33,25 +30,31 @@ def _apply_limits(self) -> None:
     def _apply_polynomial_modulo_limit(self) -> None:
         """Replaces `self._coefs` with the remainder of division `self._coefs / (x^n+1)`."""
         # this is an optimal version of polynomial long division
-        while len(self._coefs) >= n+1:
+        coef_count = len(self._coefs)
+        while coef_count >= n+1:
             self._coefs[-n-1] -= self._coefs[-1]
             self._coefs[-1] = 0
             while self._coefs[-1] == 0:
                 self._coefs.pop()
+                coef_count -= 1
 
     def __add__(self, other: "PolynomialRing") -> "PolynomialRing":
         result = []
-        for i in range(max(len(self.coefs), len(other.coefs))):
-            self_coef = self.coefs[i] if i < len(self.coefs) else 0
-            other_coef = other.coefs[i] if i < len(other.coefs) else 0
+        self_length = len(self._coefs)
+        other_length = len(other.coefs)
+        for i in range(max(self_length, other_length)):
+            self_coef = self.coefs[i] if i < self_length else 0
+            other_coef = other.coefs[i] if i < other_length else 0
             result.append(self_coef + other_coef)
         return PolynomialRing(result)
 
     def __sub__(self, other: "PolynomialRing") -> "PolynomialRing":
         result = []
-        for i in range(max(len(self.coefs), len(other.coefs))):
-            self_coef = self.coefs[i] if i < len(self.coefs) else 0
-            other_coef = other.coefs[i] if i < len(other.coefs) else 0
+        self_length = len(self._coefs)
+        other_length = len(other.coefs)
+        for i in range(max(self_length, other_length)):
+            self_coef = self.coefs[i] if i < self_length else 0
+            other_coef = other.coefs[i] if i < other_length else 0
             result.append(self_coef - other_coef)
         return PolynomialRing(result)