KZG updates from das branch

eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/KzgRetriever.java

@@ -25,10 +25,11 @@ public class KzgRetriever {
   private static final Map<String, String> TRUSTED_SETUP_FILES_BY_NETWORK = Maps.newHashMap();
 
   public static KZG getKzgWithLoadedTrustedSetup(final Spec spec, final String network) {
-    if (!spec.isMilestoneSupported(SpecMilestone.DENEB)) {
-      return KZG.NOOP;
+    if (spec.isMilestoneSupported(SpecMilestone.DENEB)
+        || spec.isMilestoneSupported(SpecMilestone.ELECTRA)) {
+      return getKzgWithLoadedTrustedSetup(network);
     }
-    return getKzgWithLoadedTrustedSetup(network);
+    return KZG.NOOP;
   }
 
   public static KZG getKzgWithLoadedTrustedSetup(final String network) {
@@ -43,7 +44,7 @@ public static KZG getKzgWithLoadedTrustedSetup(final String network) {
                         () ->
                             new IllegalArgumentException(
                                 "No trusted setup configured for " + network)));
-    final KZG kzg = KZG.getInstance();
+    final KZG kzg = KZG.getInstance(false);
     kzg.loadTrustedSetup(trustedSetupFile);
     return kzg;
   }

ethereum/networks/src/main/java/tech/pegasys/teku/networks/Eth2NetworkConfiguration.java

@@ -71,6 +71,8 @@ public class Eth2NetworkConfiguration {
 
   public static final int DEFAULT_ASYNC_P2P_MAX_QUEUE = DEFAULT_MAX_QUEUE_SIZE;
 
+  public static final boolean DEFAULT_RUST_KZG_ENABLED = false;
+
   // at least 5, but happily up to 12
   public static final int DEFAULT_VALIDATOR_EXECUTOR_THREADS =
       Math.max(5, Math.min(Runtime.getRuntime().availableProcessors(), 12));
@@ -114,6 +116,7 @@ public class Eth2NetworkConfiguration {
   private final boolean forkChoiceLateBlockReorgEnabled;
   private final boolean forkChoiceUpdatedAlwaysSendPayloadAttributes;
   private final int pendingAttestationsMaxQueue;
+  private final boolean rustKzgEnabled;
 
   private Eth2NetworkConfiguration(
       final Spec spec,
@@ -141,7 +144,8 @@ private Eth2NetworkConfiguration(
       final int asyncBeaconChainMaxQueue,
       final boolean forkChoiceLateBlockReorgEnabled,
       final boolean forkChoiceUpdatedAlwaysSendPayloadAttributes,
-      final int pendingAttestationsMaxQueue) {
+      final int pendingAttestationsMaxQueue,
+      final boolean rustKzgEnabled) {
     this.spec = spec;
     this.constants = constants;
     this.stateBoostrapConfig = stateBoostrapConfig;
@@ -172,6 +176,7 @@ private Eth2NetworkConfiguration(
     this.forkChoiceUpdatedAlwaysSendPayloadAttributes =
         forkChoiceUpdatedAlwaysSendPayloadAttributes;
     this.pendingAttestationsMaxQueue = pendingAttestationsMaxQueue;
+    this.rustKzgEnabled = rustKzgEnabled;
 
     LOG.debug(
         "P2P async queue - {} threads, max queue size {} ", asyncP2pMaxThreads, asyncP2pMaxQueue);
@@ -293,6 +298,10 @@ public boolean isForkChoiceUpdatedAlwaysSendPayloadAttributes() {
     return forkChoiceUpdatedAlwaysSendPayloadAttributes;
   }
 
+  public boolean isRustKzgEnabled() {
+    return rustKzgEnabled;
+  }
+
   @Override
   public String toString() {
     return constants;
@@ -316,6 +325,7 @@ public boolean equals(final Object o) {
         && forkChoiceLateBlockReorgEnabled == that.forkChoiceLateBlockReorgEnabled
         && forkChoiceUpdatedAlwaysSendPayloadAttributes
             == that.forkChoiceUpdatedAlwaysSendPayloadAttributes
+        && rustKzgEnabled == that.rustKzgEnabled
         && Objects.equals(spec, that.spec)
         && Objects.equals(constants, that.constants)
         && Objects.equals(stateBoostrapConfig, that.stateBoostrapConfig)
@@ -362,7 +372,8 @@ public int hashCode() {
         asyncBeaconChainMaxQueue,
         asyncP2pMaxQueue,
         forkChoiceLateBlockReorgEnabled,
-        forkChoiceUpdatedAlwaysSendPayloadAttributes);
+        forkChoiceUpdatedAlwaysSendPayloadAttributes,
+        rustKzgEnabled);
   }
 
   public static class Builder {
@@ -400,6 +411,7 @@ public static class Builder {
     private boolean forkChoiceUpdatedAlwaysSendPayloadAttributes =
         DEFAULT_FORK_CHOICE_UPDATED_ALWAYS_SEND_PAYLOAD_ATTRIBUTES;
     private OptionalInt pendingAttestationsMaxQueue = OptionalInt.empty();
+    private boolean rustKzgEnabled = DEFAULT_RUST_KZG_ENABLED;
 
     public void spec(final Spec spec) {
       this.spec = spec;
@@ -498,7 +510,8 @@ public Eth2NetworkConfiguration build() {
           asyncBeaconChainMaxQueue.orElse(DEFAULT_ASYNC_BEACON_CHAIN_MAX_QUEUE),
           forkChoiceLateBlockReorgEnabled,
           forkChoiceUpdatedAlwaysSendPayloadAttributes,
-          pendingAttestationsMaxQueue.orElse(DEFAULT_MAX_QUEUE_PENDING_ATTESTATIONS));
+          pendingAttestationsMaxQueue.orElse(DEFAULT_MAX_QUEUE_PENDING_ATTESTATIONS),
+          rustKzgEnabled);
     }
 
     private void validateCommandLineParameters() {
@@ -737,6 +750,11 @@ public Builder epochsStoreBlobs(final String epochsStoreBlobs) {
       return this;
     }
 
+    public Builder rustKzgEnabled(final boolean rustKzgEnabled) {
+      this.rustKzgEnabled = rustKzgEnabled;
+      return this;
+    }
+
     public Builder applyNetworkDefaults(final String networkName) {
       Eth2Network.fromStringLenient(networkName)
           .ifPresentOrElse(

ethereum/spec/src/main/java/tech/pegasys/teku/spec/executionlayer/ExecutionLayerChannelStub.java

@@ -124,7 +124,7 @@ public ExecutionLayerChannelStub(
     final KZG kzg;
     if (spec.isMilestoneSupported(SpecMilestone.DENEB)) {
       // trusted setup loading will be handled by the BeaconChainController
-      kzg = KZG.getInstance();
+      kzg = KZG.getInstance(false);
     } else {
       kzg = KZG.NOOP;
     }

ethereum/spec/src/property-test/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/KzgResolver.java

@@ -51,7 +51,7 @@ private KZG getKzgWithTrustedSetup() {
 
   private static class KzgAutoLoadFree implements Store.CloseOnReset {
 
-    private final KZG kzg = KZG.getInstance();
+    private final KZG kzg = KZG.getInstance(false);
 
     private KzgAutoLoadFree() {
       TrustedSetupLoader.loadTrustedSetupForTests(kzg);

gradle/versions.gradle

@@ -35,6 +35,7 @@ dependencyManagement {
     dependency 'io.libp2p:jvm-libp2p:1.2.2-RELEASE'
     dependency 'tech.pegasys:jblst:0.3.12'
     dependency 'io.consensys.protocols:jc-kzg-4844:2.1.1'
+    dependency 'io.github.crate-crypto:java-eth-kzg:0.5.2'
 
     dependency 'org.hdrhistogram:HdrHistogram:2.2.2'
 

infrastructure/kzg/build.gradle

@@ -5,6 +5,7 @@ dependencies {
   implementation 'io.consensys.tuweni:tuweni-bytes'
   implementation 'io.consensys.tuweni:tuweni-ssz'
   implementation 'io.consensys.protocols:jc-kzg-4844'
+  implementation "io.github.crate-crypto:java-eth-kzg"
   implementation 'commons-io:commons-io'
 
   testFixturesImplementation 'com.google.guava:guava'

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

@@ -13,9 +13,14 @@
 
 package tech.pegasys.teku.kzg;
 
+import static ethereum.ckzg4844.CKZG4844JNI.BYTES_PER_CELL;
+import static ethereum.ckzg4844.CKZG4844JNI.BYTES_PER_COMMITMENT;
+
 import ethereum.ckzg4844.CKZG4844JNI;
+import ethereum.ckzg4844.CellsAndProofs;
 import java.util.List;
 import java.util.Optional;
+import java.util.stream.IntStream;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
 import org.apache.tuweni.bytes.Bytes;
@@ -149,4 +154,58 @@ public KZGProof computeBlobKzgProof(final Bytes blob, final KZGCommitment kzgCom
           "Failed to compute KZG proof for blob with commitment " + kzgCommitment, ex);
     }
   }
+
+  @Override
+  public List<KZGCellAndProof> computeCellsAndProofs(final Bytes blob) {
+    final CellsAndProofs cellsAndProofs =
+        CKZG4844JNI.computeCellsAndKzgProofs(blob.toArrayUnsafe());
+    final List<KZGCell> cells = KZGCell.splitBytes(Bytes.wrap(cellsAndProofs.getCells()));
+    final 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 KZGCellAndProof(cells.get(i), proofs.get(i)))
+        .toList();
+  }
+
+  @Override
+  public boolean verifyCellProofBatch(
+      final List<KZGCommitment> commitments,
+      final List<KZGCellWithColumnId> cellWithIdList,
+      final List<KZGProof> proofs) {
+    if (commitments.size() != cellWithIdList.size() || cellWithIdList.size() != proofs.size()) {
+      throw new KZGException("Cells, proofs and commitments sizes should match");
+    }
+    return CKZG4844JNI.verifyCellKzgProofBatch(
+        CKZG4844Utils.flattenBytes(
+            commitments.stream()
+                .map(kzgCommitment -> (Bytes) kzgCommitment.getBytesCompressed())
+                .toList(),
+            commitments.size() * BYTES_PER_COMMITMENT),
+        cellWithIdList.stream()
+            .mapToLong(cellWithIds -> cellWithIds.columnId().id().longValue())
+            .toArray(),
+        CKZG4844Utils.flattenBytes(
+            cellWithIdList.stream().map(cellWithIds -> cellWithIds.cell().bytes()).toList(),
+            cellWithIdList.size() * BYTES_PER_CELL),
+        CKZG4844Utils.flattenProofs(proofs));
+  }
+
+  @Override
+  public List<KZGCellAndProof> recoverCellsAndProofs(final List<KZGCellWithColumnId> cells) {
+    final long[] cellIds = cells.stream().mapToLong(c -> c.columnId().id().longValue()).toArray();
+    final byte[] cellBytes =
+        CKZG4844Utils.flattenBytes(
+            cells.stream().map(c -> c.cell().bytes()).toList(), cells.size() * BYTES_PER_CELL);
+    final CellsAndProofs cellsAndProofs = CKZG4844JNI.recoverCellsAndKzgProofs(cellIds, cellBytes);
+    final List<KZGCell> fullCells = KZGCell.splitBytes(Bytes.wrap(cellsAndProofs.getCells()));
+    final List<KZGProof> fullProofs = KZGProof.splitBytes(Bytes.wrap(cellsAndProofs.getProofs()));
+    if (fullCells.size() != fullProofs.size()) {
+      throw new KZGException("Cells and proofs size differ");
+    }
+    return IntStream.range(0, fullCells.size())
+        .mapToObj(i -> new KZGCellAndProof(fullCells.get(i), fullProofs.get(i)))
+        .toList();
+  }
 }

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

@@ -27,12 +27,12 @@
 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;
 
 class CKZG4844Utils {
-
   private static final int MAX_BYTES_TO_FLATTEN = 100_663_296; // ~100.66 MB or 768 blobs
 
   public static byte[] flattenBlobs(final List<Bytes> blobs) {
@@ -59,6 +59,16 @@ public static byte[] flattenG2Points(final List<Bytes> g2Points) {
     return flattenBytes(g2Points, BYTES_PER_G2 * g2Points.size());
   }
 
+  static List<Bytes> bytesChunked(final Bytes bytes, final 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);
@@ -99,7 +109,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/KZG.java

@@ -13,6 +13,7 @@
 
 package tech.pegasys.teku.kzg;
 
+import java.math.BigInteger;
 import java.util.List;
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes48;
@@ -22,12 +23,16 @@
  * entry-point for all KZG operations in Teku.
  */
 public interface KZG {
-
+  BigInteger BLS_MODULUS =
+      new BigInteger(
+          "52435875175126190479447740508185965837690552500527637822603658699938581184513");
   int BYTES_PER_G1 = 48;
   int BYTES_PER_G2 = 96;
+  int CELLS_PER_EXT_BLOB = 128;
+  int FIELD_ELEMENTS_PER_BLOB = 4096;
 
-  static KZG getInstance() {
-    return CKZG4844.getInstance();
+  static KZG getInstance(final boolean rustKzgEnabled) {
+    return rustKzgEnabled ? RustWithCKZG.getInstance() : CKZG4844.getInstance();
   }
 
   KZG NOOP =
@@ -65,6 +70,24 @@ public KZGProof computeBlobKzgProof(final Bytes blob, final KZGCommitment kzgCom
             throws KZGException {
           return KZGProof.fromBytesCompressed(Bytes48.ZERO);
         }
+
+        @Override
+        public List<KZGCellAndProof> computeCellsAndProofs(Bytes blob) {
+          throw new RuntimeException("Not implemented");
+        }
+
+        @Override
+        public boolean verifyCellProofBatch(
+            List<KZGCommitment> commitments,
+            List<KZGCellWithColumnId> cellWithIDs,
+            List<KZGProof> proofs) {
+          return false;
+        }
+
+        @Override
+        public List<KZGCellAndProof> recoverCellsAndProofs(List<KZGCellWithColumnId> cells) {
+          throw new RuntimeException("Not implemented");
+        }
       };
 
   void loadTrustedSetup(String trustedSetupFile) throws KZGException;
@@ -81,4 +104,15 @@ boolean verifyBlobKzgProofBatch(
   KZGCommitment blobToKzgCommitment(Bytes blob) throws KZGException;
 
   KZGProof computeBlobKzgProof(Bytes blob, KZGCommitment kzgCommitment) throws KZGException;
+
+  // Fulu PeerDAS methods
+
+  List<KZGCellAndProof> computeCellsAndProofs(Bytes blob);
+
+  boolean verifyCellProofBatch(
+      List<KZGCommitment> commitments,
+      List<KZGCellWithColumnId> cellWithIDs,
+      List<KZGProof> proofs);
+
+  List<KZGCellAndProof> recoverCellsAndProofs(List<KZGCellWithColumnId> cells);
 }

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

@@ -0,0 +1,28 @@
+/*
+ * Copyright Consensys Software Inc., 2024
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+
+package tech.pegasys.teku.kzg;
+
+import static ethereum.ckzg4844.CKZG4844JNI.BYTES_PER_CELL;
+
+import java.util.List;
+import org.apache.tuweni.bytes.Bytes;
+
+public record KZGCell(Bytes bytes) {
+
+  static final KZGCell ZERO = new KZGCell(Bytes.wrap(new byte[BYTES_PER_CELL]));
+
+  static List<KZGCell> splitBytes(final Bytes bytes) {
+    return CKZG4844Utils.bytesChunked(bytes, BYTES_PER_CELL).stream().map(KZGCell::new).toList();
+  }
+}

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

@@ -0,0 +1,16 @@
+/*
+ * Copyright Consensys Software Inc., 2024
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+ * an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations under the License.
+ */
+
+package tech.pegasys.teku.kzg;
+
+public record KZGCellAndProof(KZGCell cell, KZGProof proof) {}