Add storage into genesis state

Author: Aurélien Richez

src/main/scala/io/iohk/ethereum/blockchain/data/GenesisDataLoader.scala

@@ -86,14 +86,14 @@ class GenesisDataLoader(blockchain: Blockchain, blockchainConfig: BlockchainConf
     } yield ()
   }
 
-  def loadGenesisData(genesisData: GenesisData): Try[Unit] = {
+  def loadGenesisData(genesisData: GenesisData, storageRootHashes: Map[Address, ByteString] = Map.empty): Try[Unit] = {
     import MerklePatriciaTrie.defaultByteArraySerializable
 
     val stateStorage = blockchain.getStateStorage
     val storage = stateStorage.getReadOnlyStorage
     val initalRootHash = MerklePatriciaTrie.EmptyRootHash
 
-    val stateMptRootHash = getGenesisStateRoot(genesisData, initalRootHash, storage)
+    val stateMptRootHash = getGenesisStateRoot(genesisData, initalRootHash, storage, storageRootHashes)
     val header: BlockHeader = prepareHeader(genesisData, stateMptRootHash)
 
     log.debug(s"Prepared genesis header: $header")
@@ -122,20 +122,27 @@ class GenesisDataLoader(blockchain: Blockchain, blockchainConfig: BlockchainConf
     }
   }
 
-  private def getGenesisStateRoot(genesisData: GenesisData, initalRootHash: Array[Byte], storage: MptStorage) = {
+  private def getGenesisStateRoot(
+      genesisData: GenesisData,
+      initalRootHash: Array[Byte],
+      storage: MptStorage,
+      storageRootHashes: Map[Address, ByteString]
+  ) = {
     import MerklePatriciaTrie.defaultByteArraySerializable
 
     genesisData.alloc.zipWithIndex.foldLeft(initalRootHash) { case (rootHash, ((address, genesisAccount), _)) =>
       val mpt = MerklePatriciaTrie[Array[Byte], Account](rootHash, storage)
       val paddedAddress = address.reverse.padTo(addressLength, "0").reverse.mkString
+
       val stateRoot = mpt
         .put(
           crypto.kec256(Hex.decode(paddedAddress)),
           Account(
             nonce = genesisAccount.nonce
               .getOrElse(blockchainConfig.accountStartNonce),
             balance = genesisAccount.balance,
-            codeHash = genesisAccount.code.map(codeValue => crypto.kec256(codeValue)).getOrElse(Account.EmptyCodeHash)
+            codeHash = genesisAccount.code.map(codeValue => crypto.kec256(codeValue)).getOrElse(Account.EmptyCodeHash),
+            storageRoot = storageRootHashes.getOrElse(Address(Hex.decode(paddedAddress)), Account.EmptyStorageRootHash)
           )
         )
         .getRootHash

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

@@ -157,10 +157,14 @@ class TestService(
 
     //save account codes to world state
     storeGenesisAccountCodes(newBlockchainConfig, genesisData.alloc)
-    storeGenesisAccountStorageData(newBlockchainConfig, genesisData.alloc)
+    val storesRootHash = storeGenesisAccountStorageData(newBlockchainConfig, genesisData.alloc)
     // update test ledger with new config
     testLedgerWrapper.blockchainConfig = newBlockchainConfig
 
+    // remove current genesis (Try because it may not exist)
+    Try(blockchain.removeBlock(blockchain.genesisHeader.hash, withState = false))
+    genesisDataLoader.loadGenesisData(genesisData, storesRootHash)
+
     accountAddresses = genesisData.alloc.keys.toList
     accountRangeOffset = 0
 
@@ -181,7 +185,7 @@ class TestService(
     InMemoryWorldStateProxy.persistState(worldToPersist)
   }
 
-  private def storeGenesisAccountStorageData(config: BlockchainConfig, accounts: Map[String, GenesisAccount]): Unit = {
+  private def storeGenesisAccountStorageData(config: BlockchainConfig, accounts: Map[String, GenesisAccount]) = {
     val genesisBlock = blockchain.getBlockByNumber(0).get
     val world =
       blockchain.getWorldStateProxy(0, UInt256.Zero, genesisBlock.header.stateRoot, false, config.ethCompatibleStorage)
@@ -199,7 +203,9 @@ class TestService(
       updatedWorld
     })
 
-    InMemoryWorldStateProxy.persistState(worldToPersist)
+    InMemoryWorldStateProxy.persistState(worldToPersist).contractStorages.map { case (address, values) =>
+      address -> ByteString(values.inner.getRootHash)
+    }
   }
 
   def mineBlocks(request: MineBlocksRequest): ServiceResponse[MineBlocksResponse] = {