Fix roots computations

da/kzg_rs/common.py

@@ -12,10 +12,11 @@ G2 = G2Uncompressed
 
 BYTES_PER_FIELD_ELEMENT = 32
 BLS_MODULUS = eth2spec.eip7594.mainnet.BLS_MODULUS
+PRIMITIVE_ROOT: int = 7
 GLOBAL_PARAMETERS: List[G1]
 GLOBAL_PARAMETERS_G2: List[G2]
 # secret is fixed but this should come from a different synchronization protocol
 GLOBAL_PARAMETERS, GLOBAL_PARAMETERS_G2 = map(list, generate_setup(4096, 8, 1987))
 ROOTS_OF_UNITY: Tuple[int] = compute_roots_of_unity(
-    7, 4096, BLS_MODULUS
+    PRIMITIVE_ROOT, 4096, BLS_MODULUS
 )

da/kzg_rs/fk20.py

@@ -3,13 +3,14 @@ from typing import List, Sequence
 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.common import G1, BLS_MODULUS, PRIMITIVE_ROOT
 from da.kzg_rs.fft import fft, ifft, fft_g1, ifft_g1
 from da.kzg_rs.poly import Polynomial
+from da.kzg_rs.roots import compute_roots_of_unity
 from da.kzg_rs.utils import is_power_of_two
 
 
-def toeplitz1(global_parameters: List[G1], roots_of_unity: Sequence[int], polynomial_degree: int) -> List[G1]:
+def toeplitz1(global_parameters: List[G1], polynomial_degree: int) -> List[G1]:
     """
     This part can be precomputed for different global_parameters lengths depending on polynomial degree of powers of two.
     :param global_parameters:
@@ -17,8 +18,8 @@ def toeplitz1(global_parameters: List[G1], roots_of_unity: Sequence[int], polyno
     :param polynomial_degree:
     :return:
     """
-    assert len(roots_of_unity) >= 2 * polynomial_degree
     assert len(global_parameters) >= polynomial_degree
+    roots_of_unity = compute_roots_of_unity(PRIMITIVE_ROOT, polynomial_degree*2, BLS_MODULUS)
     global_parameters = global_parameters[:polynomial_degree]
     # algorithm only works on powers of 2 for dft computations
     assert is_power_of_two(len(global_parameters))
@@ -28,21 +29,22 @@ def toeplitz1(global_parameters: List[G1], roots_of_unity: Sequence[int], polyno
     return vector_x_extended_fft
 
 
-def toeplitz2(coefficients: List[G1], roots_of_unity: Sequence[int], extended_vector: Sequence[G1]) -> List[G1]:
+def toeplitz2(coefficients: List[G1], extended_vector: Sequence[G1]) -> List[G1]:
     assert is_power_of_two(len(coefficients))
+    roots_of_unity = compute_roots_of_unity(PRIMITIVE_ROOT, len(coefficients), BLS_MODULUS)
     toeplitz_coefficients_fft = fft(coefficients, roots_of_unity, BLS_MODULUS)
     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]:
+def toeplitz3(h_extended_fft: Sequence[G1], polynomial_degree: int) -> List[G1]:
+    roots_of_unity = compute_roots_of_unity(PRIMITIVE_ROOT, len(h_extended_fft), BLS_MODULUS)
     return ifft_g1(h_extended_fft, roots_of_unity, BLS_MODULUS)[:polynomial_degree]
 
 
 def fk20_generate_proofs(
-        polynomial: Polynomial, global_parameters: List[G1], roots_of_unity: Sequence[int]
+        polynomial: Polynomial, global_parameters: List[G1]
 ) -> List[Proof]:
     polynomial_degree = len(polynomial)
-    assert len(roots_of_unity) >= 2 * polynomial_degree
     assert len(global_parameters) >= polynomial_degree
     assert is_power_of_two(len(polynomial))
 
@@ -50,16 +52,19 @@ 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)
+    roots_of_unity = compute_roots_of_unity(PRIMITIVE_ROOT, polynomial_degree, BLS_MODULUS)
     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)
+    extended_vector = toeplitz1(global_parameters, polynomial_degree)
     # 2 - Build circulant matrix with the polynomial coefficients (reversed N..n, and padded)
     toeplitz_coefficients = [
-        polynomial.coefficients[-1], *(0 for _ in range(polynomial_degree+1)), *polynomial.coefficients[1:-1]
+        polynomial.coefficients[-1],
+        *(BLSFieldElement(0) for _ in range(polynomial_degree+1)),
+        *polynomial.coefficients[1:-1]
     ]
-    h_extended_vector = toeplitz2(toeplitz_coefficients, roots_of_unity[:len(extended_vector)], extended_vector)
+    h_extended_vector = toeplitz2(toeplitz_coefficients, extended_vector)
     # 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)
+    h_vector = toeplitz3(h_extended_vector, polynomial_degree)
     # 4 - proof are the dft of the h vector
-    proofs = fft_g1(h_vector, roots_of_unity[:polynomial_degree], BLS_MODULUS)
+    proofs = fft_g1(h_vector, roots_of_unity, BLS_MODULUS)
     proofs = [Proof(bls.G1_to_bytes48(proof)) for proof in proofs]
     return proofs

da/kzg_rs/test_fk20.py

@@ -3,7 +3,7 @@ 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 .common import BLS_MODULUS, BYTES_PER_FIELD_ELEMENT, GLOBAL_PARAMETERS, PRIMITIVE_ROOT
 from .roots import compute_roots_of_unity
 
 
@@ -19,11 +19,11 @@ class TestFK20(TestCase):
 
     def test_fk20(self):
         for size in [16, 32, 64, 128]:
-            roots_of_unity = compute_roots_of_unity(2, size*2, BLS_MODULUS)
+            roots_of_unity = compute_roots_of_unity(PRIMITIVE_ROOT, 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)
+            fk20_proofs = fk20_generate_proofs(polynomial, GLOBAL_PARAMETERS)
             self.assertEqual(len(proofs), len(fk20_proofs))
             self.assertEqual(proofs, fk20_proofs)