move hashing and preImage functions into a PreImageProxy class

[skip ci]

Signed-off-by: garyschulte <garyschulte@gmail.com>

Author: garyschulte

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

@@ -36,9 +36,6 @@
 import java.util.Map;
 import java.util.NavigableMap;
 import java.util.Objects;
-import java.util.Optional;
-import java.util.TreeMap;
-import java.util.stream.Collectors;
 
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;

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

@@ -1,2 +1,86 @@
-package org.hyperledger.besu.ethereum.bonsai.storage;public interface BonsaiPreImageProxy {
+/*
+ * 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);
+
+  /** PreImageProxy which does not store or cache preImages and only implements hashing. */
+  class NoOpPreImageProxy implements BonsaiPreImageProxy {
+
+    @Override
+    public Hash hashAndSavePreImage(final Bytes value) {
+      return Hash.hash(value);
+    }
+
+    @Override
+    public Optional<UInt256> getStorageTrieKeyPreimage(final Bytes32 trieKey) {
+      return Optional.empty();
+    }
+
+    @Override
+    public Optional<Address> getAccountTrieKeyPreimage(final Bytes32 trieKey) {
+      return Optional.empty();
+    }
+
+    @Override
+    public Updater updater() {
+      throw new UnsupportedOperationException("NoOpPreImageProxy does not implement an updater");
+    }
+  }
+
+  /**
+   * A caching PreImageProxy suitable for ReferenceTestWorldState which saves hashes in an unbounded
+   * BiMap.
+   */
+  class BonsaiReferenceTestPreImageProxy extends NoOpPreImageProxy {
+    BiMap<Hash, Bytes> preImageCache = HashBiMap.create();
+
+    @Override
+    public 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);
+    }
+  }
 }

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

@@ -44,13 +44,15 @@ public BonsaiSnapshotWorldStateKeyValueStorage(
       final BonsaiWorldStateKeyValueStorage parentWorldStateStorage,
       final SnappedKeyValueStorage segmentedWorldStateStorage,
       final KeyValueStorage trieLogStorage,
-      final ObservableMetricsSystem metricsSystem) {
+      final ObservableMetricsSystem metricsSystem,
+      final BonsaiPreImageProxy preImageProxy) {
     super(
         parentWorldStateStorage.flatDbMode,
         parentWorldStateStorage.flatDbReaderStrategy,
         segmentedWorldStateStorage,
         trieLogStorage,
-        metricsSystem);
+        metricsSystem,
+        preImageProxy);
     this.parentWorldStateStorage = parentWorldStateStorage;
     this.subscribeParentId = parentWorldStateStorage.subscribe(this);
   }
@@ -62,7 +64,8 @@ public BonsaiSnapshotWorldStateKeyValueStorage(
         worldStateStorage,
         ((SnappableKeyValueStorage) worldStateStorage.composedWorldStateStorage).takeSnapshot(),
         worldStateStorage.trieLogStorage,
-        metricsSystem);
+        metricsSystem,
+        worldStateStorage.preImageProxy);
   }
 
   private boolean isClosedGet() {

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

@@ -19,12 +19,14 @@
 import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.CODE_STORAGE;
 import static org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier.TRIE_BRANCH_STORAGE;
 
-import org.apache.tuweni.units.bigints.UInt256;
+import org.hyperledger.besu.datatypes.Address;
 import org.hyperledger.besu.datatypes.Hash;
 import org.hyperledger.besu.datatypes.StorageSlotKey;
+import org.hyperledger.besu.ethereum.bonsai.BonsaiAccount;
 import org.hyperledger.besu.ethereum.bonsai.storage.flat.FlatDbReaderStrategy;
 import org.hyperledger.besu.ethereum.bonsai.storage.flat.FullFlatDbReaderStrategy;
 import org.hyperledger.besu.ethereum.bonsai.storage.flat.PartialFlatDbReaderStrategy;
+import org.hyperledger.besu.ethereum.bonsai.worldview.BonsaiWorldView;
 import org.hyperledger.besu.ethereum.storage.StorageProvider;
 import org.hyperledger.besu.ethereum.storage.keyvalue.KeyValueSegmentIdentifier;
 import org.hyperledger.besu.ethereum.trie.MerkleTrie;
@@ -33,6 +35,7 @@
 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.evm.worldstate.WorldState;
 import org.hyperledger.besu.metrics.ObservableMetricsSystem;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorage;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorageTransaction;
@@ -41,25 +44,28 @@
 import org.hyperledger.besu.util.Subscribers;
 
 import java.nio.charset.StandardCharsets;
+import java.util.Comparator;
 import java.util.List;
 import java.util.Map;
 import java.util.NavigableMap;
 import java.util.Optional;
 import java.util.TreeMap;
 import java.util.concurrent.atomic.AtomicBoolean;
-import java.util.function.Function;
 import java.util.function.Predicate;
 import java.util.function.Supplier;
 import java.util.stream.Collectors;
+import java.util.stream.Stream;
 
 import org.apache.tuweni.bytes.Bytes;
 import org.apache.tuweni.bytes.Bytes32;
+import org.apache.tuweni.units.bigints.UInt256;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 @SuppressWarnings("unused")
 public class BonsaiWorldStateKeyValueStorage implements WorldStateStorage, AutoCloseable {
-  Bytes32 BYTES32_MAX_VALUE = Bytes32.fromHexString("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff");
+  Bytes32 BYTES32_MAX_VALUE =
+      Bytes32.fromHexString("0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff");
   private static final Logger LOG = LoggerFactory.getLogger(BonsaiWorldStateKeyValueStorage.class);
 
   // 0x776f726c64526f6f74
@@ -85,18 +91,25 @@ public class BonsaiWorldStateKeyValueStorage implements WorldStateStorage, AutoC
 
   protected final Subscribers<BonsaiStorageSubscriber> subscribers = Subscribers.create();
 
-  // no-op default hash preImage mapper:
-  protected Function<Hash, Optional<Bytes>> hashPreImageMapper = __ -> Optional.empty();
+  final BonsaiPreImageProxy preImageProxy;
 
   public BonsaiWorldStateKeyValueStorage(
       final StorageProvider provider, final ObservableMetricsSystem metricsSystem) {
+    this(provider, metricsSystem, new BonsaiPreImageProxy.NoOpPreImageProxy());
+  }
+
+  public BonsaiWorldStateKeyValueStorage(
+      final StorageProvider provider,
+      final ObservableMetricsSystem metricsSystem,
+      final BonsaiPreImageProxy preImageProxy) {
     this.composedWorldStateStorage =
         provider.getStorageBySegmentIdentifiers(
             List.of(
                 ACCOUNT_INFO_STATE, CODE_STORAGE, ACCOUNT_STORAGE_STORAGE, TRIE_BRANCH_STORAGE));
     this.trieLogStorage =
         provider.getStorageBySegmentIdentifier(KeyValueSegmentIdentifier.TRIE_LOG_STORAGE);
     this.metricsSystem = metricsSystem;
+    this.preImageProxy = preImageProxy;
     loadFlatDbStrategy();
   }
 
@@ -105,12 +118,14 @@ public BonsaiWorldStateKeyValueStorage(
       final FlatDbReaderStrategy flatDbReaderStrategy,
       final SegmentedKeyValueStorage composedWorldStateStorage,
       final KeyValueStorage trieLogStorage,
-      final ObservableMetricsSystem metricsSystem) {
+      final ObservableMetricsSystem metricsSystem,
+      final BonsaiPreImageProxy preImageProxy) {
     this.flatDbMode = flatDbMode;
     this.flatDbReaderStrategy = flatDbReaderStrategy;
     this.composedWorldStateStorage = composedWorldStateStorage;
     this.trieLogStorage = trieLogStorage;
     this.metricsSystem = metricsSystem;
+    this.preImageProxy = preImageProxy;
   }
 
   public void loadFlatDbStrategy() {
@@ -267,31 +282,42 @@ public Map<Bytes32, Bytes> streamFlatStorages(
             composedWorldStateStorage, accountHash, startKeyHash, endKeyHash, max);
   }
 
-  /**
-   * Provide a function to map hash values to an Optional wrapping their corresponding preImage
-   * or empty if the preImage is not available.
-   *
-   * @param hashPreImageMapper preImage mapper function
-   */
-  public void setPreImageMapper(final Function<Hash, Optional<Bytes>> hashPreImageMapper) {
-    this.hashPreImageMapper = hashPreImageMapper;
-  }
-
-  public NavigableMap<Bytes32, AccountStorageEntry> storageEntriesFrom(final Hash addressHash,
-      final Bytes32 startKeyHash, final int limit) {
-        return streamFlatStorages(addressHash, startKeyHash, BYTES32_MAX_VALUE, limit)
-            .entrySet()
-            .stream()
-            .collect(Collectors.toMap(
+  public NavigableMap<Bytes32, AccountStorageEntry> storageEntriesFrom(
+      final Hash addressHash, final Bytes32 startKeyHash, final int limit) {
+    return streamFlatStorages(addressHash, startKeyHash, BYTES32_MAX_VALUE, limit)
+        .entrySet()
+        // map back to slot keys using preImage provider:
+        .stream()
+        .collect(
+            Collectors.toMap(
                 e -> e.getKey(),
-                e -> AccountStorageEntry.create(
-                    UInt256.fromBytes(e.getValue()),
-                    Hash.wrap(e.getKey()),
-                    hashPreImageMapper.apply(Hash.wrap(e.getKey()))
-                        .map(UInt256::fromBytes)),
-                (a,b) -> a,
-                TreeMap::new
-            ));
+                e ->
+                    AccountStorageEntry.create(
+                        UInt256.fromBytes(e.getValue()),
+                        Hash.wrap(e.getKey()),
+                        preImageProxy.getStorageTrieKeyPreimage(e.getKey())),
+                (a, b) -> a,
+                TreeMap::new));
+  }
+
+  public Stream<WorldState.StreamableAccount> streamAccounts(
+      final BonsaiWorldView context, final Bytes32 startKeyHash, final int limit) {
+    return streamFlatAccounts(startKeyHash, BYTES32_MAX_VALUE, Long.MAX_VALUE)
+        .entrySet()
+        // map back to addresses using preImage provider:
+        .stream()
+        .map(
+            entry ->
+                preImageProxy
+                    .getAccountTrieKeyPreimage(entry.getKey())
+                    .map(
+                        address ->
+                            new WorldState.StreamableAccount(
+                                Optional.of(address),
+                                BonsaiAccount.fromRLP(context, address, entry.getValue(), false))))
+        .filter(Optional::isPresent)
+        .map(Optional::get)
+        .sorted(Comparator.comparing(account -> account.getAddress().orElse(Address.ZERO)));
   }
 
   @Override

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

@@ -30,16 +30,17 @@ public BonsaiWorldStateLayerStorage(final BonsaiWorldStateKeyValueStorage parent
         new LayeredKeyValueStorage(parent.composedWorldStateStorage),
         parent.trieLogStorage,
         parent,
-        parent.metricsSystem);
+        parent.metricsSystem,
+        parent.preImageProxy);
   }
 
   public BonsaiWorldStateLayerStorage(
       final SnappedKeyValueStorage composedWorldStateStorage,
       final KeyValueStorage trieLogStorage,
       final BonsaiWorldStateKeyValueStorage parent,
-      final ObservableMetricsSystem metricsSystem) {
-    super(parent, composedWorldStateStorage, trieLogStorage, metricsSystem);
-    setPreImageMapper(parent.hashPreImageMapper);
+      final ObservableMetricsSystem metricsSystem,
+      final BonsaiPreImageProxy preImageProxy) {
+    super(parent, composedWorldStateStorage, trieLogStorage, metricsSystem, preImageProxy);
   }
 
   @Override
@@ -53,6 +54,7 @@ public BonsaiWorldStateLayerStorage clone() {
         ((LayeredKeyValueStorage) composedWorldStateStorage).clone(),
         trieLogStorage,
         parentWorldStateStorage,
-        metricsSystem);
+        metricsSystem,
+        preImageProxy);
   }
 }