From e9d07261d34d7da104d047451fbd2b5c7e17d283 Mon Sep 17 00:00:00 2001
From: Stefan Bratanov <stefan.bratanov93@gmail.com>
Date: Fri, 20 Oct 2023 17:01:06 +0100
Subject: [PATCH] better way

---
 .../phase0/forkchoice/ForkChoiceTestExecutor.java | 15 +++++++++++++++
 .../versions/deneb/helpers/MiscHelpersDeneb.java  |  4 ++++
 .../tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java  |  9 ---------
 3 files changed, 19 insertions(+), 9 deletions(-)

diff --git a/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java b/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java
index a2fed8b2ddc..ffaf775700f 100644
--- a/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java
+++ b/eth-reference-tests/src/referenceTest/java/tech/pegasys/teku/reference/phase0/forkchoice/ForkChoiceTestExecutor.java
@@ -42,10 +42,13 @@
 import tech.pegasys.teku.infrastructure.async.SafeFuture;
 import tech.pegasys.teku.infrastructure.async.eventthread.InlineEventThread;
 import tech.pegasys.teku.infrastructure.unsigned.UInt64;
+import tech.pegasys.teku.kzg.KZG;
 import tech.pegasys.teku.kzg.KZGProof;
 import tech.pegasys.teku.reference.TestDataUtils;
 import tech.pegasys.teku.reference.TestExecutor;
 import tech.pegasys.teku.spec.Spec;
+import tech.pegasys.teku.spec.SpecMilestone;
+import tech.pegasys.teku.spec.SpecVersion;
 import tech.pegasys.teku.spec.datastructures.attestation.ValidatableAttestation;
 import tech.pegasys.teku.spec.datastructures.blobs.versions.deneb.Blob;
 import tech.pegasys.teku.spec.datastructures.blocks.BeaconBlock;
@@ -63,7 +66,9 @@
 import tech.pegasys.teku.spec.executionlayer.ExecutionLayerChannelStub;
 import tech.pegasys.teku.spec.executionlayer.ExecutionPayloadStatus;
 import tech.pegasys.teku.spec.executionlayer.PayloadStatus;
+import tech.pegasys.teku.spec.logic.common.helpers.MiscHelpers;
 import tech.pegasys.teku.spec.logic.common.statetransition.results.BlockImportResult;
+import tech.pegasys.teku.spec.logic.versions.deneb.helpers.MiscHelpersDeneb;
 import tech.pegasys.teku.statetransition.forkchoice.ForkChoice;
 import tech.pegasys.teku.statetransition.forkchoice.ForkChoiceStateProvider;
 import tech.pegasys.teku.statetransition.forkchoice.MergeTransitionBlockValidator;
@@ -156,6 +161,8 @@ spec, new SignedBlockAndState(anchorBlock, anchorState)),
               + "\nProtoarray data:\n"
               + protoArrayData,
           e);
+    } finally {
+      freeTrustedSetupIfRequired(spec);
     }
   }
 
@@ -482,6 +489,14 @@ private void assertCheckpoint(
         .isEqualTo(new Checkpoint(expectedEpoch, expectedRoot));
   }
 
+  private void freeTrustedSetupIfRequired(final Spec spec) {
+    Optional.ofNullable(spec.forMilestone(SpecMilestone.DENEB))
+        .map(SpecVersion::miscHelpers)
+        .flatMap(MiscHelpers::toVersionDeneb)
+        .map(MiscHelpersDeneb::getKzg)
+        .ifPresent(KZG::freeTrustedSetup);
+  }
+
   @SuppressWarnings({"unchecked", "TypeParameterUnusedInFormals"})
   private static <T> T get(final Map<String, Object> yamlData, final String key) {
     return (T) yamlData.get(key);
diff --git a/ethereum/spec/src/main/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/MiscHelpersDeneb.java b/ethereum/spec/src/main/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/MiscHelpersDeneb.java
index ffda11cd491..7af5a85cd83 100644
--- a/ethereum/spec/src/main/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/MiscHelpersDeneb.java
+++ b/ethereum/spec/src/main/java/tech/pegasys/teku/spec/logic/versions/deneb/helpers/MiscHelpersDeneb.java
@@ -60,6 +60,10 @@ private static KZG initKZG(final SpecConfigDeneb config) {
     return kzg;
   }
 
+  public KZG getKzg() {
+    return kzg;
+  }
+
   /**
    * Performs complete data availability check <a
    * href="https://github.com/ethereum/consensus-specs/blob/dev/specs/deneb/fork-choice.md#is_data_available">is_data_available</a>
diff --git a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java
index e9e0b748ea7..10627580348 100644
--- a/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java
+++ b/infrastructure/kzg/src/main/java/tech/pegasys/teku/kzg/ckzg4844/CKZG4844.java
@@ -63,15 +63,6 @@ public synchronized void loadTrustedSetup(final String trustedSetupFile) throws
     }
 
     try {
-      // allow loading different trusted setup
-      loadedTrustedSetupFile.ifPresent(
-          currentTrustedSetupFile -> {
-            LOG.debug(
-                "Replacing current trusted setup {} with {}",
-                currentTrustedSetupFile,
-                trustedSetupFile);
-            freeTrustedSetup();
-          });
       final TrustedSetup trustedSetup = CKZG4844Utils.parseTrustedSetupFile(trustedSetupFile);
       final List<Bytes> g1Points = trustedSetup.g1Points();
       final List<Bytes> g2Points = trustedSetup.g2Points();