[ETCM-964] Fix preimage cache to make vmArithmeticTest pass

src/main/scala/io/iohk/ethereum/jsonrpc/TestService.scala

@@ -117,8 +117,7 @@ class TestService(
     consensusConfig: ConsensusConfig,
     testModeComponentsProvider: TestModeComponentsProvider,
     initialConfig: BlockchainConfig,
-    transactionMappingStorage: TransactionMappingStorage,
-    preimageCache: collection.concurrent.Map[ByteString, UInt256]
+    transactionMappingStorage: TransactionMappingStorage
 )(implicit
     scheduler: Scheduler
 ) extends Logger {
@@ -131,6 +130,8 @@ class TestService(
   private var currentConfig: BlockchainConfig = initialConfig
   private var blockTimestamp: Long = 0
   private var sealEngine: SealEngineType = SealEngineType.NoReward
+  private val preimageCache: collection.concurrent.Map[ByteString, UInt256] =
+    new collection.concurrent.TrieMap[ByteString, UInt256]()
 
   def setChainParams(request: SetChainParamsRequest): ServiceResponse[SetChainParamsResponse] = {
     currentConfig = buildNewConfig(request.chainParams.blockchainParams)
@@ -235,7 +236,9 @@ class TestService(
     def mineBlock(): Task[Unit] = {
       getBlockForMining(blockchain.getBestBlock().get)
         .flatMap(blockForMining =>
-          testModeComponentsProvider.blockImport(currentConfig, sealEngine).importBlock(blockForMining.block)
+          testModeComponentsProvider
+            .blockImport(currentConfig, preimageCache, sealEngine)
+            .importBlock(blockForMining.block)
         )
         .map { res =>
           log.info("Block mining result: " + res)
@@ -271,7 +274,7 @@ class TestService(
         Task.now(Left(JsonRpcError(-1, "block validation failed!", None)))
       case Success(value) =>
         testModeComponentsProvider
-          .blockImport(currentConfig, sealEngine)
+          .blockImport(currentConfig, preimageCache, sealEngine)
           .importBlock(value)
           .flatMap(handleResult)
     }

src/main/scala/io/iohk/ethereum/ledger/BlockExecution.scala

@@ -63,15 +63,7 @@ class BlockExecution(
       parentHeader <- blockchainReader
         .getBlockHeaderByHash(block.header.parentHash)
         .toRight(MissingParentError) // Should not never occur because validated earlier
-      initialWorld = InMemoryWorldStateProxy(
-        evmCodeStorage = evmCodeStorage,
-        blockchain.getBackingMptStorage(block.header.number),
-        (number: BigInt) => blockchainReader.getBlockHeaderByNumber(number).map(_.hash),
-        accountStartNonce = blockchainConfig.accountStartNonce,
-        stateRootHash = parentHeader.stateRoot,
-        noEmptyAccounts = EvmConfig.forBlock(parentHeader.number, blockchainConfig).noEmptyAccounts,
-        ethCompatibleStorage = blockchainConfig.ethCompatibleStorage
-      )
+      initialWorld = buildInitialWorld(block, parentHeader)
       execResult <- executeBlockTransactions(block, initialWorld)
       worldToPersist <- Either
         .catchOnly[MPTException](blockPreparator.payBlockReward(block, execResult.worldState))
@@ -81,6 +73,18 @@ class BlockExecution(
     } yield execResult.copy(worldState = worldPersisted)
   }
 
+  protected def buildInitialWorld(block: Block, parentHeader: BlockHeader): InMemoryWorldStateProxy = {
+    InMemoryWorldStateProxy(
+      evmCodeStorage = evmCodeStorage,
+      blockchain.getBackingMptStorage(block.header.number),
+      (number: BigInt) => blockchainReader.getBlockHeaderByNumber(number).map(_.hash),
+      accountStartNonce = blockchainConfig.accountStartNonce,
+      stateRootHash = parentHeader.stateRoot,
+      noEmptyAccounts = EvmConfig.forBlock(parentHeader.number, blockchainConfig).noEmptyAccounts,
+      ethCompatibleStorage = blockchainConfig.ethCompatibleStorage
+    )
+  }
+
   /** This function runs transactions
     *
     * @param block the block with transactions to run
@@ -144,7 +148,8 @@ class BlockExecution(
     * @param blocks   blocks to be executed
     * @param parentChainWeight parent weight
     *
-    * @return a list of blocks in incremental order that were correctly executed and an optional [[BlockExecutionError]]
+    * @return a list of blocks in incremental order that were correctly executed and an optional
+    *         [[io.iohk.ethereum.ledger.BlockExecutionError]]
     */
   def executeAndValidateBlocks(
       blocks: List[Block],

src/main/scala/io/iohk/ethereum/nodebuilder/NodeBuilder.scala

@@ -427,8 +427,7 @@ trait TestServiceBuilder {
       consensusConfig,
       testModeComponentsProvider,
       blockchainConfig,
-      storagesInstance.storages.transactionMappingStorage,
-      preimages
+      storagesInstance.storages.transactionMappingStorage
     )(scheduler)
 }
 

src/main/scala/io/iohk/ethereum/testmode/TestBlockchainBuilder.scala

@@ -1,17 +1,12 @@
 package io.iohk.ethereum.testmode
 
 import akka.util.ByteString
-import io.iohk.ethereum.crypto
 import io.iohk.ethereum.domain.{BlockchainImpl, UInt256}
-import io.iohk.ethereum.ledger.InMemoryWorldStateProxy
 import io.iohk.ethereum.nodebuilder.{BlockchainBuilder, StorageBuilder}
 
 trait TestBlockchainBuilder extends BlockchainBuilder {
   self: StorageBuilder =>
 
-  lazy val preimages: collection.concurrent.Map[ByteString, UInt256] =
-    new collection.concurrent.TrieMap[ByteString, UInt256]()
-
   override lazy val blockchain: BlockchainImpl = {
     val storages = storagesInstance.storages
     new BlockchainImpl(

src/main/scala/io/iohk/ethereum/testmode/TestModeBlockExecution.scala

@@ -0,0 +1,39 @@
+package io.iohk.ethereum.testmode
+
+import akka.util.ByteString
+import io.iohk.ethereum.crypto
+import io.iohk.ethereum.db.storage.EvmCodeStorage
+import io.iohk.ethereum.domain.{Block, BlockHeader, BlockchainImpl, BlockchainReader, UInt256}
+import io.iohk.ethereum.ledger.{BlockExecution, BlockPreparator, BlockValidation, InMemoryWorldStateProxy}
+import io.iohk.ethereum.utils.BlockchainConfig
+import io.iohk.ethereum.vm.EvmConfig
+
+class TestModeBlockExecution(
+    blockchain: BlockchainImpl,
+    blockchainReader: BlockchainReader,
+    evmCodeStorage: EvmCodeStorage,
+    blockchainConfig: BlockchainConfig,
+    blockPreparator: BlockPreparator,
+    blockValidation: BlockValidation,
+    saveStoragePreimage: (UInt256) => Unit
+) extends BlockExecution(
+      blockchain,
+      blockchainReader,
+      evmCodeStorage,
+      blockchainConfig,
+      blockPreparator,
+      blockValidation
+    ) {
+
+  override protected def buildInitialWorld(block: Block, parentHeader: BlockHeader): InMemoryWorldStateProxy =
+    TestModeWorldStateProxy(
+      evmCodeStorage = evmCodeStorage,
+      nodesKeyValueStorage = blockchain.getBackingMptStorage(block.header.number),
+      getBlockHashByNumber = (number: BigInt) => blockchainReader.getBlockHeaderByNumber(number).map(_.hash),
+      accountStartNonce = blockchainConfig.accountStartNonce,
+      stateRootHash = parentHeader.stateRoot,
+      noEmptyAccounts = EvmConfig.forBlock(parentHeader.number, blockchainConfig).noEmptyAccounts,
+      ethCompatibleStorage = blockchainConfig.ethCompatibleStorage,
+      saveStoragePreimage = saveStoragePreimage
+    )
+}

src/main/scala/io/iohk/ethereum/testmode/TestModeComponentsProvider.scala

@@ -1,9 +1,11 @@
 package io.iohk.ethereum.testmode
 
+import akka.util.ByteString
 import io.iohk.ethereum.consensus.difficulty.DifficultyCalculator
 import io.iohk.ethereum.consensus.{Consensus, ConsensusConfig}
+import io.iohk.ethereum.crypto
 import io.iohk.ethereum.db.storage.EvmCodeStorage
-import io.iohk.ethereum.domain.{BlockchainImpl, BlockchainReader}
+import io.iohk.ethereum.domain.{BlockchainImpl, BlockchainReader, UInt256}
 import io.iohk.ethereum.ledger.VMImpl
 import io.iohk.ethereum.ledger.StxLedger
 import io.iohk.ethereum.utils.BlockchainConfig
@@ -26,18 +28,23 @@ class TestModeComponentsProvider(
     vm: VMImpl
 ) {
 
-  def blockImport(blockchainConfig: BlockchainConfig, sealEngine: SealEngineType): BlockImport = {
+  def blockImport(
+      blockchainConfig: BlockchainConfig,
+      preimageCache: collection.concurrent.Map[ByteString, UInt256],
+      sealEngine: SealEngineType
+  ): BlockImport = {
     val blockQueue = BlockQueue(blockchain, syncConfig)
     val consensuz = consensus(blockchainConfig, sealEngine)
     val blockValidation = new BlockValidation(consensuz, blockchainReader, blockQueue)
     val blockExecution =
-      new BlockExecution(
+      new TestModeBlockExecution(
         blockchain,
         blockchainReader,
         evmCodeStorage,
         blockchainConfig,
         consensuz.blockPreparator,
-        blockValidation
+        blockValidation,
+        (key: UInt256) => preimageCache.put(crypto.kec256(key.bytes), key)
       )
 
     new BlockImport(