Bonsai based reference test worldstate

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiAccount.java

@@ -219,7 +219,7 @@ public UInt256 getOriginalStorageValue(final UInt256 key) {
   @Override
   public NavigableMap<Bytes32, AccountStorageEntry> storageEntriesFrom(
       final Bytes32 startKeyHash, final int limit) {
-    throw new RuntimeException("Bonsai Tries does not currently support enumerating storage");
+    return context.getWorldStateStorage().storageEntriesFrom(this.addressHash, startKeyHash, limit);
   }
 
   public Bytes serializeAccount() {

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/BonsaiWorldStateProvider.java

@@ -102,8 +102,8 @@ public BonsaiWorldStateProvider(
             pluginContext);
     this.blockchain = blockchain;
     this.worldStateStorage = worldStateStorage;
-    this.persistedState = new BonsaiWorldState(this, worldStateStorage);
     this.cachedMerkleTrieLoader = cachedMerkleTrieLoader;
+    this.persistedState = new BonsaiWorldState(this, worldStateStorage);
     blockchain
         .getBlockHeader(persistedState.getWorldStateBlockHash())
         .ifPresent(

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/storage/BonsaiPreImageProxy.java

@@ -0,0 +1,68 @@
+/*
+ * Copyright Hyperledger Besu Contributors.
+ *
+ * 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.bonsai.storage;
+
+import org.hyperledger.besu.datatypes.Address;
+import org.hyperledger.besu.datatypes.Hash;
+import org.hyperledger.besu.ethereum.worldstate.WorldStatePreimageStorage;
+
+import java.util.Optional;
+
+import com.google.common.collect.BiMap;
+import com.google.common.collect.HashBiMap;
+import org.apache.tuweni.bytes.Bytes;
+import org.apache.tuweni.bytes.Bytes32;
+import org.apache.tuweni.units.bigints.UInt256;
+
+/** Acts as both a Hasher and PreImageStorage for Bonsai storage format. */
+public interface BonsaiPreImageProxy extends WorldStatePreimageStorage {
+  /**
+   * If this value is not already present, save in preImage store and return the hash value.
+   *
+   * @param value value to hash
+   * @return Hash of value
+   */
+  Hash hashAndSavePreImage(Bytes value);
+
+  /**
+   * A caching PreImageProxy suitable for ReferenceTestWorldState which saves hashes in an unbounded
+   * BiMap.
+   */
+  class BonsaiReferenceTestPreImageProxy implements BonsaiPreImageProxy {
+    BiMap<Hash, Bytes> preImageCache = HashBiMap.create();
+
+    @Override
+    public synchronized Hash hashAndSavePreImage(final Bytes value) {
+      return preImageCache.inverse().computeIfAbsent(value, Hash::hash);
+    }
+
+    @Override
+    public Optional<UInt256> getStorageTrieKeyPreimage(final Bytes32 trieKey) {
+      return Optional.ofNullable(preImageCache.get(trieKey)).map(UInt256::fromBytes);
+    }
+
+    @Override
+    public Optional<Address> getAccountTrieKeyPreimage(final Bytes32 trieKey) {
+      return Optional.ofNullable(preImageCache.get(trieKey)).map(Address::wrap);
+    }
+
+    @Override
+    public Updater updater() {
+      throw new UnsupportedOperationException(
+          "BonsaiReferenceTestPreImageProxy does not implement an updater");
+    }
+  }
+}

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/storage/BonsaiWorldStateKeyValueStorage.java

@@ -31,6 +31,7 @@
 import org.hyperledger.besu.ethereum.worldstate.FlatDbMode;
 import org.hyperledger.besu.ethereum.worldstate.StateTrieAccountValue;
 import org.hyperledger.besu.ethereum.worldstate.WorldStateStorage;
+import org.hyperledger.besu.evm.account.AccountStorageEntry;
 import org.hyperledger.besu.metrics.ObservableMetricsSystem;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorage;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorageTransaction;
@@ -41,6 +42,7 @@
 import java.nio.charset.StandardCharsets;
 import java.util.List;
 import java.util.Map;
+import java.util.NavigableMap;
 import java.util.Optional;
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.Predicate;
@@ -53,7 +55,6 @@
 
 @SuppressWarnings("unused")
 public class BonsaiWorldStateKeyValueStorage implements WorldStateStorage, AutoCloseable {
-
   private static final Logger LOG = LoggerFactory.getLogger(BonsaiWorldStateKeyValueStorage.class);
 
   // 0x776f726c64526f6f74
@@ -250,6 +251,11 @@ public Map<Bytes32, Bytes> streamFlatStorages(
             composedWorldStateStorage, accountHash, startKeyHash, endKeyHash, max);
   }
 
+  public NavigableMap<Bytes32, AccountStorageEntry> storageEntriesFrom(
+      final Hash addressHash, final Bytes32 startKeyHash, final int limit) {
+    throw new RuntimeException("Bonsai Tries does not currently support enumerating storage");
+  }
+
   @Override
   public Optional<Bytes> getNodeData(final Bytes location, final Bytes32 hash) {
     return Optional.empty();

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/storage/flat/FlatDbReaderStrategy.java

@@ -135,10 +135,10 @@ public Map<Bytes32, Bytes> streamAccountFlatDatabase(
       final long max) {
     final Stream<Pair<Bytes32, Bytes>> pairStream =
         storage
-            .streamFromKey(ACCOUNT_INFO_STATE, startKeyHash.toArrayUnsafe())
+            .streamFromKey(
+                ACCOUNT_INFO_STATE, startKeyHash.toArrayUnsafe(), endKeyHash.toArrayUnsafe())
             .limit(max)
-            .map(pair -> new Pair<>(Bytes32.wrap(pair.getKey()), Bytes.wrap(pair.getValue())))
-            .takeWhile(pair -> pair.getFirst().compareTo(endKeyHash) <= 0);
+            .map(pair -> new Pair<>(Bytes32.wrap(pair.getKey()), Bytes.wrap(pair.getValue())));
 
     final TreeMap<Bytes32, Bytes> collected =
         pairStream.collect(
@@ -157,15 +157,14 @@ public Map<Bytes32, Bytes> streamStorageFlatDatabase(
         storage
             .streamFromKey(
                 ACCOUNT_STORAGE_STORAGE,
-                Bytes.concatenate(accountHash, startKeyHash).toArrayUnsafe())
-            .takeWhile(pair -> Bytes.wrap(pair.getKey()).slice(0, Hash.SIZE).equals(accountHash))
+                Bytes.concatenate(accountHash, startKeyHash).toArrayUnsafe(),
+                Bytes.concatenate(accountHash, endKeyHash).toArrayUnsafe())
             .limit(max)
             .map(
                 pair ->
                     new Pair<>(
                         Bytes32.wrap(Bytes.wrap(pair.getKey()).slice(Hash.SIZE)),
-                        RLP.encodeValue(Bytes.wrap(pair.getValue()).trimLeadingZeros())))
-            .takeWhile(pair -> pair.getFirst().compareTo(endKeyHash) <= 0);
+                        RLP.encodeValue(Bytes.wrap(pair.getValue()).trimLeadingZeros())));
 
     final TreeMap<Bytes32, Bytes> collected =
         pairStream.collect(

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/bonsai/worldview/BonsaiWorldState.java

@@ -27,6 +27,7 @@
 import org.hyperledger.besu.ethereum.bonsai.BonsaiAccount;
 import org.hyperledger.besu.ethereum.bonsai.BonsaiValue;
 import org.hyperledger.besu.ethereum.bonsai.BonsaiWorldStateProvider;
+import org.hyperledger.besu.ethereum.bonsai.cache.CachedMerkleTrieLoader;
 import org.hyperledger.besu.ethereum.bonsai.storage.BonsaiSnapshotWorldStateKeyValueStorage;
 import org.hyperledger.besu.ethereum.bonsai.storage.BonsaiWorldStateKeyValueStorage;
 import org.hyperledger.besu.ethereum.bonsai.storage.BonsaiWorldStateKeyValueStorage.BonsaiStorageSubscriber;
@@ -65,51 +66,52 @@ public class BonsaiWorldState
 
   private static final Logger LOG = LoggerFactory.getLogger(BonsaiWorldState.class);
 
-  private BonsaiWorldStateKeyValueStorage worldStateStorage;
+  protected BonsaiWorldStateKeyValueStorage worldStateStorage;
 
-  private final BonsaiWorldStateProvider archive;
-  private final BonsaiWorldStateUpdateAccumulator accumulator;
+  protected final CachedMerkleTrieLoader cachedMerkleTrieLoader;
+  protected final TrieLogManager trieLogManager;
+  private BonsaiWorldStateUpdateAccumulator accumulator;
 
-  private Hash worldStateRootHash;
+  protected Hash worldStateRootHash;
   Hash worldStateBlockHash;
-
   private boolean isFrozen;
 
   public BonsaiWorldState(
       final BonsaiWorldStateProvider archive,
       final BonsaiWorldStateKeyValueStorage worldStateStorage) {
-    this.archive = archive;
+    this(worldStateStorage, archive.getCachedMerkleTrieLoader(), archive.getTrieLogManager());
+  }
+
+  protected BonsaiWorldState(
+      final BonsaiWorldStateKeyValueStorage worldStateStorage,
+      final CachedMerkleTrieLoader cachedMerkleTrieLoader,
+      final TrieLogManager trieLogManager) {
     this.worldStateStorage = worldStateStorage;
-    worldStateRootHash =
+    this.worldStateRootHash =
         Hash.wrap(
             Bytes32.wrap(worldStateStorage.getWorldStateRootHash().orElse(Hash.EMPTY_TRIE_HASH)));
-    worldStateBlockHash =
+    this.worldStateBlockHash =
         Hash.wrap(Bytes32.wrap(worldStateStorage.getWorldStateBlockHash().orElse(Hash.ZERO)));
-    accumulator =
+    this.accumulator =
         new BonsaiWorldStateUpdateAccumulator(
             this,
             (addr, value) ->
-                archive
-                    .getCachedMerkleTrieLoader()
-                    .preLoadAccount(getWorldStateStorage(), worldStateRootHash, addr),
+                cachedMerkleTrieLoader.preLoadAccount(
+                    getWorldStateStorage(), worldStateRootHash, addr),
             (addr, value) ->
-                archive
-                    .getCachedMerkleTrieLoader()
-                    .preLoadStorageSlot(getWorldStateStorage(), addr, value));
+                cachedMerkleTrieLoader.preLoadStorageSlot(getWorldStateStorage(), addr, value));
+    this.cachedMerkleTrieLoader = cachedMerkleTrieLoader;
+    this.trieLogManager = trieLogManager;
   }
 
-  public BonsaiWorldState(
-      final BonsaiWorldStateProvider archive,
-      final BonsaiWorldStateKeyValueStorage worldStateStorage,
-      final BonsaiWorldStateUpdateAccumulator updater) {
-    this.archive = archive;
-    this.worldStateStorage = worldStateStorage;
-    this.worldStateRootHash =
-        Hash.wrap(
-            Bytes32.wrap(worldStateStorage.getWorldStateRootHash().orElse(Hash.EMPTY_TRIE_HASH)));
-    this.worldStateBlockHash =
-        Hash.wrap(Bytes32.wrap(worldStateStorage.getWorldStateBlockHash().orElse(Hash.ZERO)));
-    this.accumulator = updater;
+  /**
+   * Having a protected method to override the accumulator solves the chicken-egg problem of needing
+   * a worldstate reference (this) when construction the Accumulator.
+   *
+   * @param accumulator accumulator to use.
+   */
+  protected void setAccumulator(final BonsaiWorldStateUpdateAccumulator accumulator) {
+    this.accumulator = accumulator;
   }
 
   /**
@@ -130,10 +132,6 @@ public Hash getWorldStateRootHash() {
     return worldStateRootHash;
   }
 
-  public BonsaiWorldStateProvider getArchive() {
-    return archive;
-  }
-
   @Override
   public boolean isPersisted() {
     return isPersisted(worldStateStorage);
@@ -189,9 +187,7 @@ private Hash calculateRootHash(
     final StoredMerklePatriciaTrie<Bytes, Bytes> accountTrie =
         createTrie(
             (location, hash) ->
-                archive
-                    .getCachedMerkleTrieLoader()
-                    .getAccountStateTrieNode(worldStateStorage, location, hash),
+                cachedMerkleTrieLoader.getAccountStateTrieNode(worldStateStorage, location, hash),
             worldStateRootHash);
 
     // for manicured tries and composting, collect branches here (not implemented)
@@ -223,7 +219,7 @@ private void updateTheAccounts(
       final BonsaiAccount updatedAccount = bonsaiValue.getUpdated();
       try {
         if (updatedAccount == null) {
-          final Hash addressHash = Hash.hash(accountKey);
+          final Hash addressHash = hashAndSavePreImage(accountKey);
           accountTrie.remove(addressHash);
           maybeStateUpdater.ifPresent(
               bonsaiUpdater -> bonsaiUpdater.removeAccountInfoState(addressHash));
@@ -232,7 +228,7 @@ private void updateTheAccounts(
           final Bytes accountValue = updatedAccount.serializeAccount();
           maybeStateUpdater.ifPresent(
               bonsaiUpdater ->
-                  bonsaiUpdater.putAccountInfoState(Hash.hash(accountKey), accountValue));
+                  bonsaiUpdater.putAccountInfoState(hashAndSavePreImage(accountKey), accountValue));
           accountTrie.put(addressHash, accountValue);
         }
       } catch (MerkleTrieException e) {
@@ -277,10 +273,8 @@ private void updateAccountStorageState(
       final StoredMerklePatriciaTrie<Bytes, Bytes> storageTrie =
           createTrie(
               (location, key) ->
-                  archive
-                      .getCachedMerkleTrieLoader()
-                      .getAccountStorageTrieNode(
-                          worldStateStorage, updatedAddressHash, location, key),
+                  cachedMerkleTrieLoader.getAccountStorageTrieNode(
+                      worldStateStorage, updatedAddressHash, location, key),
               storageRoot);
 
       // for manicured tries and composting, collect branches here (not implemented)
@@ -405,7 +399,6 @@ public void persist(final BlockHeader blockHeader) {
         }
         saveTrieLog =
             () -> {
-              final TrieLogManager trieLogManager = archive.getTrieLogManager();
               trieLogManager.saveTrieLog(localCopy, newWorldStateRootHash, blockHeader, this);
               // not save a frozen state in the cache
               if (!isFrozen) {
@@ -626,4 +619,9 @@ private void closeFrozenStorage() {
       // no op
     }
   }
+
+  protected Hash hashAndSavePreImage(final Bytes value) {
+    // by default do not save has preImages
+    return Hash.hash(value);
+  }
 }