From 228f3903014fa3fdaa6fd483bc209cafc3e2b4a8 Mon Sep 17 00:00:00 2001
From: Georgios Konstantopoulos <me@gakonst.com>
Date: Sun, 19 Mar 2023 18:31:39 -0700
Subject: [PATCH] fix(init-genesis): insert bytecode/storage (#1853)

---
 crates/primitives/src/account.rs     |  6 ++++++
 crates/staged-sync/src/utils/init.rs | 30 ++++++++++++++++++++++++----
 2 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/crates/primitives/src/account.rs b/crates/primitives/src/account.rs
index b91796129c..9228350384 100644
--- a/crates/primitives/src/account.rs
+++ b/crates/primitives/src/account.rs
@@ -62,6 +62,12 @@ impl Bytecode {
         Self(RevmBytecode::new_raw(bytes))
     }
 
+    /// Create new bytecode from raw bytes and its hash.
+    pub fn new_raw_with_hash(bytes: Bytes, code_hash: H256) -> Self {
+        let revm_bytecode = unsafe { RevmBytecode::new_raw_with_hash(bytes, code_hash) };
+        Self(revm_bytecode)
+    }
+
     /// Set the hash of the inner bytecode.
     pub fn with_code_hash(mut self, code_hash: H256) -> Self {
         self.0.hash = code_hash;
diff --git a/crates/staged-sync/src/utils/init.rs b/crates/staged-sync/src/utils/init.rs
index 7ee1d76382..1038214231 100644
--- a/crates/staged-sync/src/utils/init.rs
+++ b/crates/staged-sync/src/utils/init.rs
@@ -1,11 +1,11 @@
 use reth_db::{
-    cursor::DbCursorRO,
+    cursor::{DbCursorRO, DbCursorRW},
     database::Database,
     mdbx::{Env, WriteMap},
     tables,
     transaction::{DbTx, DbTxMut},
 };
-use reth_primitives::{Account, ChainSpec, H256};
+use reth_primitives::{keccak256, Account, Bytecode, ChainSpec, StorageEntry, H256};
 use std::{path::Path, sync::Arc};
 use tracing::debug;
 
@@ -60,17 +60,39 @@ pub fn init_genesis<DB: Database>(
     debug!("Writing genesis block.");
     let tx = db.tx_mut()?;
 
+    let mut account_cursor = tx.cursor_write::<tables::PlainAccountState>()?;
+    let mut storage_cursor = tx.cursor_write::<tables::PlainStorageState>()?;
+    let mut bytecode_cursor = tx.cursor_write::<tables::Bytecodes>()?;
+
     // Insert account state
     for (address, account) in &genesis.alloc {
-        tx.put::<tables::PlainAccountState>(
+        let mut bytecode_hash = None;
+        // insert bytecode hash
+        if let Some(code) = &account.code {
+            let hash = keccak256(code.as_ref());
+            bytecode_cursor.upsert(hash, Bytecode::new_raw_with_hash(code.0.clone(), hash))?;
+            bytecode_hash = Some(hash);
+        }
+        // insert plain account.
+        account_cursor.upsert(
             *address,
             Account {
                 nonce: account.nonce.unwrap_or_default(),
                 balance: account.balance,
-                bytecode_hash: None,
+                bytecode_hash,
             },
         )?;
+        // insert plain storages
+        if let Some(storage) = &account.storage {
+            for (&key, &value) in storage {
+                storage_cursor.upsert(*address, StorageEntry { key, value: value.into() })?
+            }
+        }
     }
+    // Drop all cursor so we can commit the changes at the end of the fn.
+    drop(account_cursor);
+    drop(storage_cursor);
+    drop(bytecode_cursor);
 
     // Insert header
     tx.put::<tables::CanonicalHeaders>(0, hash)?;