Implement encode rows and test

da/common.py

@@ -11,11 +11,11 @@ class Chunk(Bytes32):
     pass
 
 
-class Column(List[Chunk]):
+class Column(bytes):
     pass
 
 
-class Row(List[Chunk]):
+class Row(bytes):
     pass
 
 

da/encoder.py

@@ -1,11 +1,12 @@
 from dataclasses import dataclass
-from itertools import batched
+from itertools import batched, chain
 from typing import List, Sequence
-from eth2spec.eip7594.mainnet import KZGCommitment as Commitment, KZGProof as Proof
+from eth2spec.eip7594.mainnet import KZGCommitment as Commitment, KZGProof as Proof, BLSFieldElement
 
-from da.common import ChunksMatrix
+from da.common import ChunksMatrix, Chunk
 from da.kzg_rs import kzg, rs, poly
-from da.kzg_rs.common import GLOBAL_PARAMETERS
+from da.kzg_rs.common import GLOBAL_PARAMETERS, ROOTS_OF_UNITY
+
 
 @dataclass
 class DAEncoderParams:
@@ -15,7 +16,7 @@ class DAEncoderParams:
 
 @dataclass
 class EncodedData:
-    data: bytearray
+    data: bytes
     extended_matrix: ChunksMatrix
     row_commitments: List[Commitment]
     row_proofs: List[List[Proof]]
@@ -28,16 +29,26 @@ class DAEncoder:
     def __init__(self, params: DAEncoderParams):
         self.params = params
 
-    def _chunkify_data(self, data: bytearray) -> ChunksMatrix:
+    def _chunkify_data(self, data: bytes) -> ChunksMatrix:
         size: int = self.params.column_count * self.params.bytes_per_field_element
-        return ChunksMatrix(batched(data, size))
+        return ChunksMatrix(bytes(b) for b in batched(data, size))
 
     @staticmethod
-    def _compute_row_kzg_commitments(rows: Sequence[bytearray]) -> List[Commitment]:
+    def _compute_row_kzg_commitments(rows: Sequence[bytes]) -> List[Commitment]:
         return [kzg.bytes_to_commitment(row, GLOBAL_PARAMETERS) for row in rows]
 
     def _rs_encode_rows(self, chunks_matrix: ChunksMatrix) -> ChunksMatrix:
-        ...
+        def __rs_encode_row(row: bytes) -> bytes:
+            polynomial = kzg.bytes_to_polynomial(row)
+            return bytes(
+                chain.from_iterable(
+                    Chunk(BLSFieldElement.to_bytes(
+                        x,
+                        length=self.params.bytes_per_field_element, byteorder="big"
+                    )) for x in rs.encode(polynomial, 2, ROOTS_OF_UNITY)
+                )
+            )
+        return ChunksMatrix(__rs_encode_row(row) for row in chunks_matrix)
 
     def _compute_rows_proofs(self, chunks_matrix: ChunksMatrix, row_commitments: List[Commitment]) -> List[List[Proof]]:
         ...
@@ -57,7 +68,7 @@ class DAEncoder:
     ) -> List[Proof]:
         ...
 
-    def encode(self, data: bytearray) -> EncodedData:
+    def encode(self, data: bytes) -> EncodedData:
         chunks_matrix = self._chunkify_data(data)
         row_commitments = self._compute_row_kzg_commitments(chunks_matrix)
         extended_matrix = self._rs_encode_rows(chunks_matrix)

da/kzg_rs/kzg.py

@@ -9,12 +9,12 @@ from .common import BYTES_PER_FIELD_ELEMENT, G1, BLS_MODULUS, GLOBAL_PARAMETERS_
 from .poly import Polynomial
 
 
-def bytes_to_polynomial(bytes: bytearray) -> Polynomial:
+def bytes_to_polynomial(b: bytes) -> Polynomial:
     """
     Convert bytes to list of BLS field scalars.
     """
-    assert len(bytes) % BYTES_PER_FIELD_ELEMENT == 0
-    eval_form = [int(bytes_to_bls_field(b)) for b in batched(bytes, int(BYTES_PER_FIELD_ELEMENT))]
+    assert len(b) % BYTES_PER_FIELD_ELEMENT == 0
+    eval_form = [int(bytes_to_bls_field(b)) for b in batched(b, int(BYTES_PER_FIELD_ELEMENT))]
     return Polynomial.from_evaluations(eval_form, BLS_MODULUS)
 
 
@@ -33,7 +33,7 @@ def g1_linear_combination(polynomial: Polynomial[BLSFieldElement], global_parame
     return Commitment(bls.G1_to_bytes48(point))
 
 
-def bytes_to_commitment(b: bytearray, global_parameters: Sequence[G1]) -> Commitment:
+def bytes_to_commitment(b: bytes, global_parameters: Sequence[G1]) -> Commitment:
     poly = bytes_to_polynomial(b)
     return g1_linear_combination(poly, global_parameters)
 

da/kzg_rs/poly.py

@@ -11,7 +11,7 @@ class Polynomial[T]:
 
     @classmethod
     def from_evaluations(cls, evaluations: Sequence[T], modulus) -> Self:
-        coefficients = intt(evaluations, prime=modulus)[:len(evaluations)]
+        coefficients = intt(evaluations, prime=modulus)
         return cls(coefficients, modulus)
 
     def __repr__(self):

da/kzg_rs/rs.py

@@ -9,7 +9,7 @@ from .poly import Polynomial
 ExtendedData = Sequence[BLSFieldElement]
 
 
-def encode(polynomial: Polynomial, factor: int, roots_of_unity: Sequence[BLSFieldElement]) -> ExtendedData:
+def encode(polynomial: Polynomial, factor: int, roots_of_unity: Sequence[int]) -> ExtendedData:
     """
     Encode a polynomial extending to the given factor
     Parameters:

da/test_encoder.py

@@ -1,21 +1,20 @@
-from itertools import chain
+from itertools import chain, batched
 from random import randrange
-from typing import List
 from unittest import TestCase
 
 from da import encoder
 from da.encoder import DAEncoderParams, Commitment, DAEncoder
-from eth2spec.eip7594.mainnet import BYTES_PER_FIELD_ELEMENT
-
-from da.kzg_rs.common import BLS_MODULUS
+from eth2spec.eip7594.mainnet import BYTES_PER_FIELD_ELEMENT, BLSFieldElement
 
+from da.kzg_rs.common import BLS_MODULUS, ROOTS_OF_UNITY
+from da.kzg_rs import kzg, rs
 
 class TestEncoder(TestCase):
 
     def setUp(self):
-        self.params: DAEncoderParams = DAEncoderParams(column_count=10, bytes_per_field_element=32)
+        self.params: DAEncoderParams = DAEncoderParams(column_count=16, bytes_per_field_element=32)
         self.encoder: DAEncoder = DAEncoder(self.params)
-        self.elements = 100
+        self.elements = 16
         self.data = bytearray(
             chain.from_iterable(
                 randrange(BLS_MODULUS).to_bytes(length=self.params.bytes_per_field_element, byteorder='big')
@@ -47,7 +46,14 @@ class TestEncoder(TestCase):
         self.assertEqual(len(commitments), len(chunks_matrix))
 
     def test_rs_encode_rows(self):
-        pass
+        chunks_matrix = self.encoder._chunkify_data(self.data)
+        extended_chunks_matrix = self.encoder._rs_encode_rows(chunks_matrix)
+        for r1, r2 in zip(chunks_matrix, extended_chunks_matrix):
+            r2 = [BLSFieldElement.from_bytes(x) for x in batched(r2, self.params.bytes_per_field_element)]
+            poly_1 = kzg.bytes_to_polynomial(r1)
+            # we check against decoding so we now the encoding was properly done
+            poly_2 = rs.decode(r2, ROOTS_OF_UNITY, len(poly_1))
+            self.assertEqual(poly_1, poly_2)
 
     def test_compute_rows_proofs(self):
         pass