add contract code update in trie log - bonsai trie (#1963)

Signed-off-by: Karim TAAM <t2am.karim@gmail.com>

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

@@ -18,8 +18,11 @@ package org.hyperledger.besu.ethereum.bonsai;
 
 import org.hyperledger.besu.ethereum.core.Account;
 import org.hyperledger.besu.ethereum.core.Address;
+import org.hyperledger.besu.ethereum.core.BlockHeader;
 import org.hyperledger.besu.ethereum.core.Hash;
+import org.hyperledger.besu.ethereum.core.MutableWorldState;
 import org.hyperledger.besu.ethereum.core.WorldState;
+import org.hyperledger.besu.ethereum.core.WorldUpdater;
 import org.hyperledger.besu.ethereum.worldstate.StateTrieAccountValue;
 
 import java.util.HashMap;
@@ -32,8 +35,9 @@ import org.apache.tuweni.bytes.Bytes32;
 import org.apache.tuweni.units.bigints.UInt256;
 
 /** A World State backed first by trie log layer and then by another world state. */
-public class BonsaiLayeredWorldState implements BonsaiWorldView, WorldState {
+public class BonsaiLayeredWorldState implements MutableWorldState, BonsaiWorldView, WorldState {
 
+  private final BonsaiWorldStateArchive worldStateArchive;
   private final BonsaiWorldView parent;
   protected final long height;
   protected final TrieLogLayer trieLog;
@@ -41,10 +45,12 @@ public class BonsaiLayeredWorldState implements BonsaiWorldView, WorldState {
   private final Hash worldStateRootHash;
 
   BonsaiLayeredWorldState(
+      final BonsaiWorldStateArchive worldStateArchive,
       final BonsaiWorldView parent,
       final long height,
       final Hash worldStateRootHash,
       final TrieLogLayer trieLog) {
+    this.worldStateArchive = worldStateArchive;
     this.parent = parent;
     this.height = height;
     this.worldStateRootHash = worldStateRootHash;
@@ -73,6 +79,15 @@ public class BonsaiLayeredWorldState implements BonsaiWorldView, WorldState {
       if (maybeCode.isPresent()) {
         return maybeCode;
       }
+      final Optional<Hash> maybeLastChangedCodeLocation =
+          currentLayer.trieLog.getLastChangedCodeLocation(address);
+      if (maybeLastChangedCodeLocation.isPresent()) {
+        final Optional<TrieLogLayer> maybeTrieLogLayer =
+            worldStateArchive.getTrieLogLayer(maybeLastChangedCodeLocation.get());
+        if (maybeTrieLogLayer.isPresent()) {
+          return maybeTrieLogLayer.get().getCode(address);
+        }
+      }
       if (currentLayer.parent == null) {
         currentLayer = null;
       } else if (currentLayer.parent instanceof BonsaiLayeredWorldState) {
@@ -218,4 +233,20 @@ public class BonsaiLayeredWorldState implements BonsaiWorldView, WorldState {
   public Stream<StreamableAccount> streamAccounts(final Bytes32 startKeyHash, final int limit) {
     throw new UnsupportedOperationException("Bonsai does not support pruning and debug RPCs");
   }
+
+  @Override
+  public MutableWorldState copy() {
+    throw new UnsupportedOperationException(
+        "Bonsai Tries does not support direct duplication of the persisted tries.");
+  }
+
+  @Override
+  public void persist(final BlockHeader blockHeader) {
+    throw new UnsupportedOperationException("Layered worldState can not be persisted.");
+  }
+
+  @Override
+  public WorldUpdater updater() {
+    return new BonsaiWorldStateUpdater(this);
+  }
 }

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

@@ -32,6 +32,7 @@ import org.hyperledger.besu.ethereum.worldstate.WorldStateStorage;
 import org.hyperledger.besu.plugin.services.exception.StorageException;
 import org.hyperledger.besu.plugin.services.storage.KeyValueStorageTransaction;
 
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Function;
@@ -56,6 +57,8 @@ public class BonsaiPersistedWorldState implements MutableWorldState, BonsaiWorld
   private Hash worldStateRootHash;
   private Hash worldStateBlockHash;
 
+  private final Map<Address, Hash> contractCodeChangesHistory;
+
   public BonsaiPersistedWorldState(
       final BonsaiWorldStateArchive archive,
       final BonsaiWorldStateKeyValueStorage worldStateStorage) {
@@ -67,6 +70,12 @@ public class BonsaiPersistedWorldState implements MutableWorldState, BonsaiWorld
     worldStateBlockHash =
         Hash.wrap(Bytes32.wrap(worldStateStorage.getWorldStateBlockHash().orElse(Hash.ZERO)));
     updater = new BonsaiWorldStateUpdater(this);
+    contractCodeChangesHistory =
+        worldStateStorage
+            .getTrieLog(worldStateBlockHash)
+            .map(TrieLogLayer::fromBytes)
+            .map(TrieLogLayer::getContractCodeChangesHistory)
+            .orElse(new HashMap<>());
   }
 
   public BonsaiWorldStateArchive getArchive() {
@@ -249,13 +258,15 @@ public class BonsaiPersistedWorldState implements MutableWorldState, BonsaiWorld
             .put(WORLD_BLOCK_HASH_KEY, worldStateBlockHash.toArrayUnsafe());
         if (originalBlockHash.equals(blockHeader.getParentHash())) {
           LOG.debug("Writing Trie Log for {}", worldStateBlockHash);
-          final TrieLogLayer trieLog = updater.generateTrieLog(worldStateBlockHash);
+          final TrieLogLayer trieLog =
+              updater.generateTrieLog(worldStateBlockHash, contractCodeChangesHistory);
           trieLog.freeze();
           archive.addLayeredWorldState(
               new BonsaiLayeredWorldState(
-                  this, blockHeader.getNumber(), worldStateRootHash, trieLog));
+                  getArchive(), this, blockHeader.getNumber(), worldStateRootHash, trieLog));
           final BytesValueRLPOutput rlpLog = new BytesValueRLPOutput();
           trieLog.writeTo(rlpLog);
+
           stateUpdater
               .getTrieLogStorageTransaction()
               .put(worldStateBlockHash.toArrayUnsafe(), rlpLog.encoded().toArrayUnsafe());

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

@@ -73,7 +73,7 @@ public class BonsaiWorldStateArchive implements WorldStateArchive {
     layeredWorldStates.put(worldState.blockHash(), worldState);
   }
 
-  private Optional<TrieLogLayer> getTrieLogLayer(final Hash blockHash) {
+  public Optional<TrieLogLayer> getTrieLogLayer(final Hash blockHash) {
     if (layeredWorldStates.containsKey(blockHash)) {
       return Optional.of(layeredWorldStates.get(blockHash).getTrieLog());
     } else {
@@ -88,6 +88,22 @@ public class BonsaiWorldStateArchive implements WorldStateArchive {
         || worldStateStorage.isWorldStateAvailable(rootHash, blockHash);
   }
 
+  @Override
+  public Optional<MutableWorldState> getWorldState(final Hash rootHash, final Hash blockHash) {
+    if (layeredWorldStates.containsKey(blockHash)) {
+      return Optional.of(layeredWorldStates.get(blockHash));
+    } else {
+      final Optional<TrieLogLayer> trieLogLayer = getTrieLogLayer(blockHash);
+      if (trieLogLayer.isPresent()) {
+        final BlockHeader header = blockchain.getBlockHeader(persistedState.blockHash()).get();
+        return Optional.of(
+            new BonsaiLayeredWorldState(
+                this, persistedState, header.getNumber(), blockHash, trieLogLayer.get()));
+      }
+    }
+    return Optional.empty();
+  }
+
   @Override
   public Optional<MutableWorldState> getMutable(final Hash rootHash, final Hash blockHash) {
     if (blockHash.equals(persistedState.blockHash())) {

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

@@ -335,8 +335,9 @@ public class BonsaiWorldStateUpdater extends AbstractWorldUpdater<BonsaiWorldVie
     return results;
   }
 
-  public TrieLogLayer generateTrieLog(final Hash blockHash) {
-    final TrieLogLayer layer = new TrieLogLayer();
+  public TrieLogLayer generateTrieLog(
+      final Hash blockHash, final Map<Address, Hash> contractCodeChangesHistory) {
+    final TrieLogLayer layer = new TrieLogLayer(contractCodeChangesHistory);
     importIntoTrieLog(layer, blockHash);
     return layer;
   }
@@ -373,7 +374,8 @@ public class BonsaiWorldStateUpdater extends AbstractWorldUpdater<BonsaiWorldVie
       layer.addCodeChange(
           updatedCode.getKey(),
           updatedCode.getValue().getOriginal(),
-          updatedCode.getValue().getUpdated());
+          updatedCode.getValue().getUpdated(),
+          blockHash);
     }
 
     for (final Map.Entry<Address, Map<Hash, BonsaiValue<UInt256>>> updatesStorage :

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

@@ -51,13 +51,19 @@ public class TrieLogLayer {
   private final Map<Address, BonsaiValue<StateTrieAccountValue>> accounts;
   private final Map<Address, BonsaiValue<Bytes>> code;
   private final Map<Address, Map<Hash, BonsaiValue<UInt256>>> storage;
+  private final Map<Address, Hash> contractCodeChangesHistory;
   private boolean frozen = false;
 
   TrieLogLayer() {
     // TODO when tuweni fixes zero length byte comparison consider TreeMap
-    accounts = new HashMap<>();
-    code = new HashMap<>();
-    storage = new HashMap<>();
+    this(new HashMap<>());
+  }
+
+  public TrieLogLayer(final Map<Address, Hash> contractCodeChangesHistory) {
+    this.accounts = new HashMap<>();
+    this.code = new HashMap<>();
+    this.storage = new HashMap<>();
+    this.contractCodeChangesHistory = contractCodeChangesHistory;
   }
 
   /** Locks the layer so no new changes can be added; */
@@ -82,12 +88,14 @@ public class TrieLogLayer {
     accounts.put(address, new BonsaiValue<>(oldValue, newValue));
   }
 
-  void addCodeChange(final Address address, final Bytes oldValue, final Bytes newValue) {
+  void addCodeChange(
+      final Address address, final Bytes oldValue, final Bytes newValue, final Hash blockHash) {
     checkState(!frozen, "Layer is Frozen");
     code.put(
         address,
         new BonsaiValue<>(
             oldValue == null ? Bytes.EMPTY : oldValue, newValue == null ? Bytes.EMPTY : newValue));
+    contractCodeChangesHistory.put(address, blockHash);
   }
 
   void addStorageChange(
@@ -149,6 +157,12 @@ public class TrieLogLayer {
         newLayer.storage.put(address, storageChanges);
       }
 
+      if (input.nextIsNull()) {
+        input.skipNext();
+      } else {
+        newLayer.contractCodeChangesHistory.put(address, Hash.wrap(input.readBytes32()));
+      }
+
       // TODO add trie nodes
 
       // lenient leave list for forward compatible additions.
@@ -167,6 +181,7 @@ public class TrieLogLayer {
     addresses.addAll(accounts.keySet());
     addresses.addAll(code.keySet());
     addresses.addAll(storage.keySet());
+    addresses.addAll(contractCodeChangesHistory.keySet());
 
     output.startList(); // container
     output.writeBytes(blockHash);
@@ -204,6 +219,13 @@ public class TrieLogLayer {
         output.endList();
       }
 
+      final Hash blockHash = contractCodeChangesHistory.get(address);
+      if (blockHash == null) {
+        output.writeNull();
+      } else {
+        output.writeBytes(blockHash);
+      }
+
       // TODO write trie nodes
 
       output.endList(); // this change
@@ -223,6 +245,10 @@ public class TrieLogLayer {
     return storage.entrySet().stream();
   }
 
+  public Map<Address, Hash> getContractCodeChangesHistory() {
+    return contractCodeChangesHistory;
+  }
+
   boolean hasStorageChanges(final Address address) {
     return storage.containsKey(address);
   }
@@ -244,6 +270,14 @@ public class TrieLogLayer {
     return Optional.ofNullable(code.get(address)).map(BonsaiValue::getUpdated);
   }
 
+  public Optional<Hash> getLastChangedCodeLocation(final Address address) {
+    if (contractCodeChangesHistory.containsKey(address)
+        && !contractCodeChangesHistory.get(address).isEmpty()) {
+      return Optional.of(contractCodeChangesHistory.get(address));
+    }
+    return Optional.empty();
+  }
+
   Optional<UInt256> getStorageBySlotHash(final Address address, final Hash slotHash) {
     return Optional.ofNullable(storage.get(address))
         .map(i -> i.get(slotHash))

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/transaction/TransactionSimulator.java

@@ -141,7 +141,6 @@ public class TransactionSimulator {
     if (header == null) {
       return Optional.empty();
     }
-
     final MutableWorldState publicWorldState =
         worldStateArchive.getMutable(header.getStateRoot(), header.getHash()).orElse(null);
 

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/worldstate/DefaultWorldStateArchive.java

@@ -54,6 +54,11 @@ public class DefaultWorldStateArchive implements WorldStateArchive {
     return worldStateStorage.isWorldStateAvailable(rootHash, blockHash);
   }
 
+  @Override
+  public Optional<MutableWorldState> getWorldState(final Hash rootHash, final Hash blockHash) {
+    return getMutable(rootHash, blockHash);
+  }
+
   @Override
   public Optional<MutableWorldState> getMutable(final Hash rootHash, final Hash blockHash) {
     if (!worldStateStorage.isWorldStateAvailable(rootHash, blockHash)) {

ethereum/core/src/main/java/org/hyperledger/besu/ethereum/worldstate/WorldStateArchive.java

@@ -35,6 +35,8 @@ public interface WorldStateArchive {
 
   boolean isWorldStateAvailable(Hash rootHash, Hash blockHash);
 
+  Optional<MutableWorldState> getWorldState(Hash rootHash, Hash blockHash);
+
   Optional<MutableWorldState> getMutable(Hash rootHash, Hash blockHash);
 
   MutableWorldState getMutable();

ethereum/core/src/test-support/java/org/hyperledger/besu/ethereum/core/InMemoryStorageProvider.java

@@ -14,7 +14,9 @@
  */
 package org.hyperledger.besu.ethereum.core;
 
+import org.hyperledger.besu.ethereum.bonsai.BonsaiWorldStateArchive;
 import org.hyperledger.besu.ethereum.bonsai.BonsaiWorldStateKeyValueStorage;
+import org.hyperledger.besu.ethereum.chain.Blockchain;
 import org.hyperledger.besu.ethereum.chain.BlockchainStorage;
 import org.hyperledger.besu.ethereum.chain.DefaultBlockchain;
 import org.hyperledger.besu.ethereum.chain.MutableBlockchain;
@@ -61,6 +63,11 @@ public class InMemoryStorageProvider implements StorageProvider {
         new WorldStatePreimageKeyValueStorage(new InMemoryKeyValueStorage()));
   }
 
+  public static BonsaiWorldStateArchive createInMemoryWorldStateArchive(
+      final Blockchain blockchain) {
+    return new BonsaiWorldStateArchive(new InMemoryStorageProvider(), blockchain);
+  }
+
   public static MutableWorldState createInMemoryWorldState() {
     final InMemoryStorageProvider provider = new InMemoryStorageProvider();
     return new DefaultMutableWorldState(