Fix fk20 and tests

da/kzg_rs/fk20.py

@@ -1,9 +1,10 @@
 from typing import List, Sequence
 
-from eth2spec.deneb.mainnet import KZGProof as Proof
+from eth2spec.deneb.mainnet import KZGProof as Proof, BLSFieldElement
+from eth2spec.utils import bls
 
 from da.kzg_rs.common import G1, BLS_MODULUS
-from da.kzg_rs.fft import fft, ifft
+from da.kzg_rs.fft import fft, ifft, fft_g1, ifft_g1
 from da.kzg_rs.poly import Polynomial
 from da.kzg_rs.utils import is_power_of_two
 
@@ -22,19 +23,19 @@ def toeplitz1(global_parameters: List[G1], roots_of_unity: Sequence[int], polyno
     # algorithm only works on powers of 2 for dft computations
     assert is_power_of_two(len(global_parameters))
     roots_of_unity = roots_of_unity[:2*polynomial_degree]
-    vector_x_extended = global_parameters + [G1(0) for _ in range(len(global_parameters))]
-    vector_x_extended_fft = fft(vector_x_extended, roots_of_unity, BLS_MODULUS)
+    vector_x_extended = global_parameters + [bls.multiply(bls.Z1(), 0) for _ in range(len(global_parameters))]
+    vector_x_extended_fft = fft_g1(vector_x_extended, roots_of_unity, BLS_MODULUS)
     return vector_x_extended_fft
 
 
 def toeplitz2(coefficients: List[G1], roots_of_unity: Sequence[int], extended_vector: Sequence[G1]) -> List[G1]:
     assert is_power_of_two(len(coefficients))
     toeplitz_coefficients_fft = fft(coefficients, roots_of_unity, BLS_MODULUS)
-    return [v*c for v, c in zip(extended_vector, toeplitz_coefficients_fft)]
+    return [bls.multiply(v, c) for v, c in zip(extended_vector, toeplitz_coefficients_fft)]
 
 
 def toeplitz3(h_extended_fft: Sequence[G1], roots_of_unity: Sequence[int], polynomial_degree: int) -> List[G1]:
-    return ifft(h_extended_fft, roots_of_unity, BLS_MODULUS)[:polynomial_degree]
+    return ifft_g1(h_extended_fft, roots_of_unity, BLS_MODULUS)[:polynomial_degree]
 
 
 def fk20_generate_proofs(
@@ -49,7 +50,7 @@ def fk20_generate_proofs(
     # 1.1 y = dft([s^d-1, s^d-2, ..., s, 1, *[0 for _ in len(polynomial)]])
     # 1.2 z = dft([*[0 for _ in len(polynomial)], f1, f2, ..., fd])
     # 1.3 u = y * v * roots_of_unity(len(polynomial)*2)
-    global_parameters = [*global_parameters[polynomial_degree-2::-1], 0]
+    global_parameters = [*global_parameters[polynomial_degree-2::-1], bls.multiply(bls.Z1(), 0)]
     extended_vector = toeplitz1(global_parameters, roots_of_unity[:polynomial_degree*2], polynomial_degree)
     # 2 - Build circulant matrix with the polynomial coefficients (reversed N..n, and padded)
     toeplitz_coefficients = [
@@ -59,5 +60,6 @@ def fk20_generate_proofs(
     # 3 - Perform fft and nub the tail half as it is padding
     h_vector = toeplitz3(h_extended_vector, roots_of_unity[:len(h_extended_vector)], polynomial_degree)
     # 4 - proof are the dft of the h vector
-    proofs = fft(h_vector, roots_of_unity[:polynomial_degree], BLS_MODULUS)
+    proofs = fft_g1(h_vector, roots_of_unity[:polynomial_degree], BLS_MODULUS)
+    proofs = [Proof(bls.G1_to_bytes48(proof)) for proof in proofs]
     return proofs

da/kzg_rs/test_fft.py

@@ -1,14 +1,14 @@
 from unittest import TestCase
 
+from da.kzg_rs.roots import compute_roots_of_unity
 from da.kzg_rs.common import BLS_MODULUS
 from fft import fft, ifft
 
 
 class TestFFT(TestCase):
     def test_fft_ifft(self):
-        for size in [256, 512, 1024, 2048, 4096]:
-            roots_of_unity = [pow(23674694431658770659612952115660802947967373701506253797663184111817857449850, i, BLS_MODULUS) for i in range(size)]
+        for size in [16, 32, 64, 128, 256, 512, 1024, 2048, 4096]:
+            roots_of_unity = compute_roots_of_unity(2, size, BLS_MODULUS)
             vals = list(x for x in range(size))
             vals_fft = fft(vals, roots_of_unity, BLS_MODULUS)
             self.assertEqual(vals, ifft(vals_fft, roots_of_unity, BLS_MODULUS))
-

da/kzg_rs/test_fk20.py

@@ -1,5 +1,28 @@
+from itertools import chain
 from unittest import TestCase
 import random
+from .fk20 import fk20_generate_proofs
+from .kzg import generate_element_proof, bytes_to_polynomial
+from .common import BLS_MODULUS, BYTES_PER_FIELD_ELEMENT, GLOBAL_PARAMETERS
+from .roots import compute_roots_of_unity
+
 
 class TestFK20(TestCase):
-    def test_toeplizt1(self):
+    @staticmethod
+    def rand_bytes(n_chunks=1024):
+        return bytes(
+            chain.from_iterable(
+                int.to_bytes(random.randrange(BLS_MODULUS), length=BYTES_PER_FIELD_ELEMENT)
+                for _ in range(n_chunks)
+            )
+        )
+
+    def test_fk20(self):
+        for size in [16, 32, 64, 128]:
+            roots_of_unity = compute_roots_of_unity(2, size*2, BLS_MODULUS)
+            rand_bytes = self.rand_bytes(size)
+            polynomial = bytes_to_polynomial(rand_bytes)
+            proofs = [generate_element_proof(i, polynomial, GLOBAL_PARAMETERS, roots_of_unity) for i in range(size)]
+            fk20_proofs = fk20_generate_proofs(polynomial, GLOBAL_PARAMETERS, roots_of_unity)
+            self.assertEqual(proofs, fk20_proofs)
+