Refactor from code to get advantage from the MutableRepository

- In order to get advantage from the whole flow of tracking, rollback and commits already present in the MutableRepository.
We refactored a bit the logic of transient storage, now it's simpler.

- The task isn't done yet, more tests needs to be added to validate that this is working properly.

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

@@ -30,16 +30,19 @@
 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.HashSet;
 import java.util.Iterator;
@@ -54,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;
     }
 
@@ -334,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
@@ -345,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
@@ -413,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/vm/program/Program.java

@@ -40,7 +40,6 @@
 import org.ethereum.core.SignatureCache;
 import org.ethereum.core.Transaction;
 import org.ethereum.crypto.HashUtil;
-import org.ethereum.db.TransientStorageRepositoryCreator;
 import org.ethereum.util.ByteUtil;
 import org.ethereum.util.FastByteComparisons;
 import org.ethereum.vm.DataWord;
@@ -117,7 +116,6 @@ public class Program {
     private final Stack stack;
     private final Memory memory;
     private final Storage storage;
-    private final Map<RskAddress, Repository> transientStorages;
     private byte[] returnDataBuffer;
 
     private final ProgramResult result = new ProgramResult();
@@ -181,7 +179,6 @@ public Program(
         this.stack = setupProgramListener(new Stack());
         this.stack.ensureCapacity(1024); // faster?
         this.storage = setupProgramListener(new Storage(programInvoke));
-        this.transientStorages = new HashMap<>();
         this.deletedAccountsInBlock = new HashSet<>(deletedAccounts);
         this.signatureCache = signatureCache;
         precompile();
@@ -457,16 +454,6 @@ public Repository getStorage() {
         return this.storage;
     }
 
-    public Repository getTransientStorage(RskAddress addr) {
-        Repository current = transientStorages.get(addr);
-        if(current == null) {
-            current = TransientStorageRepositoryCreator.createNewTransientStorage();
-            transientStorages.put(addr, current);
-        }
-        return current;
-    }
-
-
     @SuppressWarnings("ThrowableResultOfMethodCallIgnored")
     public void createContract(DataWord value, DataWord memStart, DataWord memSize) {
         RskAddress senderAddress = new RskAddress(getOwnerAddress());
@@ -1000,9 +987,7 @@ private void storageSave(byte[] key, byte[] val) {
     }
 
     public void transientStorageSave(DataWord key, DataWord value) {
-        RskAddress addr = getOwnerRskAddress();
-        Repository storage = getTransientStorage(addr);
-        storage.addStorageRow(addr, key, value);
+        getStorage().addTransientStorageRow(getOwnerRskAddress(), key, value);
     }
 
     private RskAddress getOwnerRskAddress() {
@@ -1116,10 +1101,7 @@ public DataWord storageLoad(DataWord key) {
     }
 
     public DataWord transientStorageLoad(DataWord key) {
-        RskAddress addr = getOwnerRskAddress();
-        Repository currentTransientStorage = getTransientStorage(addr);
-
-        return currentTransientStorage.getStorageValue(addr, 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);
+    }
 }