Implement new KZG methods

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/CKZG4844.java

@@ -14,12 +14,18 @@
 package tech.pegasys.teku.kzg;
 
 import ethereum.ckzg4844.CKZG4844JNI;
+
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.IntStream;
+
+import ethereum.ckzg4844.CellsAndProofs;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.tuweni.bytes.Bytes;
 
+import static ethereum.ckzg4844.CKZG4844JNI.BYTES_PER_CELL;
+
 /**
  * Wrapper around jc-kzg-4844
  *
@@ -49,7 +55,9 @@ private CKZG4844() {
     }
   }
 
-  /** Only one trusted setup at a time can be loaded. */
+  /**
+   * Only one trusted setup at a time can be loaded.
+   */
   @Override
   public synchronized void loadTrustedSetup(final String trustedSetupFile) throws KZGException {
     if (loadedTrustedSetupFile.isPresent()
@@ -145,4 +153,41 @@ public KZGProof computeBlobKzgProof(final Bytes blob, final KZGCommitment kzgCom
           "Failed to compute KZG proof for blob with commitment " + kzgCommitment, ex);
     }
   }
+
+  @Override
+  public List<Cell> computeCells(Bytes blob) {
+    byte[] cellBytes = CKZG4844JNI.computeCells(blob.toArrayUnsafe());
+    return Cell.splitBytes(Bytes.wrap(cellBytes));
+  }
+
+  @Override
+  public List<CellAndProof> computeCellsAndProofs(Bytes blob) {
+    CellsAndProofs cellsAndProofs = CKZG4844JNI.computeCellsAndProofs(blob.toArrayUnsafe());
+    List<Cell> cells = Cell.splitBytes(Bytes.wrap(cellsAndProofs.getCells()));
+    List<KZGProof> proofs = KZGProof.splitBytes(Bytes.wrap(cellsAndProofs.getProofs()));
+    if (cells.size() != proofs.size()) throw new KZGException("Cells and proofs size differ");
+    return IntStream.range(0, cells.size())
+        .mapToObj(i -> new CellAndProof(cells.get(i), proofs.get(i)))
+        .toList();
+  }
+
+  @Override
+  public boolean verifyCellProof(KZGCommitment commitment, CellWithID cellWithID, KZGProof proof) {
+    return CKZG4844JNI.verifyCellProof(
+        commitment.toArrayUnsafe(),
+        cellWithID.id().id().longValue(),
+        cellWithID.cell().bytes().toArrayUnsafe(),
+        proof.toArrayUnsafe());
+  }
+
+  @Override
+  public List<Cell> recoverCells(List<CellWithID> cells) {
+    long[] cellIds = cells.stream().mapToLong(c -> c.id().id().longValue()).toArray();
+    byte[] cellBytes = CKZG4844Utils.flattenBytes(
+        cells.stream().map(c -> c.cell().bytes()).toList(),
+        cells.size() * BYTES_PER_CELL
+    );
+    byte[] recovered = CKZG4844JNI.recoverCells(cellIds, cellBytes);
+    return Cell.splitBytes(Bytes.wrap(recovered));
+  }
 }

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/CKZG4844Utils.java

@@ -24,10 +24,14 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.function.Function;
+import java.util.stream.IntStream;
+
 import org.apache.tuweni.bytes.Bytes;
 import tech.pegasys.teku.infrastructure.http.UrlSanitizer;
 import tech.pegasys.teku.infrastructure.io.resource.ResourceLoader;
 
+import static ethereum.ckzg4844.CKZG4844JNI.BYTES_PER_CELL;
+
 class CKZG4844Utils {
 
   private static final int MAX_BYTES_TO_FLATTEN = 100_663_296; // ~100.66 MB or 768 blobs
@@ -56,6 +60,16 @@ public static byte[] flattenG2Points(final List<Bytes> g2Points) {
     return flattenBytes(g2Points, CKZG4844JNI.BYTES_PER_G2 * g2Points.size());
   }
 
+  static List<Bytes> bytesChunked(Bytes bytes, int chunkSize) {
+    if (bytes.size() % chunkSize != 0) {
+      throw new IllegalArgumentException("Invalid bytes size: " + bytes.size());
+    }
+    return IntStream.range(0, bytes.size() / chunkSize)
+        .map(i -> i * chunkSize)
+        .mapToObj(startIdx -> bytes.slice(startIdx, chunkSize))
+        .toList();
+  }
+
   public static TrustedSetup parseTrustedSetupFile(final String trustedSetupFile)
       throws IOException {
     final String sanitizedTrustedSetup = UrlSanitizer.sanitizePotentialUrl(trustedSetupFile);
@@ -90,7 +104,7 @@ public static TrustedSetup parseTrustedSetupFile(final String trustedSetupFile)
     }
   }
 
-  private static byte[] flattenBytes(final List<Bytes> toFlatten, final int expectedSize) {
+  static byte[] flattenBytes(final List<Bytes> toFlatten, final int expectedSize) {
     return flattenBytes(toFlatten, Bytes::toArrayUnsafe, expectedSize);
   }
 

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/Cell.java

@@ -0,0 +1,21 @@
+package tech.pegasys.teku.kzg;
+
+import ethereum.ckzg4844.CKZG4844JNI;
+import org.apache.tuweni.bytes.Bytes;
+
+import java.util.List;
+import java.util.stream.IntStream;
+
+import static ethereum.ckzg4844.CKZG4844JNI.BYTES_PER_CELL;
+
+public record Cell(Bytes bytes) {
+
+  static Cell ZERO = new Cell(Bytes.wrap(new byte[BYTES_PER_CELL]));
+
+  static List<Cell> splitBytes(Bytes bytes) {
+    return CKZG4844Utils.bytesChunked(bytes, BYTES_PER_CELL)
+        .stream()
+        .map(Cell::new)
+        .toList();
+  }
+}

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/CellAndProof.java

@@ -0,0 +1,7 @@
+package tech.pegasys.teku.kzg;
+
+public record CellAndProof(
+    Cell cell,
+    KZGProof proof
+) {
+}

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/CellID.java

@@ -0,0 +1,14 @@
+package tech.pegasys.teku.kzg;
+
+import tech.pegasys.teku.infrastructure.unsigned.UInt64;
+
+public record CellID(UInt64 id) {
+
+  static CellID fromCellColumnIndex(int idx) {
+    return new CellID(UInt64.valueOf(idx));
+  }
+
+  int getColumnIndex() {
+    return id.intValue();
+  }
+}

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/CellWithID.java

@@ -0,0 +1,11 @@
+package tech.pegasys.teku.kzg;
+
+public record CellWithID(
+    Cell cell,
+    CellID id
+) {
+
+  static CellWithID fromCellAndColumn(Cell cell, int index) {
+    return new CellWithID(cell, CellID.fromCellColumnIndex(index));
+  }
+}

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZG.java

@@ -14,9 +14,13 @@
 package tech.pegasys.teku.kzg;
 
 import java.util.List;
+import java.util.stream.Stream;
+
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes48;
 
+import static ethereum.ckzg4844.CKZG4844JNI.CELLS_PER_BLOB;
+
 /**
  * This interface specifies all the KZG functions needed for the Deneb specification and is the
  * entry-point for all KZG operations in Teku.
@@ -31,10 +35,12 @@ static KZG getInstance() {
       new KZG() {
 
         @Override
-        public void loadTrustedSetup(final String trustedSetupFile) throws KZGException {}
+        public void loadTrustedSetup(final String trustedSetupFile) throws KZGException {
+        }
 
         @Override
-        public void freeTrustedSetup() throws KZGException {}
+        public void freeTrustedSetup() throws KZGException {
+        }
 
         @Override
         public boolean verifyBlobKzgProof(
@@ -62,6 +68,38 @@ public KZGProof computeBlobKzgProof(final Bytes blob, final KZGCommitment kzgCom
             throws KZGException {
           return KZGProof.fromBytesCompressed(Bytes48.ZERO);
         }
+
+        @Override
+        public List<Cell> computeCells(Bytes blob) {
+          List<Cell> blobCells = Cell.splitBytes(blob);
+          return Stream.concat(
+              blobCells.stream(),
+              Stream.generate(() -> Cell.ZERO).limit(blobCells.size())
+          ).toList();
+        }
+
+        @Override
+        public List<CellAndProof> computeCellsAndProofs(Bytes blob) {
+          return computeCells(blob)
+              .stream()
+              .map(cell -> new CellAndProof(cell, KZGProof.fromBytesCompressed(Bytes48.ZERO)))
+              .toList();
+        }
+
+        @Override
+        public boolean verifyCellProof(KZGCommitment commitment, CellWithID cellWithID, KZGProof proof) {
+          return true;
+        }
+
+        @Override
+        public List<Cell> recoverCells(List<CellWithID> cells) {
+          if (cells.size() < CELLS_PER_BLOB)
+            throw new IllegalArgumentException("Can't recover from " + cells.size() + " cells");
+          return cells.stream()
+              .map(CellWithID::cell)
+              .limit(CELLS_PER_BLOB)
+              .toList();
+        }
       };
 
   void loadTrustedSetup(String trustedSetupFile) throws KZGException;
@@ -78,4 +116,16 @@ boolean verifyBlobKzgProofBatch(
   KZGCommitment blobToKzgCommitment(Bytes blob) throws KZGException;
 
   KZGProof computeBlobKzgProof(Bytes blob, KZGCommitment kzgCommitment) throws KZGException;
+
+  // EIP-7594 methods
+
+  List<Cell> computeCells(Bytes blob);
+
+  List<CellAndProof> computeCellsAndProofs(Bytes blob);
+
+  boolean verifyCellProof(KZGCommitment commitment, CellWithID cellWithID, KZGProof proof);
+
+  // TODO veryCellProofBatch()
+
+  List<Cell> recoverCells(List<CellWithID> cells);
 }

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGException.java

@@ -18,4 +18,8 @@ public class KZGException extends RuntimeException {
   public KZGException(final String message, final Throwable cause) {
     super(message, cause);
   }
+
+  public KZGException(final String message) {
+    super(message);
+  }
 }

infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/KZGProof.java

@@ -14,8 +14,12 @@
 package tech.pegasys.teku.kzg;
 
 import static com.google.common.base.Preconditions.checkArgument;
+import static ethereum.ckzg4844.CKZG4844JNI.BYTES_PER_CELL;
+import static ethereum.ckzg4844.CKZG4844JNI.BYTES_PER_PROOF;
 
 import ethereum.ckzg4844.CKZG4844JNI;
+
+import java.util.List;
 import java.util.Objects;
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes48;
@@ -29,12 +33,12 @@ public static KZGProof fromHexString(final String hexString) {
 
   public static KZGProof fromSSZBytes(final Bytes bytes) {
     checkArgument(
-        bytes.size() == CKZG4844JNI.BYTES_PER_PROOF,
-        "Expected " + CKZG4844JNI.BYTES_PER_PROOF + " bytes but received %s.",
+        bytes.size() == BYTES_PER_PROOF,
+        "Expected " + BYTES_PER_PROOF + " bytes but received %s.",
         bytes.size());
     return SSZ.decode(
         bytes,
-        reader -> new KZGProof(Bytes48.wrap(reader.readFixedBytes(CKZG4844JNI.BYTES_PER_PROOF))));
+        reader -> new KZGProof(Bytes48.wrap(reader.readFixedBytes(BYTES_PER_PROOF))));
   }
 
   public static KZGProof fromBytesCompressed(final Bytes48 bytes) throws IllegalArgumentException {
@@ -45,6 +49,13 @@ public static KZGProof fromArray(final byte[] bytes) {
     return fromBytesCompressed(Bytes48.wrap(bytes));
   }
 
+  static List<KZGProof> splitBytes(Bytes bytes) {
+    return CKZG4844Utils.bytesChunked(bytes, BYTES_PER_PROOF)
+        .stream()
+        .map(b -> new KZGProof(Bytes48.wrap(b)))
+        .toList();
+  }
+
   private final Bytes48 bytesCompressed;
 
   public KZGProof(final Bytes48 bytesCompressed) {