Adding transient storage maps with TLOAD/TSTORE implementation

rskj-core/src/main/java/org/ethereum/core/Repository.java

@@ -27,7 +27,7 @@
 
 import java.math.BigInteger;
 
-public interface Repository extends RepositorySnapshot {
+public interface Repository extends RepositorySnapshot, TransientRepository {
     Trie getTrie();
 
     /**

rskj-core/src/main/java/org/ethereum/core/TransientRepository.java

@@ -0,0 +1,40 @@
+/*
+ * This file is part of RskJ
+ * Copyright (C) 2024 RSK Labs Ltd.
+ * (derived from ethereumJ library, Copyright (c) 2016 <ether.camp>)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package org.ethereum.core;
+
+import co.rsk.core.RskAddress;
+import org.ethereum.vm.DataWord;
+
+import javax.annotation.Nullable;
+
+public interface TransientRepository  {
+
+    void addTransientStorageRow(RskAddress addr, DataWord key, DataWord value);
+
+    void addTransientStorageBytes(RskAddress addr, DataWord key, byte[] value);
+
+    void clearTransientStorage();
+
+    @Nullable
+    DataWord getTransientStorageValue(RskAddress addr, DataWord key);
+
+    @Nullable
+    byte[] getTransientStorageBytes(RskAddress addr, DataWord key);
+}

rskj-core/src/main/java/org/ethereum/db/MutableRepository.java

@@ -25,19 +25,29 @@
 import co.rsk.crypto.Keccak256;
 import co.rsk.db.MutableTrieCache;
 import co.rsk.db.MutableTrieImpl;
-import co.rsk.trie.*;
+import co.rsk.trie.IterationElement;
+import co.rsk.trie.MutableTrie;
+import co.rsk.trie.Trie;
+import co.rsk.trie.TrieKeySlice;
+import co.rsk.trie.TrieStore;
+import co.rsk.trie.TrieStoreImpl;
 import com.google.common.annotations.VisibleForTesting;
 import org.ethereum.core.AccountState;
 import org.ethereum.core.Repository;
 import org.ethereum.crypto.HashUtil;
 import org.ethereum.crypto.Keccak256Helper;
+import org.ethereum.datasource.HashMapDB;
 import org.ethereum.vm.DataWord;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import javax.annotation.Nonnull;
+import javax.annotation.Nullable;
 import java.math.BigInteger;
-import java.util.*;
+import java.util.HashSet;
+import java.util.Iterator;
+import java.util.Optional;
+import java.util.Set;
 
 public class MutableRepository implements Repository {
     private static final Logger logger = LoggerFactory.getLogger("repository");
@@ -47,21 +57,32 @@ public class MutableRepository implements Repository {
 
     private final TrieKeyMapper trieKeyMapper;
     private final MutableTrie mutableTrie;
+    private MutableTrie transientTrie;
     private final IReadWrittenKeysTracker tracker;
 
     public MutableRepository(TrieStore trieStore, Trie trie) {
-        this(new MutableTrieImpl(trieStore, trie));
+        this(new MutableTrieImpl(trieStore, trie), new MutableTrieImpl(new TrieStoreImpl(new HashMapDB()), new Trie()));
     }
 
     public MutableRepository(MutableTrie mutableTrie) {
+       this(mutableTrie, new MutableTrieImpl(new TrieStoreImpl(new HashMapDB()), new Trie()));
+    }
+
+    public MutableRepository(MutableTrie mutableTrie, IReadWrittenKeysTracker tracker) {
+        this(mutableTrie, new MutableTrieImpl(new TrieStoreImpl(new HashMapDB()), new Trie()), tracker);
+    }
+
+    public MutableRepository(MutableTrie mutableTrie, MutableTrie transientTrie) {
         this.trieKeyMapper = new TrieKeyMapper();
         this.mutableTrie = mutableTrie;
+        this.transientTrie = transientTrie;
         this.tracker = new DummyReadWrittenKeysTracker();
     }
 
-    public MutableRepository(MutableTrie mutableTrie, IReadWrittenKeysTracker tracker) {
+    public MutableRepository(MutableTrie mutableTrie, MutableTrie transientTrie, IReadWrittenKeysTracker tracker) {
         this.trieKeyMapper = new TrieKeyMapper();
         this.mutableTrie = mutableTrie;
+        this.transientTrie = transientTrie;
         this.tracker = tracker;
     }
 
@@ -327,7 +348,7 @@ public synchronized Set<RskAddress> getAccountsKeys() {
     // To start tracking, a new repository is created, with a MutableTrieCache in the middle
     @Override
     public synchronized Repository startTracking() {
-        return new MutableRepository(new MutableTrieCache(mutableTrie), tracker);
+        return new MutableRepository(new MutableTrieCache(mutableTrie), new MutableTrieCache(transientTrie), tracker);
     }
 
     @Override
@@ -338,11 +359,13 @@ public void save() {
     @Override
     public synchronized void commit() {
         mutableTrie.commit();
+        transientTrie.commit();
     }
 
     @Override
     public synchronized void rollback() {
         mutableTrie.rollback();
+        transientTrie.rollback();
     }
 
     @Override
@@ -406,4 +429,48 @@ private Optional<Keccak256> internalGetValueHash(byte[] key) {
         tracker.addNewReadKey(new ByteArrayWrapper(key));
         return mutableTrie.getValueHash(key);
     }
+
+    @Override
+    public void addTransientStorageRow(RskAddress addr, DataWord key, DataWord value) {
+        addTransientStorageBytes(addr, key, value.getByteArrayForStorage());
+    }
+
+    @Override
+    public void addTransientStorageBytes(RskAddress addr, DataWord key, byte[] value) {
+        byte[] triekey = trieKeyMapper.getAccountStorageKey(addr, key);
+
+        // Special case: if the value is an empty vector, we pass "null" which commands the trie to remove the item.
+        // Note that if the call comes from addStorageRow(), this method will already have replaced 0 by null, so the
+        // conversion here only applies if this is called directly. If suppose this only occurs in tests, but it can
+        // also occur in precompiled contracts that store data directly using this method.
+        if (value == null || value.length == 0) {
+            transientTrie.put(triekey, null);
+        } else {
+            transientTrie.put(triekey, value);
+        }
+    }
+
+    @Override
+    public void clearTransientStorage() {
+        this.transientTrie =  new MutableTrieImpl(new TrieStoreImpl(new HashMapDB()), new Trie());
+    }
+
+    @Nullable
+    @Override
+    public DataWord getTransientStorageValue(RskAddress addr, DataWord key) {
+        byte[] triekey = trieKeyMapper.getAccountStorageKey(addr, key);
+        byte[] value = transientTrie.get(triekey);
+        if (value == null) {
+            return null;
+        }
+
+        return DataWord.valueOf(value);
+    }
+
+    @Nullable
+    @Override
+    public byte[] getTransientStorageBytes(RskAddress addr, DataWord key) {
+        byte[] triekey = trieKeyMapper.getAccountStorageKey(addr, key);
+        return transientTrie.get(triekey);
+    }
 }

rskj-core/src/main/java/org/ethereum/db/TransientStorageRepositoryCreator.java

@@ -0,0 +1,36 @@
+/*
+ * This file is part of RskJ
+ * Copyright (C) 2024 RSK Labs Ltd.
+ * (derived from ethereumJ library, Copyright (c) 2016 <ether.camp>)
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU Lesser General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+ * GNU Lesser General Public License for more details.
+ *
+ * You should have received a copy of the GNU Lesser General Public License
+ * along with this program. If not, see <http://www.gnu.org/licenses/>.
+ */
+
+package org.ethereum.db;
+
+import co.rsk.db.MutableTrieImpl;
+import co.rsk.trie.Trie;
+import co.rsk.trie.TrieStoreImpl;
+import org.ethereum.core.Repository;
+import org.ethereum.datasource.HashMapDB;
+
+public class TransientStorageRepositoryCreator {
+
+    private TransientStorageRepositoryCreator() {
+    }
+
+    public static Repository createNewTransientStorage() {
+        return new MutableRepository(new MutableTrieImpl(new TrieStoreImpl(new HashMapDB()), new Trie()));
+    }
+}

rskj-core/src/main/java/org/ethereum/vm/VM.java

@@ -1327,30 +1327,30 @@ else if (oldValue != null && newValue.isZero()) {
     }
 
     protected void doTLOAD(){
-
         DataWord key = program.stackPop();
-        DataWord address = program.getOwnerAddress();
-
         if (isLogEnabled) {
-            logger.info("Executing TLOAD with parameters: address={} | key = {}", address, key);
+            logger.info("Executing TLOAD with parameters: key = {}",  key);
         }
+        DataWord val = program.transientStorageLoad(key);
 
-        program.transientStorageSave(key, address);
-        // key could be returned to the pool, but storageLoad semantics should be checked
+        if (val == null) {
+            val = DataWord.ZERO;
+        }
+
+        program.stackPush(val);
+        // key could be returned to the pool, but transientStorageLoad semantics should be checked
         // to make sure storageLoad always gets a copy, not a reference.
         program.step();
     }
 
     protected void doTSTORE(){
-
+        DataWord address = program.stackPop();
         DataWord value = program.stackPop();
-        DataWord address = program.getOwnerAddress();
-        DataWord key = DataWord.ZERO;
 
         if (isLogEnabled) {
-            logger.info("Executing TSTORE with parameters: address={} | key = {} | value = {}",  address, key, value);
+            logger.info("Executing TSTORE with parameters: address={} | value = {}",  address, value);
         }
-        program.transientStorageLoad(address, key, value);
+        program.transientStorageSave(address, value);
         program.step();
     }
 
@@ -1967,9 +1967,16 @@ protected void executeOpcode() {
             break;
             case OpCodes.OP_SSTORE: doSSTORE();
                 break;
-            case OpCodes.OP_TLOAD: doTLOAD();
+            case OpCodes.OP_TLOAD:
+                if (!activations.isActive(RSKIP446)) {
+                    throw Program.ExceptionHelper.invalidOpCode(program);
+                }
+                doTLOAD();
                 break;
             case OpCodes.OP_TSTORE: doTSTORE();
+                if (!activations.isActive(RSKIP446)) {
+                    throw Program.ExceptionHelper.invalidOpCode(program);
+                }
                 break;
             case OpCodes.OP_JUMP: doJUMP();
             break;

rskj-core/src/main/java/org/ethereum/vm/program/Program.java

@@ -72,8 +72,10 @@
 import java.math.BigInteger;
 import java.util.Arrays;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Iterator;
+import java.util.Map;
 import java.util.Optional;
 import java.util.Set;
 
@@ -95,7 +97,6 @@ public class Program {
     private static final Logger logger = LoggerFactory.getLogger("VM");
     private static final Logger gasLogger = LoggerFactory.getLogger("gas");
 
-
     public static final long MAX_MEMORY = (1<<30);
 
     //Max size for stack checks
@@ -985,8 +986,8 @@ private void storageSave(byte[] key, byte[] val) {
         getStorage().addStorageRow(getOwnerRskAddress(), keyWord, valWord);
     }
 
-    public void transientStorageSave(DataWord key, DataWord address) {
-
+    public void transientStorageSave(DataWord key, DataWord value) {
+        getStorage().addTransientStorageRow(getOwnerRskAddress(), key, value);
     }
 
     private RskAddress getOwnerRskAddress() {
@@ -1099,8 +1100,8 @@ public DataWord storageLoad(DataWord key) {
         return getStorage().getStorageValue(getOwnerRskAddress(), key);
     }
 
-    public void transientStorageLoad(DataWord address, DataWord key, DataWord value) {
-
+    public DataWord transientStorageLoad(DataWord key) {
+        return getStorage().getTransientStorageValue(getOwnerRskAddress(), key);
     }
 
     public DataWord getPrevHash() {

rskj-core/src/main/java/org/ethereum/vm/program/Storage.java

@@ -30,6 +30,7 @@
 import org.ethereum.vm.program.listener.ProgramListener;
 import org.ethereum.vm.program.listener.ProgramListenerAware;
 
+import javax.annotation.Nullable;
 import java.math.BigInteger;
 import java.util.Iterator;
 import java.util.Set;
@@ -223,4 +224,31 @@ public byte[] getRoot() {
     public void updateAccountState(RskAddress addr, AccountState accountState) {
         throw new UnsupportedOperationException();
     }
+
+    @Override
+    public void addTransientStorageRow(RskAddress addr, DataWord key, DataWord value) {
+        repository.addTransientStorageRow(addr, key, value);
+    }
+
+    @Override
+    public void addTransientStorageBytes(RskAddress addr, DataWord key, byte[] value) {
+        repository.addTransientStorageBytes(addr, key, value);
+    }
+
+    @Override
+    public void clearTransientStorage() {
+        repository.clearTransientStorage();
+    }
+
+    @Nullable
+    @Override
+    public DataWord getTransientStorageValue(RskAddress addr, DataWord key) {
+        return repository.getTransientStorageValue(addr, key);
+    }
+
+    @Nullable
+    @Override
+    public byte[] getTransientStorageBytes(RskAddress addr, DataWord key) {
+        return repository.getTransientStorageBytes(addr, key);
+    }
 }