feat: Implement compute_cells peerdas test case (#1281)

Author: Kayden-ML

crates/common/consensus/beacon/src/matrix_entry.rs

@@ -1,18 +1,9 @@
-use alloy_primitives::B256;
 use anyhow::{Ok, Result, anyhow, ensure};
-use discv5::enr::k256::{
-    Scalar,
-    elliptic_curve::{Field, PrimeField},
-};
 use ream_consensus_misc::{
-    constants::beacon::{
-        BYTES_PER_BLOB, BYTES_PER_FIELD_ELEMENT, CELLS_PER_EXT_BLOB, FIELD_ELEMENTS_PER_BLOB,
-        FIELD_ELEMENTS_PER_CELL, FIELD_ELEMENTS_PER_EXT_BLOB,
-    },
-    polynomial_commitments::kzg_proof::KZGProof,
+    constants::beacon::CELLS_PER_EXT_BLOB, polynomial_commitments::kzg_proof::KZGProof,
 };
 use ream_execution_rpc_types::get_blobs::Blob;
-use rust_eth_kzg::{Cell as KZGCell, CellIndex, DASContext, KZGProof as Proof};
+use rust_eth_kzg::{Cell as KZGCell, DASContext, KZGProof as Proof};
 use ssz_types::FixedVector;
 
 use crate::data_column_sidecar::Cell;
@@ -144,231 +135,25 @@ fn convert_kzg_proof(kzg_proof: Proof) -> KZGProof {
     KZGProof::from(kzg_proof)
 }
 
-/// Convert untrusted bytes to a trusted and validated BLS scalar field element.
-/// This function does not accept inputs greater than the BLS modulus.
-pub fn bytes_to_bls_field(b: B256) -> anyhow::Result<Scalar> {
-    ensure!(b.len() == 32, "Invalid input length for field element");
-
-    let field_element = Scalar::from_repr((*b).into())
-        .into_option()
-        .ok_or_else(|| anyhow!("Bytes exceed the field modulus"))?;
-    Ok(field_element)
-}
-
-/// Convert a blob to list of BLS field scalars.
-pub fn blob_to_polynomial(blob: Blob) -> anyhow::Result<[Scalar; FIELD_ELEMENTS_PER_BLOB]> {
-    let mut polynomial = [Scalar::default(); FIELD_ELEMENTS_PER_BLOB];
-
-    for (i, polynomial_element) in polynomial.iter_mut().enumerate() {
-        let start = i * BYTES_PER_FIELD_ELEMENT;
-        let end = (i + 1) * BYTES_PER_FIELD_ELEMENT;
-
-        let chunk: B256 = blob.inner[start..end]
-            .try_into()
-            .map_err(|err| anyhow!("Invalid chunk size at index {err:?}"))?;
-
-        *polynomial_element = bytes_to_bls_field(chunk)?;
-    }
-
-    Ok(polynomial)
-}
-
-/// Return ``x`` to power of [0, n-1], if n > 0. When n==0, an empty array is returned.
-pub fn compute_powers(x: Scalar, n: u64) -> anyhow::Result<Vec<Scalar>> {
-    let mut powers = Vec::with_capacity(n as usize);
-    let mut current_power = Scalar::ONE;
-
-    for _ in 0..n {
-        powers.push(current_power);
-        current_power *= x;
-    }
-
-    Ok(powers)
-}
-
-/// Return roots of unity of ``order``.
-pub fn compute_roots_of_unity(order: u64) -> anyhow::Result<Vec<Scalar>> {
-    ensure!(
-        order <= (1 << Scalar::S),
-        "Order exceeds maximum supported by field"
-    );
-    let ratio = (1 << Scalar::S) / order;
-    let exponent: [u64; 4] = [ratio, 0, 0, 0];
-
-    let root = Scalar::ROOT_OF_UNITY.pow_vartime(exponent);
-
-    Ok(compute_powers(root, order)?)
-}
-
-fn _fft_field(vals: Vec<Scalar>, roots_of_unity: Vec<Scalar>) -> anyhow::Result<Vec<Scalar>> {
-    let n = vals.len();
-    if n <= 1 {
-        return Ok(vals);
-    }
-    let mut evens = Vec::with_capacity(n / 2);
-    let mut odds = Vec::with_capacity(n / 2);
-    for (i, val) in vals.into_iter().enumerate() {
-        if i % 2 == 0 {
-            evens.push(val);
-        } else {
-            odds.push(val);
-        }
-    }
-    let next_roots: Vec<Scalar> = roots_of_unity.iter().step_by(2).cloned().collect();
-    let left = _fft_field(evens, next_roots.clone())?;
-    let right = _fft_field(odds, next_roots)?;
-
-    let mut result = vec![Scalar::ZERO; n];
-    for i in 0..(n / 2) {
-        let root = roots_of_unity[i];
-        let y_times_root = right[i] * root;
-        result[i] = left[i] + y_times_root;
-        result[i + (n / 2)] = left[i] - y_times_root;
-    }
-
-    Ok(result)
-}
-
-pub fn fft_field(
-    vals: Vec<Scalar>,
-    roots_of_unity: Vec<Scalar>,
-    inv: bool,
-) -> anyhow::Result<Vec<Scalar>> {
-    if inv {
-        let n = vals.len() as u64;
-        let invlen = Scalar::from(n).invert().unwrap();
-
-        let mut reversed_roots = Vec::with_capacity(roots_of_unity.len());
-        if !roots_of_unity.is_empty() {
-            reversed_roots.push(roots_of_unity[0]);
-            reversed_roots.extend(roots_of_unity.iter().skip(1).rev());
-        }
-
-        let fft_result = _fft_field(vals, reversed_roots)?;
-        let result = fft_result.into_iter().map(|x| x * invlen).collect();
-
-        Ok(result)
-    } else {
-        Ok(_fft_field(vals, roots_of_unity)?)
-    }
-}
-
-/// Reverse the bit order of an integer ``n``.
-pub fn reverse_bits(n: usize, order: usize) -> anyhow::Result<usize> {
-    ensure!(order.is_power_of_two(), "Order must be a power of two");
-
-    let width = order.trailing_zeros();
-    let mut result = 0;
-    let mut temp_n = n;
-
-    for _ in 0..width {
-        result = (result << 1) | (temp_n & 1);
-        temp_n >>= 1;
-    }
-
-    Ok(result)
-}
-
-/// Return a copy with bit-reversed permutation. The permutation is an involution (inverts itself).
-///
-/// The input and output are a sequence of generic type ``T`` objects.
-pub fn bit_reversal_permutation<T: Clone>(sequence: Vec<T>) -> anyhow::Result<Vec<T>> {
-    let n = sequence.len();
-    let mut result = Vec::with_capacity(n);
-
-    for i in 0..n {
-        let rev_idx = reverse_bits(i, n)?;
-        result.push(sequence[rev_idx].clone());
-    }
-
-    Ok(result)
-}
-
-/// Interpolates a polynomial (given in evaluation form) to a polynomial in coefficient form.
-pub fn polynomial_eval_to_coeff(
-    polynomial: [Scalar; FIELD_ELEMENTS_PER_BLOB],
-) -> anyhow::Result<Vec<Scalar>> {
-    let roots_of_unity = compute_roots_of_unity(FIELD_ELEMENTS_PER_BLOB as u64)?;
-    let rearranged = bit_reversal_permutation(polynomial.to_vec())?;
-
-    fft_field(rearranged, roots_of_unity, true)
-}
-
-/// Get the coset for a given ``cell_index``.
-/// Precisely, consider the group of roots of unity of order FIELD_ELEMENTS_PER_CELL *
-/// CELLS_PER_EXT_BLOB. Let G = {1, g, g^2, ...} denote its subgroup of order
-/// FIELD_ELEMENTS_PER_CELL. Then, the coset is defined as h * G = {h, hg, hg^2, ...}.
-/// This function, returns the coset.
-pub fn coset_for_cell(cell_index: CellIndex) -> anyhow::Result<[Scalar; FIELD_ELEMENTS_PER_CELL]> {
-    ensure!(
-        cell_index < CELLS_PER_EXT_BLOB,
-        "Cell index great then CELLS_PER_EXT_BLOB"
-    );
-    let roots = compute_roots_of_unity(FIELD_ELEMENTS_PER_EXT_BLOB as u64)?;
-    let roots_of_unity_brp = bit_reversal_permutation(roots)?;
-
-    let start = (cell_index as usize) * FIELD_ELEMENTS_PER_CELL;
-    let end = start + FIELD_ELEMENTS_PER_CELL;
-
-    let coset_slice = &roots_of_unity_brp[start..end];
-
-    let coset: [Scalar; FIELD_ELEMENTS_PER_CELL] = coset_slice
-        .try_into()
-        .map_err(|err| anyhow!("Slice length mismatch for coset conversion {err:?}"))?;
-
-    Ok(coset)
-}
-
-/// Evaluate a coefficient form polynomial at ``z`` using Horner's schema.
-pub fn evaluate_polynomialcoeff(polynomial_coeff: &[Scalar], z: Scalar) -> anyhow::Result<Scalar> {
-    let mut y = Scalar::ZERO;
-    for coef in polynomial_coeff.iter().rev() {
-        y = (y * z) + *coef;
-    }
-
-    Ok(y)
-}
-
-pub fn bls_field_to_bytes(x: Scalar) -> B256 {
-    let bytes: [u8; 32] = x.to_bytes().into();
-    B256::from(bytes)
-}
-
-/// Convert a trusted ``CosetEval`` into an untrusted ``Cell``.
-pub fn coset_evals_to_cell(coset_evals: [Scalar; FIELD_ELEMENTS_PER_CELL]) -> anyhow::Result<Cell> {
-    let mut cell_bytes = Vec::with_capacity(FIELD_ELEMENTS_PER_CELL * 32);
-    for eval in coset_evals {
-        let bytes: B256 = bls_field_to_bytes(eval);
-        cell_bytes.extend_from_slice(bytes.as_slice());
-    }
-
-    let cell = cell_bytes
+pub fn compute_cells(
+    blob: Blob,
+    das_context: &DASContext,
+) -> anyhow::Result<[Cell; CELLS_PER_EXT_BLOB as usize]> {
+    let blob_data: Vec<u8> = blob.inner.into();
+    let blob_bytes: &[u8; 131072] = blob_data
+        .as_slice()
         .try_into()
-        .map_err(|err| anyhow!("Failed to convert bytes to Cell FixedVector {err:?}"))?;
-    Ok(cell)
-}
-
-pub fn compute_cells(blob: Blob) -> anyhow::Result<[Cell; CELLS_PER_EXT_BLOB as usize]> {
-    ensure!(blob.inner.len() == BYTES_PER_BLOB, "Invalid blob length");
-    let polynomial = blob_to_polynomial(blob)?;
-    let polynomial_coeff = polynomial_eval_to_coeff(polynomial)?;
-    let mut cells = Vec::with_capacity(CELLS_PER_EXT_BLOB as usize);
+        .map_err(|err| anyhow!("Invalid blob size {err:?}"))?;
 
-    for i in 0..CELLS_PER_EXT_BLOB {
-        let coset = coset_for_cell(i as CellIndex)?;
-
-        let mut ys = [Scalar::ZERO; FIELD_ELEMENTS_PER_CELL];
-        for (j, &z) in coset.iter().enumerate() {
-            ys[j] = evaluate_polynomialcoeff(&polynomial_coeff, z)?;
-        }
+    let kzg_cells = das_context
+        .compute_cells(blob_bytes)
+        .map_err(|err| anyhow!("KZG error: {err:?}"))?;
 
-        let cell = coset_evals_to_cell(ys)?;
-        cells.push(cell);
-    }
+    let cells: Vec<Cell> = kzg_cells.into_iter().map(convert_cell).collect();
 
     let final_cells: [Cell; CELLS_PER_EXT_BLOB as usize] = cells
         .try_into()
-        .map_err(|err| anyhow!("Failed to convert cells to fixed array {err:?}"))?;
+        .map_err(|err| anyhow!("Failed to convert to fixed array {err:?}"))?;
 
     Ok(final_cells)
 }

crates/common/consensus/misc/src/constants/beacon.rs

@@ -39,8 +39,6 @@ pub const ETH1_FOLLOW_DISTANCE: u64 = 2048;
 pub const EXECUTION_PAYLOAD_INDEX: u64 = 9;
 pub const FAR_FUTURE_EPOCH: u64 = 18446744073709551615;
 pub const FIELD_ELEMENTS_PER_BLOB: usize = 4096;
-pub const FIELD_ELEMENTS_PER_CELL: usize = 64;
-pub const FIELD_ELEMENTS_PER_EXT_BLOB: usize = 8192;
 pub const FINALIZED_CHECKPOINT_INDEX: u64 = 20;
 pub const FULU_FORK_EPOCH: u64 = 411392;
 pub const GENESIS_SLOT: u64 = 0;

testing/ef-tests/src/macros/peerdas.rs

@@ -15,7 +15,7 @@ macro_rules! test_peerdas_kzg {
                 use ream_execution_rpc_types::get_blobs::Blob;
                 use rust_eth_kzg::{DASContext, UsePrecomp, TrustedSetup};
                 use ream_polynomial_commitments::handlers::verify_cell_kzg_proof_batch;
-                use ream_consensus_beacon::matrix_entry::{compute_cells_and_kzg_proofs, recover_cells_and_kzg_proofs};
+                use ream_consensus_beacon::matrix_entry::{compute_cells_and_kzg_proofs, recover_cells_and_kzg_proofs, compute_cells};
 
                 #[derive(Debug, Deserialize)]
                 pub struct KzgInput {
@@ -104,6 +104,21 @@ macro_rules! test_peerdas_kzg {
                                             )).unwrap()
                                         }).map_err(|err| anyhow!("Failed to get recover_cells_and_kzg_proofs {err:?}"))
                                     },
+                                    "compute_cells" => {
+                                        let blob_hex = input.blob.as_ref().expect("No blob");
+                                        let blob = Blob::from_ssz_bytes(
+                                            &const_hex::decode(blob_hex.strip_prefix("0x").unwrap_or(blob_hex))?
+                                        ).map_err(|err| anyhow!("Failed to decode {err:?}"))?;
+
+                                        compute_cells(blob, &context).map(|cells| {
+                                            let output: Vec<String> = cells
+                                                .iter()
+                                                .map(|cell| const_hex::encode_prefixed(cell.as_ref()))
+                                                .collect();
+
+                                            to_value(output).unwrap()
+                                        }).map_err(|err| anyhow!("Failed to execute compute_cells: {err:?}"))
+                                    },
                                     _ => Err(anyhow!("Unknown function path: {path}")),
                                 }
                             })();

testing/ef-tests/tests/tests.rs

@@ -226,3 +226,4 @@ test_sanity_blocks!(test_finality, "finality/finality");
 test_peerdas_kzg!(verify_cell_kzg_proof_batch);
 test_peerdas_kzg!(compute_cells_and_kzg_proofs);
 test_peerdas_kzg!(recover_cells_and_kzg_proofs);
+test_peerdas_kzg!(compute_cells);