From 5d3182eace8c633d60d18caee21907c966f4f7f4 Mon Sep 17 00:00:00 2001 From: Daniel Sanchez Quiros Date: Tue, 5 Mar 2024 12:16:28 +0100 Subject: [PATCH] Added columns property to chunks matrix Added test for columns --- da/common.py | 4 ++-- da/dispersal.py | 2 +- da/encoder.py | 2 +- da/test_common.py | 12 ++++++++++++ 4 files changed, 16 insertions(+), 4 deletions(-) create mode 100644 da/test_common.py diff --git a/da/common.py b/da/common.py index bd4b58f..8a61dcd 100644 --- a/da/common.py +++ b/da/common.py @@ -20,9 +20,9 @@ class Row(List[Chunk]): class ChunksMatrix(List[Row]): + @property def columns(self) -> Generator[List[Chunk], None, None]: - # TODO: yield columns - yield None + yield from map(list, zip(*self)) diff --git a/da/dispersal.py b/da/dispersal.py index 6419100..1c8271f 100644 --- a/da/dispersal.py +++ b/da/dispersal.py @@ -13,7 +13,7 @@ class Dispersal: def _prepare_data(self, encoded_data: EncodedData) -> Generator[DABlob, None, None]: assert len(encoded_data.row_commitments) == len(self.nodes) assert len(encoded_data.row_proofs) == len(self.nodes) - columns = encoded_data.extended_matrix.columns() + columns = encoded_data.extended_matrix.columns column_commitments = encoded_data.column_commitments row_commitments = encoded_data.row_commitments rows_proofs = encoded_data.row_proofs diff --git a/da/encoder.py b/da/encoder.py index 97380e9..431b561 100644 --- a/da/encoder.py +++ b/da/encoder.py @@ -3,7 +3,7 @@ from typing import List from eth2spec.eip7594.mainnet import KZGCommitment as Commitment, KZGProof as Proof from da.common import ChunksMatrix - +from da.kzg_rs import kzg, rs, poly @dataclass class DAEncoderParams: diff --git a/da/test_common.py b/da/test_common.py new file mode 100644 index 0000000..de9aac5 --- /dev/null +++ b/da/test_common.py @@ -0,0 +1,12 @@ +from unittest import TestCase + +from da.common import ChunksMatrix + + +class TestCommon(TestCase): + + def test_chunks_matrix_columns(self): + matrix = ChunksMatrix([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + expected = [[1, 4, 7], [2, 5, 8], [3, 6, 9]] + for c1, c2 in zip(expected, matrix.columns): + self.assertEqual(c1, c2)