EIP4844: Handle barycentric evaluation at roots of unity

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)

tests/core/pyspec/eth2spec/test/eip4844/unittests/polynomial_commitments/test_polynomial_commitments.py

@@ -1,9 +1,13 @@
+import random
+
 from eth2spec.test.context import (
     spec_state_test,
     with_eip4844_and_later,
 )
 from eth2spec.test.helpers.sharding import (
     get_sample_blob,
+    get_poly_in_both_forms,
+    eval_poly_in_coeff_form,
 )
 
 
@@ -18,3 +22,68 @@ 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_outside_domain(spec, state):
+    """
+    Test barycentric formula correctness by using it to evaluate a polynomial at a bunch of points outside 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
+    """
+    rng = random.Random(5566)
+    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_samples = 12
+
+    for _ in range(n_samples):
+        # Get a random evaluation point and make sure it's not a root of unity
+        z = rng.randint(0, spec.BLS_MODULUS - 1)
+        while z in roots_of_unity_brp:
+            z = rng.randint(0, spec.BLS_MODULUS - 1)
+
+        # 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)
+
+        # Both evaluations should agree
+        assert p_z_coeff == p_z_eval
+
+
+@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]

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 = []