diff --git a/src/kyber_py/modules/modules.py b/src/kyber_py/modules/modules.py index e33f5d2..6d19b2e 100644 --- a/src/kyber_py/modules/modules.py +++ b/src/kyber_py/modules/modules.py @@ -8,6 +8,14 @@ def __init__(self): self.matrix = MatrixKyber def decode_vector(self, input_bytes, k, d, is_ntt=False): + """ + Decode bytes into a a vector of polynomial elements. + + Each element is assumed to be encoded as a polynomial with ``d``-bit + coefficients (hence a polynomial is encoded into ``256 * d`` bits). + + A vector of length ``k`` then has ``256 * d * k`` bits. + """ # Ensure the input bytes are the correct length to create k elements with # d bits used for each coefficient if self.ring.n * d * k != len(input_bytes) * 8: @@ -32,6 +40,9 @@ def __init__(self, parent, matrix_data, transpose=False): super().__init__(parent, matrix_data, transpose=transpose) def encode(self, d): + """ + Encode every element of a matrix into bytes and concatenate + """ output = b"" for row in self._data: for ele in row: @@ -39,21 +50,34 @@ def encode(self, d): return output def compress(self, d): + """ + Compress every element of the matrix to have at most ``d`` bits + """ for row in self._data: for ele in row: ele.compress(d) return self def decompress(self, d): + """ + Perform (lossy) decompression of the polynomial assuming it has been + compressed to have at most ``d`` bits. + """ for row in self._data: for ele in row: ele.decompress(d) return self def to_ntt(self): + """ + Convert every element of the matrix into NTT form + """ data = [[x.to_ntt() for x in row] for row in self._data] return self.parent(data, transpose=self._transpose) def from_ntt(self): + """ + Convert every element of the matrix from NTT form + """ data = [[x.from_ntt() for x in row] for row in self._data] return self.parent(data, transpose=self._transpose) diff --git a/src/kyber_py/modules/modules_generic.py b/src/kyber_py/modules/modules_generic.py index 37aecc7..1353db6 100644 --- a/src/kyber_py/modules/modules_generic.py +++ b/src/kyber_py/modules/modules_generic.py @@ -1,9 +1,19 @@ class Module: def __init__(self, ring): + """ + Initialise a module over the ring ``ring``. + """ self.ring = ring self.matrix = Matrix def random_element(self, m, n): + """ + Generate a random element of the module of dimension m x n + + :param int m: the number of rows in the matrix + :param int m: the number of columns in tge matrix + :return: an element of the module with dimension `m times n` + """ elements = [ [self.ring.random_element() for _ in range(n)] for _ in range(m) ] @@ -47,7 +57,10 @@ def __call__(self, matrix_elements, transpose=False): def vector(self, elements): """ - Construct a vector with the given elements + Construct a vector given a list of elements of the module's ring + + :param list: a list of elements of the ring + :return: a vector of the module """ return self.matrix(self, [elements], transpose=True) @@ -64,6 +77,9 @@ def dim(self): """ Return the dimensions of the matrix with m rows and n columns + + :return: the dimension of the matrix ``(m, n)`` + :rtype: tuple(int, int) """ if not self._transpose: return len(self._data), len(self._data[0]) @@ -78,13 +94,13 @@ def _check_dimensions(self): def transpose(self): """ - Swap rows and columns of self + Return a matrix with the rows and columns of swapped """ return self.parent(self._data, not self._transpose) def transpose_self(self): """ - Transpose in place + Swap the rows and columns of the matrix in place """ self._transpose = not self._transpose return @@ -193,7 +209,7 @@ def __matmul__(self, other): def dot(self, other): """ - Inner product + Compute the inner product of two vectors """ if not isinstance(other, type(self)): raise TypeError("Can only perform dot product with other matrices") diff --git a/src/kyber_py/polynomials/polynomials.py b/src/kyber_py/polynomials/polynomials.py index 6876078..be6cb39 100644 --- a/src/kyber_py/polynomials/polynomials.py +++ b/src/kyber_py/polynomials/polynomials.py @@ -157,6 +157,7 @@ def _decompress_ele(self, x, d): def compress(self, d): """ Compress the polynomial by compressing each coefficient + NOTE: This is lossy compression """ self.coeffs = [self._compress_ele(c, d) for c in self.coeffs] @@ -165,6 +166,7 @@ def compress(self, d): def decompress(self, d): """ Decompress the polynomial by decompressing each coefficient + NOTE: This as compression is lossy, we have x' = decompress(compress(x)), which x' != x, but is close in magnitude. @@ -198,6 +200,9 @@ def to_ntt(self): return self.parent(coeffs, is_ntt=True) def from_ntt(self): + """ + Not supported, raises a ``TypeError`` + """ raise TypeError(f"Polynomial not in the NTT domain: {type(self) = }") @@ -207,6 +212,9 @@ def __init__(self, parent, coefficients): self.coeffs = self._parse_coefficients(coefficients) def to_ntt(self): + """ + Not supported, raises a ``TypeError`` + """ raise TypeError( f"Polynomial is already in the NTT domain: {type(self) = }" ) @@ -249,6 +257,10 @@ def _ntt_base_multiplication(a0, a1, b0, b1, zeta): return r0, r1 def _ntt_coefficient_multiplication(self, f_coeffs, g_coeffs): + """ + Given the coefficients of two polynomials compute the coefficients of + their product + """ new_coeffs = [] zetas = self.parent.ntt_zetas for i in range(64): @@ -272,7 +284,6 @@ def _ntt_coefficient_multiplication(self, f_coeffs, g_coeffs): def _ntt_multiplication(self, other): """ Number Theoretic Transform multiplication. - Only implemented (currently) for n = 256 """ new_coeffs = self._ntt_coefficient_multiplication( self.coeffs, other.coeffs diff --git a/src/kyber_py/polynomials/polynomials_generic.py b/src/kyber_py/polynomials/polynomials_generic.py index 45640da..ff1be8f 100644 --- a/src/kyber_py/polynomials/polynomials_generic.py +++ b/src/kyber_py/polynomials/polynomials_generic.py @@ -14,9 +14,16 @@ def __init__(self, q, n): self.element = Polynomial def gen(self): + """ + Return the generator `x` of the polynomial ring + """ return self([0, 1]) def random_element(self): + """ + Compute a random element of the polynomial ring with coefficients in the + canonical range: ``[0, q-1]`` + """ coefficients = [random.randint(0, self.q - 1) for _ in range(self.n)] return self(coefficients)