chore: Refactor cell recovery code

specs/_features/eip7594/polynomial-commitments-sampling.md

@@ -20,13 +20,13 @@
   - [FFTs](#ffts)
     - [`_fft_field`](#_fft_field)
     - [`fft_field`](#fft_field)
+    - [`coset_fft_field`](#coset_fft_field)
   - [Polynomials in coefficient form](#polynomials-in-coefficient-form)
     - [`polynomial_eval_to_coeff`](#polynomial_eval_to_coeff)
     - [`add_polynomialcoeff`](#add_polynomialcoeff)
     - [`neg_polynomialcoeff`](#neg_polynomialcoeff)
     - [`multiply_polynomialcoeff`](#multiply_polynomialcoeff)
     - [`divide_polynomialcoeff`](#divide_polynomialcoeff)
-    - [`shift_polynomialcoeff`](#shift_polynomialcoeff)
     - [`interpolate_polynomialcoeff`](#interpolate_polynomialcoeff)
     - [`vanishing_polynomialcoeff`](#vanishing_polynomialcoeff)
     - [`evaluate_polynomialcoeff`](#evaluate_polynomialcoeff)
@@ -44,8 +44,7 @@
     - [`verify_cell_kzg_proof_batch`](#verify_cell_kzg_proof_batch)
 - [Reconstruction](#reconstruction)
   - [`construct_vanishing_polynomial`](#construct_vanishing_polynomial)
-  - [`recover_shifted_data`](#recover_shifted_data)
-  - [`recover_original_data`](#recover_original_data)
+  - [`recover_data`](#recover_data)
   - [`recover_all_cells`](#recover_all_cells)
 
 <!-- END doctoc generated TOC please keep comment here to allow auto update -->
@@ -180,6 +179,40 @@ def fft_field(vals: Sequence[BLSFieldElement],
         return _fft_field(vals, roots_of_unity)
 ```
 
+#### `coset_fft_field`
+
+```python
+def coset_fft_field(vals: Sequence[BLSFieldElement], 
+                    roots_of_unity: Sequence[BLSFieldElement], 
+                    inv: bool=False) -> Sequence[BLSFieldElement]:
+    """
+    Computes an FFT/IFFT over a coset of the roots of unity. 
+    This is useful for when one wants to divide by a polynomial which 
+    vanishes on one or more elements in the domain.
+    """
+
+    vals = vals.copy()
+
+    # Multiply each entry in `vals` by succeeding powers of `factor`
+    # ie [vals[0] * factor^0, vals[1] * factor^1, ..., vals[n] * factor^n]
+    def shift_vals(vals: Sequence[BLSFieldElement], factor: BLSFieldElement) -> Sequence[BLSFieldElement]:
+        shift = 1
+        for i in range(len(vals)):
+            vals[i] = BLSFieldElement((int(vals[i]) * shift) % BLS_MODULUS)
+            shift = (shift * int(factor)) % BLS_MODULUS
+        return vals
+
+    # This is the coset generator; it is used to compute a FFT/IFFT over a coset of
+    # the roots of unity.
+    shift_factor = BLSFieldElement(PRIMITIVE_ROOT_OF_UNITY)
+    if inv:
+        vals = fft_field(vals, roots_of_unity, inv)
+        shift_inv = bls_modular_inverse(shift_factor)
+        return shift_vals(vals, shift_inv)
+    else:
+        vals = shift_vals(vals, shift_factor)
+        return fft_field(vals, roots_of_unity, inv)
+```
 
 ### Polynomials in coefficient form
 
@@ -257,23 +290,6 @@ def divide_polynomialcoeff(a: PolynomialCoeff, b: PolynomialCoeff) -> Polynomial
     return [x % BLS_MODULUS for x in o]
 ```
 
-#### `shift_polynomialcoeff`
-
-```python
-def shift_polynomialcoeff(polynomial_coeff: PolynomialCoeff, factor: BLSFieldElement) -> PolynomialCoeff:
-    """
-    Shift the evaluation of a polynomial in coefficient form by factor.
-    This returns a new polynomial g in coefficient form such that g(x) = f(factor * x).
-    In other words, each coefficient of f is scaled by a power of factor.
-    """
-    factor_power = 1
-    o = []
-    for p in polynomial_coeff:
-        o.append(int(p) * factor_power % BLS_MODULUS)
-        factor_power = factor_power * int(factor) % BLS_MODULUS
-    return o
-```
-
 #### `interpolate_polynomialcoeff`
 
 ```python
@@ -555,13 +571,25 @@ def verify_cell_kzg_proof_batch(row_commitments_bytes: Sequence[Bytes48],
 ### `construct_vanishing_polynomial`
 
 ```python
-def construct_vanishing_polynomial(missing_cell_ids: Sequence[CellID]) -> Tuple[
-        Sequence[BLSFieldElement],
-        Sequence[BLSFieldElement]]:
+def construct_vanishing_polynomial(missing_cell_ids: Sequence[CellID]) -> Sequence[BLSFieldElement]:
     """
-    Given the cells that are missing from the data, compute the polynomial that vanishes at every point that
+    Given the cells IDs that are missing from the data, compute the polynomial that vanishes at every point that
     corresponds to a missing field element.
     """
+
+    # If all of the cells are missing, then the vanishing polynomial
+    # can be computed as Z(x) = x^n - 1, where `n` is FIELD_ELEMENTS_PER_EXT_BLOB
+    # 
+    # Note: this makes the function complete on all inputs, however this code path should not 
+    # be hit since this method is used for data recovery and if all of the
+    # cells are missing, then we can return early.  
+    if len(missing_cell_ids) == CELLS_PER_EXT_BLOB:
+        z_x = [BLSFieldElement(0)] * (FIELD_ELEMENTS_PER_EXT_BLOB + 1)
+        z_x[0] = BLS_MODULUS - 1
+        z_x[-1] = 1
+
+        return z_x
+
     # Get the small domain
     roots_of_unity_reduced = compute_roots_of_unity(CELLS_PER_EXT_BLOB)
 
@@ -576,86 +604,61 @@ def construct_vanishing_polynomial(missing_cell_ids: Sequence[CellID]) -> Tuple[
     for i, coeff in enumerate(short_zero_poly):
         zero_poly_coeff[i * FIELD_ELEMENTS_PER_CELL] = coeff
 
-    # Compute evaluations of the extended vanishing polynomial
-    zero_poly_eval = fft_field(zero_poly_coeff,
-                               compute_roots_of_unity(FIELD_ELEMENTS_PER_EXT_BLOB))
-    zero_poly_eval_brp = bit_reversal_permutation(zero_poly_eval)
-
-    # Sanity check
-    for cell_id in range(CELLS_PER_EXT_BLOB):
-        start = cell_id * FIELD_ELEMENTS_PER_CELL
-        end = (cell_id + 1) * FIELD_ELEMENTS_PER_CELL
-        if cell_id in missing_cell_ids:
-            assert all(a == 0 for a in zero_poly_eval_brp[start:end])
-        else:  # cell_id in cell_ids
-            assert all(a != 0 for a in zero_poly_eval_brp[start:end])
-
-    return zero_poly_coeff, zero_poly_eval
+    return zero_poly_coeff
 ```
 
-### `recover_shifted_data`
+### `recover_data`
 
 ```python
-def recover_shifted_data(cell_ids: Sequence[CellID],
-                         cells: Sequence[Cell],
-                         zero_poly_eval: Sequence[BLSFieldElement],
-                         zero_poly_coeff: Sequence[BLSFieldElement],
-                         roots_of_unity_extended: Sequence[BLSFieldElement]) -> Tuple[
-                             Sequence[BLSFieldElement],
-                             Sequence[BLSFieldElement],
-                             BLSFieldElement]:
+def recover_data(cell_ids: Sequence[CellID],
+                 cells: Sequence[Cell],
+                 ) -> Sequence[BLSFieldElement]:
     """
-    Given Z(x), return polynomial Q_1(x)=(E*Z)(k*x) and Q_2(x)=Z(k*x) and k^{-1}.
+    Recover the missing evaluations for the extended blob, given at least half of the evaluations.
     """
-    shift_factor = BLSFieldElement(PRIMITIVE_ROOT_OF_UNITY)
-    shift_inv = div(BLSFieldElement(1), shift_factor)
 
+    # Get the extended domain. This will be referred to as the FFT domain.
+    roots_of_unity_extended = compute_roots_of_unity(FIELD_ELEMENTS_PER_EXT_BLOB)
+
+    # Flatten the cells into evaluations.
+    # If a cell is missing, then its evaluation is zero.
     extended_evaluation_rbo = [0] * FIELD_ELEMENTS_PER_EXT_BLOB
     for cell_id, cell in zip(cell_ids, cells):
         start = cell_id * FIELD_ELEMENTS_PER_CELL
         end = (cell_id + 1) * FIELD_ELEMENTS_PER_CELL
         extended_evaluation_rbo[start:end] = cell
     extended_evaluation = bit_reversal_permutation(extended_evaluation_rbo)
 
-    # Compute (E*Z)(x)
+    # Compute Z(x) in monomial form
+    # Z(x) is the polynomial which vanishes on all of the evaluations which are missing
+    missing_cell_ids = [CellID(cell_id) for cell_id in range(CELLS_PER_EXT_BLOB) if cell_id not in cell_ids]
+
+    zero_poly_coeff = construct_vanishing_polynomial(missing_cell_ids)
+    # Convert Z(x) to evaluation form over the FFT domain
+    zero_poly_eval = fft_field(zero_poly_coeff, roots_of_unity_extended)
+    # Compute (E*Z)(x) = E(x) * Z(x) in evaluation form over the FFT domain
     extended_evaluation_times_zero = [BLSFieldElement(int(a) * int(b) % BLS_MODULUS)
                                       for a, b in zip(zero_poly_eval, extended_evaluation)]
+    # Convert (E*Z)(x) to monomial form 
+    extended_evaluation_times_zero_coeffs = fft_field(extended_evaluation_times_zero, roots_of_unity_extended, inv=True)
 
-    extended_evaluations_fft = fft_field(extended_evaluation_times_zero, roots_of_unity_extended, inv=True)
+    # Convert (E*Z)(x) to evaluation form over a coset of the FFT domain
+    extended_evaluations_over_coset = coset_fft_field(extended_evaluation_times_zero_coeffs, roots_of_unity_extended)
 
-    # Compute (E*Z)(k*x)
-    shifted_extended_evaluation = shift_polynomialcoeff(extended_evaluations_fft, shift_factor)
-    # Compute Z(k*x)
-    shifted_zero_poly = shift_polynomialcoeff(zero_poly_coeff, shift_factor)
-
-    eval_shifted_extended_evaluation = fft_field(shifted_extended_evaluation, roots_of_unity_extended)
-    eval_shifted_zero_poly = fft_field(shifted_zero_poly, roots_of_unity_extended)
-
-    return eval_shifted_extended_evaluation, eval_shifted_zero_poly, shift_inv
-```
+    # Convert Z(x) to evaluation form over a coset of the FFT domain
+    zero_poly_over_coset = coset_fft_field(zero_poly_coeff, roots_of_unity_extended)
 
-### `recover_original_data`
-
-```python
-def recover_original_data(eval_shifted_extended_evaluation: Sequence[BLSFieldElement],
-                          eval_shifted_zero_poly: Sequence[BLSFieldElement],
-                          shift_inv: BLSFieldElement,
-                          roots_of_unity_extended: Sequence[BLSFieldElement]) -> Sequence[BLSFieldElement]:
-    """
-    Given Q_1, Q_2 and k^{-1}, compute P(x).
-    """
-    # Compute Q_3 = Q_1(x)/Q_2(x) = P(k*x)
-    eval_shifted_reconstructed_poly = [
+    # Compute H(x) = (E*Z)(x) / Z(x) in evaluation form over a coset of the FFT domain
+    reconstructed_poly_over_coset = [
         div(a, b)
-        for a, b in zip(eval_shifted_extended_evaluation, eval_shifted_zero_poly)
+        for a, b in zip(extended_evaluations_over_coset, zero_poly_over_coset)
     ]
 
-    shifted_reconstructed_poly = fft_field(eval_shifted_reconstructed_poly, roots_of_unity_extended, inv=True)
-
-    # Unshift P(k*x) by k^{-1} to get P(x)
-    reconstructed_poly = shift_polynomialcoeff(shifted_reconstructed_poly, shift_inv)
+    # Convert H(x) to monomial form
+    reconstructed_poly_coeff = coset_fft_field(reconstructed_poly_over_coset, roots_of_unity_extended, inv=True)
 
-    reconstructed_data = bit_reversal_permutation(fft_field(reconstructed_poly, roots_of_unity_extended))
+    # Convert H(x) to evaluation form over the FFT domain and bit reverse the result
+    reconstructed_data = bit_reversal_permutation(fft_field(reconstructed_poly_coeff, roots_of_unity_extended))
 
     return reconstructed_data
 ```
@@ -687,28 +690,12 @@ def recover_all_cells(cell_ids: Sequence[CellID], cells: Sequence[Cell]) -> Sequ
     for cell_id in cell_ids:
         assert cell_id < CELLS_PER_EXT_BLOB
 
-    # Get the extended domain
-    roots_of_unity_extended = compute_roots_of_unity(FIELD_ELEMENTS_PER_EXT_BLOB)
-
     # Convert cells to coset evals
     cosets_evals = [cell_to_coset_evals(cell) for cell in cells]
 
-    missing_cell_ids = [CellID(cell_id) for cell_id in range(CELLS_PER_EXT_BLOB) if cell_id not in cell_ids]
-    zero_poly_coeff, zero_poly_eval = construct_vanishing_polynomial(missing_cell_ids)
-
-    eval_shifted_extended_evaluation, eval_shifted_zero_poly, shift_inv = recover_shifted_data(
+    reconstructed_data = recover_data(
         cell_ids,
-        cosets_evals,
-        zero_poly_eval,
-        zero_poly_coeff,
-        roots_of_unity_extended,
-    )
-
-    reconstructed_data = recover_original_data(
-        eval_shifted_extended_evaluation,
-        eval_shifted_zero_poly,
-        shift_inv,
-        roots_of_unity_extended,
+        cosets_evals
     )
 
     for cell_id, coset_evals in zip(cell_ids, cosets_evals):

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

@@ -28,6 +28,61 @@ def test_fft(spec):
     assert poly_coeff_inversed == poly_coeff
 
 
+@with_eip7594_and_later
+@spec_test
+@single_phase
+def test_coset_fft(spec):
+    rng = random.Random(5566)
+
+    roots_of_unity = spec.compute_roots_of_unity(spec.FIELD_ELEMENTS_PER_BLOB)
+
+    poly_coeff = [rng.randint(0, BLS_MODULUS - 1) for _ in range(spec.FIELD_ELEMENTS_PER_BLOB)]
+
+    poly_eval = spec.coset_fft_field(poly_coeff, roots_of_unity)
+    poly_coeff_inversed = spec.coset_fft_field(poly_eval, roots_of_unity, inv=True)
+
+    assert len(poly_eval) == len(poly_coeff) == len(poly_coeff_inversed)
+    assert poly_coeff_inversed == poly_coeff
+
+
+@with_eip7594_and_later
+@spec_test
+@single_phase
+def test_construct_vanishing_polynomial(spec):
+    rng = random.Random(5566)
+
+    num_missing_cells = rng.randint(0, spec.CELLS_PER_EXT_BLOB - 1)
+    # Get a unique list of `num_missing_cells` cell IDs
+    unique_missing_cell_ids = rng.sample(range(spec.CELLS_PER_EXT_BLOB), num_missing_cells)
+
+    zero_poly_coeff = spec.construct_vanishing_polynomial(unique_missing_cell_ids)
+    roots_of_unity = spec.compute_roots_of_unity(spec.FIELD_ELEMENTS_PER_EXT_BLOB)
+    zero_poly_eval = spec.fft_field(zero_poly_coeff, roots_of_unity)
+    zero_poly_eval_brp = spec.bit_reversal_permutation(zero_poly_eval)
+
+    for cell_id in range(spec.CELLS_PER_EXT_BLOB):
+        start = cell_id * spec.FIELD_ELEMENTS_PER_CELL
+        end = (cell_id + 1) * spec.FIELD_ELEMENTS_PER_CELL
+        if cell_id in unique_missing_cell_ids:
+            assert all(a == 0 for a in zero_poly_eval_brp[start:end])
+        else:  # cell_id in cell_ids
+            assert all(a != 0 for a in zero_poly_eval_brp[start:end])
+
+
+@with_eip7594_and_later
+@spec_test
+@single_phase
+def test_construct_vanishing_polynomial_all_missing_zeros(spec):
+
+    all_cell_ids = range(spec.CELLS_PER_EXT_BLOB)
+
+    zero_poly_coeff = spec.construct_vanishing_polynomial(all_cell_ids)
+    roots_of_unity = spec.compute_roots_of_unity(spec.FIELD_ELEMENTS_PER_EXT_BLOB)
+    zero_poly_eval = spec.fft_field(zero_poly_coeff, roots_of_unity)
+
+    assert all(a == 0 for a in zero_poly_eval)
+
+
 @with_eip7594_and_later
 @spec_test
 @single_phase