[VERKLE] Preload stems

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/trie/diffbased/bonsai/worldview/BonsaiWorldStateUpdateAccumulator.java

@@ -40,7 +40,14 @@ public BonsaiWorldStateUpdateAccumulator(
       final Consumer<DiffBasedValue<BonsaiAccount>> accountPreloader,
       final Consumer<StorageSlotKey> storagePreloader,
       final EvmConfiguration evmConfiguration) {
-    super(world, accountPreloader, storagePreloader, evmConfiguration);
+    super(
+        world,
+        accountPreloader,
+        storagePreloader,
+        (__, ___) -> {
+          /*nothing to preload for the code*/
+        },
+        evmConfiguration);
   }
 
   @Override

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/trie/diffbased/common/worldview/accumulator/DiffBasedWorldStateUpdateAccumulator.java

@@ -27,6 +27,7 @@
 import org.hyperledger.besu.ethereum.trie.diffbased.common.worldview.DiffBasedWorldState;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.worldview.DiffBasedWorldView;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.worldview.accumulator.preload.AccountConsumingMap;
+import org.hyperledger.besu.ethereum.trie.diffbased.common.worldview.accumulator.preload.CodeConsumingMap;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.worldview.accumulator.preload.Consumer;
 import org.hyperledger.besu.ethereum.trie.diffbased.common.worldview.accumulator.preload.StorageConsumingMap;
 import org.hyperledger.besu.evm.account.Account;
@@ -62,9 +63,10 @@ public abstract class DiffBasedWorldStateUpdateAccumulator<ACCOUNT extends DiffB
       LoggerFactory.getLogger(DiffBasedWorldStateUpdateAccumulator.class);
   protected final Consumer<DiffBasedValue<ACCOUNT>> accountPreloader;
   protected final Consumer<StorageSlotKey> storagePreloader;
+  private final Consumer<Bytes> codePreloader;
 
   private final AccountConsumingMap<DiffBasedValue<ACCOUNT>> accountsToUpdate;
-  private final Map<Address, DiffBasedValue<Bytes>> codeToUpdate = new ConcurrentHashMap<>();
+  private final CodeConsumingMap codeToUpdate;
   private final Set<Address> storageToClear = Collections.synchronizedSet(new HashSet<>());
   protected final EvmConfiguration evmConfiguration;
 
@@ -81,11 +83,14 @@ public DiffBasedWorldStateUpdateAccumulator(
       final DiffBasedWorldView world,
       final Consumer<DiffBasedValue<ACCOUNT>> accountPreloader,
       final Consumer<StorageSlotKey> storagePreloader,
+      final Consumer<Bytes> codePreloader,
       final EvmConfiguration evmConfiguration) {
     super(world, evmConfiguration);
     this.accountsToUpdate = new AccountConsumingMap<>(new ConcurrentHashMap<>(), accountPreloader);
+    this.codeToUpdate = new CodeConsumingMap(new ConcurrentHashMap<>(), codePreloader);
     this.accountPreloader = accountPreloader;
     this.storagePreloader = storagePreloader;
+    this.codePreloader = codePreloader;
     this.isAccumulatorStateChanged = false;
     this.evmConfiguration = evmConfiguration;
   }
@@ -108,6 +113,10 @@ protected Consumer<StorageSlotKey> getStoragePreloader() {
     return storagePreloader;
   }
 
+  public Consumer<Bytes> getCodePreloader() {
+    return codePreloader;
+  }
+
   protected EvmConfiguration getEvmConfiguration() {
     return evmConfiguration;
   }

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/trie/diffbased/common/worldview/accumulator/preload/CodeConsumingMap.java

@@ -0,0 +1,53 @@
+/*
+ * Copyright contributors to Hyperledger Besu.
+ *
+ * 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.
+ *
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package org.hyperledger.besu.ethereum.trie.diffbased.common.worldview.accumulator.preload;
+
+import org.hyperledger.besu.datatypes.Address;
+import org.hyperledger.besu.ethereum.trie.diffbased.common.DiffBasedValue;
+
+import java.util.Map;
+import java.util.concurrent.ConcurrentMap;
+import javax.annotation.Nonnull;
+
+import com.google.common.collect.ForwardingMap;
+import org.apache.tuweni.bytes.Bytes;
+
+public class CodeConsumingMap extends ForwardingMap<Address, DiffBasedValue<Bytes>> {
+
+  private final ConcurrentMap<Address, DiffBasedValue<Bytes>> codes;
+  private final Consumer<Bytes> consumer;
+
+  public CodeConsumingMap(
+      final ConcurrentMap<Address, DiffBasedValue<Bytes>> codes, final Consumer<Bytes> consumer) {
+    this.codes = codes;
+    this.consumer = consumer;
+  }
+
+  @Override
+  public DiffBasedValue<Bytes> put(
+      @Nonnull final Address address, @Nonnull final DiffBasedValue<Bytes> value) {
+    consumer.process(address, value.getUpdated() != null ? value.getUpdated() : value.getPrior());
+    return codes.put(address, value);
+  }
+
+  public Consumer<Bytes> getConsumer() {
+    return consumer;
+  }
+
+  @Override
+  protected Map<Address, DiffBasedValue<Bytes>> delegate() {
+    return codes;
+  }
+}