From 3620cdc8aef64a9056f6486630a151f6398c793f Mon Sep 17 00:00:00 2001
From: danielsanchezq <sanchez.quiros.daniel@gmail.com>
Date: Tue, 11 Jun 2024 18:08:02 +0200
Subject: [PATCH] Fix roots of unity generation

---
 da/kzg_rs/common.py |  8 +++++---
 da/kzg_rs/roots.py  | 33 ++++++++++++++++++++++-----------
 2 files changed, 27 insertions(+), 14 deletions(-)

diff --git a/da/kzg_rs/common.py b/da/kzg_rs/common.py
index 9eee3c2..42f756c 100644
--- a/da/kzg_rs/common.py
+++ b/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
@@ -15,5 +15,7 @@ BLS_MODULUS = eth2spec.eip7594.mainnet.BLS_MODULUS
 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(
+    7, 4096, BLS_MODULUS
+)
diff --git a/da/kzg_rs/roots.py b/da/kzg_rs/roots.py
index ec5988b..2c2d7a7 100644
--- a/da/kzg_rs/roots.py
+++ b/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)