Da: fk20 proof generation (#95)

* Kickstart fk20

* Implement i/fft from ethspecs

* Expand test to different sizes

* Implement toeplizt

* Finish implementing fk20

* Fix roots of unity generation

* Implement fft for g1 values

* Fix fk20 and tests

* Add len assertion in test

* Fix roots computations

* Fix test

* Fix imports

* Fmt

* Docs and format

da/kzg_rs/common.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Tuple
 
 import eth2spec.eip7594.mainnet
 from py_ecc.bls.typing import G1Uncompressed, G2Uncompressed
@@ -12,8 +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(1024, 8, 1987))
-ROOTS_OF_UNITY: List[int] = compute_roots_of_unity(2, BLS_MODULUS, 4096)
+GLOBAL_PARAMETERS, GLOBAL_PARAMETERS_G2 = map(list, generate_setup(4096, 8, 1987))
+ROOTS_OF_UNITY: Tuple[int] = compute_roots_of_unity(
+    PRIMITIVE_ROOT, 4096, BLS_MODULUS
+)

da/kzg_rs/fft.py

@@ -0,0 +1,65 @@
+from typing import Sequence, List
+
+from eth2spec.deneb.mainnet import BLSFieldElement
+from eth2spec.utils import bls
+
+from da.kzg_rs.common import G1
+
+
+def fft_g1(vals: Sequence[G1], roots_of_unity: Sequence[BLSFieldElement], modulus: int) -> List[G1]:
+    if len(vals) == 1:
+        return vals
+    L = fft_g1(vals[::2], roots_of_unity[::2], modulus)
+    R = fft_g1(vals[1::2], roots_of_unity[::2], modulus)
+    o = [bls.Z1() for _ in vals]
+    for i, (x, y) in enumerate(zip(L, R)):
+        y_times_root = bls.multiply(y, roots_of_unity[i])
+        o[i] = (x + y_times_root)
+        o[i + len(L)] = x + -y_times_root
+    return o
+
+
+def ifft_g1(vals: Sequence[G1], roots_of_unity: Sequence[BLSFieldElement], modulus: int) -> List[G1]:
+    assert len(vals) == len(roots_of_unity)
+    # modular inverse
+    invlen = pow(len(vals), modulus-2, modulus)
+    return [
+        bls.multiply(x, invlen)
+        for x in fft_g1(
+            vals, [roots_of_unity[0], *roots_of_unity[:0:-1]], modulus
+        )
+    ]
+
+
+def _fft(
+        vals: Sequence[BLSFieldElement],
+        roots_of_unity: Sequence[BLSFieldElement],
+        modulus: int,
+) -> Sequence[BLSFieldElement]:
+    if len(vals) == 1:
+        return vals
+    L = _fft(vals[::2], roots_of_unity[::2], modulus)
+    R = _fft(vals[1::2], roots_of_unity[::2], modulus)
+    o = [BLSFieldElement(0) for _ in vals]
+    for i, (x, y) in enumerate(zip(L, R)):
+        y_times_root = BLSFieldElement((int(y) * int(roots_of_unity[i])) % modulus)
+        o[i] = BLSFieldElement((int(x) + y_times_root) % modulus)
+        o[i + len(L)] = BLSFieldElement((int(x) - int(y_times_root) + modulus) % modulus)
+    return o
+
+
+def fft(vals, root_of_unity, modulus):
+    assert len(vals) == len(root_of_unity)
+    return _fft(vals, root_of_unity, modulus)
+
+
+def ifft(vals, roots_of_unity, modulus):
+    assert len(vals) == len(roots_of_unity)
+    # modular inverse
+    invlen = pow(len(vals), modulus-2, modulus)
+    return [
+        BLSFieldElement((int(x) * invlen) % modulus)
+        for x in _fft(
+            vals, [roots_of_unity[0], *roots_of_unity[:0:-1]], modulus
+        )
+    ]

da/kzg_rs/fk20.py

@@ -0,0 +1,78 @@
+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, PRIMITIVE_ROOT
+from da.kzg_rs.fft import fft, 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], 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:
+    :param roots_of_unity:
+    :param polynomial_degree:
+    :return:
+    """
+    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))
+    roots_of_unity = roots_of_unity[:2*polynomial_degree]
+    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], 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], 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]
+) -> List[Proof]:
+    """
+    Generate all proofs for the polynomial points in batch.
+    This method uses the fk20 algorthm from https://eprint.iacr.org/2023/033.pdf
+    Disclaimer: It only works for polynomial degree of powers of two.
+    :param polynomial: polynomial to generate proof for
+    :param global_parameters: setup generated parameters
+    :return: list of proof for each point in the polynomial
+    """
+    polynomial_degree = len(polynomial)
+    assert len(global_parameters) >= polynomial_degree
+    assert is_power_of_two(len(polynomial))
+
+    # 1 - Build toeplitz matrix for h values
+    # 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, polynomial_degree)
+    # 2 - Build circulant matrix with the polynomial coefficients (reversed N..n, and padded)
+    toeplitz_coefficients = [
+        polynomial.coefficients[-1],
+        *(BLSFieldElement(0) for _ in range(polynomial_degree+1)),
+        *polynomial.coefficients[1:-1]
+    ]
+    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, polynomial_degree)
+    # 4 - proof are the dft of the h vector
+    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/roots.py

@@ -1,14 +1,25 @@
-def compute_roots_of_unity(primitive_root, p, n):
-    """
-    Compute the roots of unity modulo p.
+from typing import Tuple
 
-    Parameters:
-        primitive_root (int): Primitive root modulo p.
-        p (int): Modulus.
-        n (int): Number of roots of unity to compute.
 
-    Returns:
-        list: List of roots of unity modulo p.
+def compute_root_of_unity(primitive_root: int, order: int, modulus: int) -> int:
     """
-    roots_of_unity = [pow(primitive_root, i, p) for i in range(n)]
-    return roots_of_unity
+    Generate a w such that ``w**length = 1``.
+    """
+    assert (modulus - 1) % order == 0
+    return pow(primitive_root, (modulus - 1) // order, modulus)
+
+
+def compute_roots_of_unity(primitive_root: int, order: int, modulus: int) -> Tuple[int]:
+    """
+    Compute a list of roots of unity for a given order.
+    The order must divide the BLS multiplicative group order, i.e. BLS_MODULUS - 1
+    """
+    assert (modulus - 1) % order == 0
+    root_of_unity = compute_root_of_unity(primitive_root, order, modulus)
+
+    roots = []
+    current_root_of_unity = 1
+    for _ in range(order):
+        roots.append(current_root_of_unity)
+        current_root_of_unity = current_root_of_unity * root_of_unity % modulus
+    return tuple(roots)

da/kzg_rs/test_fft.py

@@ -0,0 +1,14 @@
+from unittest import TestCase
+
+from .roots import compute_roots_of_unity
+from .common import BLS_MODULUS
+from .fft import fft, ifft
+
+
+class TestFFT(TestCase):
+    def test_fft_ifft(self):
+        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

@@ -0,0 +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, PRIMITIVE_ROOT
+from .roots import compute_roots_of_unity
+
+
+class TestFK20(TestCase):
+    @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, 256]:
+            roots_of_unity = compute_roots_of_unity(PRIMITIVE_ROOT, size, 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)
+            self.assertEqual(len(proofs), len(fk20_proofs))
+            self.assertEqual(proofs, fk20_proofs)

da/kzg_rs/utils.py

@@ -0,0 +1,5 @@
+POWERS_OF_2 = {2**i for i in range(1, 32)}
+
+
+def is_power_of_two(n) -> bool:
+    return n in POWERS_OF_2

requirements.txt

@@ -1,4 +1,4 @@
-blspy==2.0.2
+blspy==2.0.3
 cffi==1.16.0
 cryptography==41.0.7
 numpy==1.26.3