From 296ba921c9c06c7aa2b5daea4c34c4407601f524 Mon Sep 17 00:00:00 2001
From: George Kadianakis <desnacked@riseup.net>
Date: Mon, 16 Jan 2023 15:57:35 +0200
Subject: [PATCH] EIP4844: Handle barycentric evaluations at roots of unity

---
 specs/eip4844/polynomial-commitments.md       |  8 +++--
 .../test_polynomial_commitments.py            | 34 +++++++++++++++++++
 .../pyspec/eth2spec/test/helpers/sharding.py  | 32 +++++++++++++++++
 3 files changed, 71 insertions(+), 3 deletions(-)

diff --git a/specs/eip4844/polynomial-commitments.md b/specs/eip4844/polynomial-commitments.md
index 8884457ed..26ee00487 100644
--- a/specs/eip4844/polynomial-commitments.md
+++ b/specs/eip4844/polynomial-commitments.md
@@ -302,11 +302,13 @@ def evaluate_polynomial_in_evaluation_form(polynomial: Polynomial,
     assert width == FIELD_ELEMENTS_PER_BLOB
     inverse_width = bls_modular_inverse(BLSFieldElement(width))
 
-    # Make sure we won't divide by zero during division
-    assert z not in ROOTS_OF_UNITY
-
     roots_of_unity_brp = bit_reversal_permutation(ROOTS_OF_UNITY)
 
+    # If we are asked to evaluate within the domain, we already know the answer
+    if z in roots_of_unity_brp:
+        eval_index = roots_of_unity_brp.index(z)
+        return BLSFieldElement(polynomial[eval_index])
+
     result = 0
     for i in range(width):
         a = BLSFieldElement(int(polynomial[i]) * int(roots_of_unity_brp[i]) % BLS_MODULUS)
diff --git a/tests/core/pyspec/eth2spec/test/eip4844/unittests/polynomial_commitments/test_polynomial_commitments.py b/tests/core/pyspec/eth2spec/test/eip4844/unittests/polynomial_commitments/test_polynomial_commitments.py
index 3e9e2cb63..75a4e822a 100644
--- a/tests/core/pyspec/eth2spec/test/eip4844/unittests/polynomial_commitments/test_polynomial_commitments.py
+++ b/tests/core/pyspec/eth2spec/test/eip4844/unittests/polynomial_commitments/test_polynomial_commitments.py
@@ -4,6 +4,8 @@ from eth2spec.test.context import (
 )
 from eth2spec.test.helpers.sharding import (
     get_sample_blob,
+    get_poly_in_both_forms,
+    eval_poly_in_coeff_form,
 )
 
 
@@ -18,3 +20,35 @@ def test_verify_kzg_proof(spec, state):
 
     y = spec.evaluate_polynomial_in_evaluation_form(polynomial, x)
     assert spec.verify_kzg_proof_impl(commitment, x, y, proof)
+
+
+@with_eip4844_and_later
+@spec_state_test
+def test_barycentric_within_domain(spec, state):
+    """
+    Test barycentric formula correctness by using it to evaluate a polynomial at all the points of its domain
+    (the roots of unity).
+
+    Then make sure that we would get the same result if we evaluated it from coefficient form without using the
+    barycentric formula
+    """
+    poly_coeff, poly_eval = get_poly_in_both_forms(spec)
+    roots_of_unity_brp = spec.bit_reversal_permutation(spec.ROOTS_OF_UNITY)
+
+    assert len(poly_coeff) == len(poly_eval) == len(roots_of_unity_brp)
+    n = len(poly_coeff)
+
+    # Iterate over the entire domain
+    for i in range(n):
+        # Grab a root of unity and use it as the evaluation point
+        z = int(roots_of_unity_brp[i])
+
+        # Get p(z) by evaluating poly in coefficient form
+        p_z_coeff = eval_poly_in_coeff_form(spec, poly_coeff, z)
+
+        # Get p(z) by evaluating poly in evaluation form
+        p_z_eval = spec.evaluate_polynomial_in_evaluation_form(poly_eval, z)
+
+        # The two evaluations should be agree and p(z) should also be the i-th "coefficient" of the polynomial in
+        # evaluation form
+        assert p_z_coeff == p_z_eval == poly_eval[i]
diff --git a/tests/core/pyspec/eth2spec/test/helpers/sharding.py b/tests/core/pyspec/eth2spec/test/helpers/sharding.py
index 4523e07f2..2ea8c94bc 100644
--- a/tests/core/pyspec/eth2spec/test/helpers/sharding.py
+++ b/tests/core/pyspec/eth2spec/test/helpers/sharding.py
@@ -66,6 +66,38 @@ def get_sample_blob(spec, rng=None):
     return spec.Blob(b)
 
 
+def eval_poly_in_coeff_form(spec, coeffs, x):
+    """
+    Evaluate a polynomial in coefficient form at 'x' using Horner's rule
+    """
+    total = 0
+    for a in reversed(coeffs):
+        total = (total * x + a) % spec.BLS_MODULUS
+    return total % spec.BLS_MODULUS
+
+
+def get_poly_in_both_forms(spec, rng=None):
+    """
+    Generate and return a random polynomial in both coefficient form and evaluation form
+    """
+    if rng is None:
+        rng = random.Random(5566)
+
+    roots_of_unity_brp = spec.bit_reversal_permutation(spec.ROOTS_OF_UNITY)
+
+    coeffs = [
+        rng.randint(0, spec.BLS_MODULUS - 1)
+        for _ in range(spec.FIELD_ELEMENTS_PER_BLOB)
+    ]
+
+    evals = [
+        eval_poly_in_coeff_form(spec, coeffs, int(z))
+        for z in roots_of_unity_brp
+    ]
+
+    return coeffs, evals
+
+
 def get_sample_opaque_tx(spec, blob_count=1, rng=None):
     blobs = []
     blob_kzg_commitments = []