chore: add recover_cells_and_compute_proofs method (#5938)

* chore: add recover_cells_and_compute_proofs method * Introduce type alias `CellsAndKzgProofs` to address type complexity. --------- Co-authored-by: Jimmy Chen <jchen.tc@gmail.com>
sigp · Jun 19, 2024 · 6ff9480 · 6ff9480
1 parent 14c0d3b
commit 6ff9480
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 48 deletions.
diff --git a/consensus/types/src/data_column_sidecar.rs b/consensus/types/src/data_column_sidecar.rs
@@ -225,15 +225,7 @@ impl<E: EthSpec> DataColumnSidecar<E> {
                     cells.push(ssz_cell_to_crypto_cell::<E>(cell)?);
                     cell_ids.push(data_column.index);
                 }
-                // recover_all_cells does not expect sorted
-                let all_cells = kzg.recover_all_cells(&cell_ids, &cells)?;
-                let blob = kzg.cells_to_blob(&all_cells)?;
-
-                // Note: This function computes all cells and proofs. According to Justin this is okay,
-                // computing a partial set may be more expensive and requires code paths that don't exist.
-                // Computing the blobs cells is technically unnecessary but very cheap. It's done here again
-                // for simplicity.
-                kzg.compute_cells_and_proofs(&blob)
+                kzg.recover_cells_and_compute_kzg_proofs(&cell_ids, &cells)
             })
             .collect::<Result<Vec<_>, KzgError>>()?;
 

diff --git a/crypto/kzg/src/lib.rs b/crypto/kzg/src/lib.rs
@@ -16,6 +16,11 @@ pub use c_kzg::{
 pub use c_kzg::{Cell, CELLS_PER_EXT_BLOB};
 use mockall::automock;
 
+pub type CellsAndKzgProofs = (
+    Box<[Cell; CELLS_PER_EXT_BLOB]>,
+    Box<[KzgProof; CELLS_PER_EXT_BLOB]>,
+);
+
 #[derive(Debug)]
 pub enum Error {
     /// An error from the underlying kzg library.
@@ -152,17 +157,7 @@ impl Kzg {
     }
 
     /// Computes the cells and associated proofs for a given `blob` at index `index`.
-    #[allow(clippy::type_complexity)]
-    pub fn compute_cells_and_proofs(
-        &self,
-        blob: &Blob,
-    ) -> Result<
-        (
-            Box<[Cell; CELLS_PER_EXT_BLOB]>,
-            Box<[KzgProof; CELLS_PER_EXT_BLOB]>,
-        ),
-        Error,
-    > {
+    pub fn compute_cells_and_proofs(&self, blob: &Blob) -> Result<CellsAndKzgProofs, Error> {
         let (cells, proofs) = c_kzg::Cell::compute_cells_and_kzg_proofs(blob, &self.trusted_setup)
             .map_err(Into::<Error>::into)?;
         let proofs = Box::new(proofs.map(|proof| KzgProof::from(proof.to_bytes().into_inner())));
@@ -196,39 +191,28 @@ impl Kzg {
         }
     }
 
-    pub fn cells_to_blob(&self, cells: &[Cell; c_kzg::CELLS_PER_EXT_BLOB]) -> Result<Blob, Error> {
+    fn cells_to_blob(&self, cells: &[Cell; c_kzg::CELLS_PER_EXT_BLOB]) -> Result<Blob, Error> {
         Ok(Blob::cells_to_blob(cells)?)
     }
 
-    pub fn recover_all_cells(
+    pub fn recover_cells_and_compute_kzg_proofs(
         &self,
         cell_ids: &[u64],
         cells: &[Cell],
-    ) -> Result<Box<[Cell; c_kzg::CELLS_PER_EXT_BLOB]>, Error> {
-        Ok(c_kzg::Cell::recover_all_cells(
-            cell_ids,
-            cells,
-            &self.trusted_setup,
-        )?)
+    ) -> Result<CellsAndKzgProofs, Error> {
+        let all_cells = c_kzg::Cell::recover_all_cells(cell_ids, cells, &self.trusted_setup)?;
+        let blob = self.cells_to_blob(&all_cells)?;
+        self.compute_cells_and_proofs(&blob)
     }
 }
 
 pub mod mock {
-    use crate::{Error, KzgProof};
+    use crate::{CellsAndKzgProofs, Error, KzgProof};
     use c_kzg::{Blob, Cell, CELLS_PER_EXT_BLOB};
 
     pub const MOCK_KZG_BYTES_PER_CELL: usize = 2048;
 
-    #[allow(clippy::type_complexity)]
-    pub fn compute_cells_and_proofs(
-        _blob: &Blob,
-    ) -> Result<
-        (
-            Box<[Cell; CELLS_PER_EXT_BLOB]>,
-            Box<[KzgProof; CELLS_PER_EXT_BLOB]>,
-        ),
-        Error,
-    > {
+    pub fn compute_cells_and_proofs(_blob: &Blob) -> Result<CellsAndKzgProofs, Error> {
         let empty_cell = Cell::new([0; MOCK_KZG_BYTES_PER_CELL]);
         Ok((
             Box::new([empty_cell; CELLS_PER_EXT_BLOB]),

diff --git a/testing/ef_tests/src/cases/kzg_compute_cells_and_kzg_proofs.rs b/testing/ef_tests/src/cases/kzg_compute_cells_and_kzg_proofs.rs
@@ -1,7 +1,6 @@
 use super::*;
 use crate::case_result::compare_result;
-use kzg::{Blob as KzgBlob, Cell};
-use kzg::{KzgProof, CELLS_PER_EXT_BLOB};
+use kzg::{Blob as KzgBlob, CellsAndKzgProofs};
 use serde::Deserialize;
 use std::marker::PhantomData;
 
@@ -62,12 +61,6 @@ impl<E: EthSpec> Case for KZGComputeCellsAndKZGProofs<E> {
                 .ok()
         });
 
-        compare_result::<
-            (
-                Box<[Cell; CELLS_PER_EXT_BLOB]>,
-                Box<[KzgProof; CELLS_PER_EXT_BLOB]>,
-            ),
-            _,
-        >(&cells_and_proofs, &expected)
+        compare_result::<CellsAndKzgProofs, _>(&cells_and_proofs, &expected)
     }
 }