Add helpers: recover_shifted_data() and recover_original_data()

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

@@ -508,29 +508,15 @@ def construct_vanishing_polynomial(cell_ids: Sequence[CellID],
 
 ### `recover_shifted_data`
 
-### `recover_original_data`
-
-### `recover_polynomial`
-
 ```python
-def recover_polynomial(cell_ids: Sequence[CellID],
-                       cells_bytes: Sequence[Vector[Bytes32, FIELD_ELEMENTS_PER_CELL]]) -> Polynomial:
-    """
-    Recovers a polynomial from 2 * FIELD_ELEMENTS_PER_CELL evaluations, half of which can be missing.
-
-    This algorithm uses FFTs to recover cells faster than using Lagrange implementation. However,
-    a faster version thanks to Qi Zhou can be found here:
-    https://github.com/ethereum/research/blob/51b530a53bd4147d123ab3e390a9d08605c2cdb8/polynomial_reconstruction/polynomial_reconstruction_danksharding.py
-
-    Public method.
-    """
-    assert len(cell_ids) == len(cells_bytes)
-
-    cells = [bytes_to_cell(cell_bytes) for cell_bytes in cells_bytes]
-    assert len(cells) >= CELLS_PER_BLOB // 2
-
-    zero_poly_coeff, zero_poly_eval, zero_poly_eval_brp = construct_vanishing_polynomial(cell_ids, cells)
-
+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]:
     extended_evaluation_rbo = [0] * (FIELD_ELEMENTS_PER_BLOB * 2)
     for cell_id, cell in zip(cell_ids, cells):
         start = cell_id * FIELD_ELEMENTS_PER_CELL
@@ -541,8 +527,6 @@ def recover_polynomial(cell_ids: Sequence[CellID],
     extended_evaluation_times_zero = [BLSFieldElement(int(a) * int(b) % BLS_MODULUS)
                                       for a, b in zip(zero_poly_eval, extended_evaluation)]
 
-    roots_of_unity_extended = compute_roots_of_unity(2 * FIELD_ELEMENTS_PER_BLOB)
-
     extended_evaluations_fft = fft_field(extended_evaluation_times_zero, roots_of_unity_extended, inv=True)
 
     shift_factor = BLSFieldElement(PRIMITIVE_ROOT_OF_UNITY)
@@ -554,6 +538,16 @@ def recover_polynomial(cell_ids: Sequence[CellID],
     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
+```
+
+### `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]:
     eval_shifted_reconstructed_poly = [
         div(a, b)
         for a, b in zip(eval_shifted_extended_evaluation, eval_shifted_zero_poly)
@@ -565,6 +559,39 @@ def recover_polynomial(cell_ids: Sequence[CellID],
 
     reconstructed_data = bit_reversal_permutation(fft_field(reconstructed_poly, roots_of_unity_extended))
 
+    return reconstructed_data
+```
+
+### `recover_polynomial`
+
+```python
+def recover_polynomial(cell_ids: Sequence[CellID],
+                       cells_bytes: Sequence[Vector[Bytes32, FIELD_ELEMENTS_PER_CELL]]) -> Polynomial:
+    """
+    Recovers a polynomial from 2 * FIELD_ELEMENTS_PER_CELL evaluations, half of which can be missing.
+
+    This algorithm uses FFTs to recover cells faster than using Lagrange implementation. However,
+    a faster version thanks to Qi Zhou can be found here:
+    https://github.com/ethereum/research/blob/51b530a53bd4147d123ab3e390a9d08605c2cdb8/polynomial_reconstruction/polynomial_reconstruction_danksharding.py
+
+    Public method.
+    """
+    assert len(cell_ids) == len(cells_bytes)
+
+    cells = [bytes_to_cell(cell_bytes) for cell_bytes in cells_bytes]
+    assert len(cells) >= CELLS_PER_BLOB // 2
+
+    roots_of_unity_extended = compute_roots_of_unity(2 * FIELD_ELEMENTS_PER_BLOB)
+
+    zero_poly_coeff, zero_poly_eval, zero_poly_eval_brp = construct_vanishing_polynomial(cell_ids, cells)
+
+    eval_shifted_extended_evaluation, eval_shifted_zero_poly, shift_inv = \
+        recover_shifted_data(cell_ids, cells, 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)
+
     for cell_id, cell in zip(cell_ids, cells):
         start = cell_id * FIELD_ELEMENTS_PER_CELL
         end = (cell_id + 1) * FIELD_ELEMENTS_PER_CELL