refactor(l1): improve the way we remove ONLY the storage of an account

crates/common/types/account_update.rs

@@ -11,6 +11,8 @@ pub struct AccountUpdate {
     pub info: Option<AccountInfo>,
     pub code: Option<Bytes>,
     pub added_storage: BTreeMap<H256, U256>,
+    /// If account was destroyed and then modified we need this for removing its storage but not the entire account.
+    pub removed_storage: bool,
     // Matches TODO in code
     // removed_storage_keys: Vec<H256>,
 }
@@ -35,6 +37,7 @@ impl AccountUpdate {
 
     pub fn merge(&mut self, other: AccountUpdate) {
         self.removed = other.removed;
+        self.removed_storage |= other.removed_storage;
         if let Some(info) = other.info {
             self.info = Some(info);
         }

crates/common/types/block_execution_witness.rs

@@ -12,7 +12,7 @@ use crate::{
 use bytes::Bytes;
 use ethereum_types::{Address, H256, U256};
 use ethrex_rlp::{decode::RLPDecode, encode::RLPEncode};
-use ethrex_trie::{NodeRLP, Trie};
+use ethrex_trie::{EMPTY_TRIE_HASH, NodeRLP, Trie};
 use rkyv::{Archive, Deserialize as RDeserialize, Serialize as RSerialize};
 use serde::de::{SeqAccess, Visitor};
 use serde::ser::SerializeSeq;
@@ -246,6 +246,9 @@ impl GuestProgramState {
                         .expect("failed to decode account state"),
                     None => AccountState::default(),
                 };
+                if update.removed_storage {
+                    account_state.storage_root = *EMPTY_TRIE_HASH;
+                }
                 if let Some(info) = &update.info {
                     account_state.nonce = info.nonce;
                     account_state.balance = info.balance;

crates/common/utils.rs

@@ -4,6 +4,8 @@ use hex::FromHexError;
 use sha3::Digest;
 use sha3::Keccak256;
 
+pub const ZERO_U256: U256 = U256([0, 0, 0, 0]);
+
 /// Converts a big endian slice to a u256, faster than `u256::from_big_endian`.
 pub fn u256_from_big_endian(slice: &[u8]) -> U256 {
     let mut padded = [0u8; 32];

crates/l2/common/src/state_diff.rs

@@ -299,6 +299,7 @@ impl StateDiff {
                     info: account_info,
                     code: diff.bytecode.clone(),
                     added_storage: diff.storage.clone().into_iter().collect(),
+                    removed_storage: false,
                 },
             );
         }

crates/storage/store.rs

@@ -396,6 +396,9 @@ impl Store {
                 Some(encoded_state) => AccountState::decode(&encoded_state)?,
                 None => AccountState::default(),
             };
+            if update.removed_storage {
+                account_state.storage_root = *EMPTY_TRIE_HASH;
+            }
             if let Some(info) = &update.info {
                 account_state.nonce = info.nonce;
                 account_state.balance = info.balance;
@@ -465,6 +468,9 @@ impl Store {
                         self.add_account_code(info.code_hash, code.clone()).await?;
                     }
                 }
+                if update.removed_storage {
+                    account_state.storage_root = *EMPTY_TRIE_HASH;
+                }
                 // Store the added storage in the account's storage trie and compute its new root
                 if !update.added_storage.is_empty() {
                     let (_witness, storage_trie) = match storage_tries.entry(update.address) {

crates/vm/levm/src/db/gen_db.rs

@@ -6,6 +6,7 @@ use ethrex_common::Address;
 use ethrex_common::H256;
 use ethrex_common::U256;
 use ethrex_common::types::Account;
+use ethrex_common::utils::ZERO_U256;
 use ethrex_common::utils::keccak;
 
 use super::Database;
@@ -163,31 +164,6 @@ impl GeneralizedDatabase {
                         "Failed to get account {address} from immutable cache",
                     ))))?;
 
-            // Edge cases:
-            //   1. Account was destroyed and created again afterwards.
-            //   2. Account was destroyed but then was sent ETH, so it's not going to be removed completely from the trie.
-            // This is a way of removing the storage of an account but keeping the info.
-            if new_state_account.status == AccountStatus::DestroyedModified {
-                // Push to account updates the removal of the account and then push the new state of the account.
-                account_updates.push(AccountUpdate::removed(*address));
-                let new_account_update = AccountUpdate {
-                    address: *address,
-                    removed: false,
-                    info: Some(new_state_account.info.clone()),
-                    code: Some(
-                        self.codes
-                            .get(&new_state_account.info.code_hash)
-                            .ok_or(VMError::Internal(InternalError::Custom(format!(
-                                "Failed to get code for account {address}"
-                            ))))?
-                            .clone(),
-                    ),
-                    added_storage: new_state_account.storage.clone(),
-                };
-                account_updates.push(new_account_update);
-                continue;
-            }
-
             let mut acc_info_updated = false;
             let mut storage_updated = false;
 
@@ -213,11 +189,22 @@ impl GeneralizedDatabase {
                     None
                 };
 
+            // Account will have only its storage removed if it was Destroyed and then modified
+            // Edge cases that can make this true:
+            //   1. Account was destroyed and created again afterwards.
+            //   2. Account was destroyed but then was sent ETH, so it's not going to be completely removed from the trie.
+            let removed_storage = new_state_account.status == AccountStatus::DestroyedModified;
+
             // 2. Storage has been updated if the current value is different from the one before execution.
             let mut added_storage = BTreeMap::new();
 
             for (key, new_value) in &new_state_account.storage {
-                let old_value = initial_state_account.storage.get(key).ok_or_else(|| { VMError::Internal(InternalError::Custom(format!("Failed to get old value from account's initial storage for address: {address}")))})?;
+                let old_value = if !removed_storage {
+                    initial_state_account.storage.get(key).ok_or_else(|| { VMError::Internal(InternalError::Custom(format!("Failed to get old value from account's initial storage for address: {address}")))})?
+                } else {
+                    // There's not an "old value" if the contract was destroyed and re-created.
+                    &ZERO_U256
+                };
 
                 if new_value != old_value {
                     added_storage.insert(*key, *new_value);
@@ -232,11 +219,11 @@ impl GeneralizedDatabase {
             };
 
             // "At the end of the transaction, any account touched by the execution of that transaction which is now empty SHALL instead become non-existent (i.e. deleted)."
-            // If the account was already empty then this is not an update
+            // ethrex is a post-Merge client, empty accounts have already been pruned from the trie on Mainnet by the Merge (see EIP-161), so we won't have any empty accounts in the trie.
             let was_empty = initial_state_account.is_empty();
             let removed = new_state_account.is_empty() && !was_empty;
 
-            if !removed && !acc_info_updated && !storage_updated {
+            if !removed && !acc_info_updated && !storage_updated && !removed_storage {
                 // Account hasn't been updated
                 continue;
             }
@@ -247,6 +234,7 @@ impl GeneralizedDatabase {
                 info,
                 code: code.cloned(),
                 added_storage,
+                removed_storage,
             };
 
             account_updates.push(account_update);

tooling/ef_tests/state/runner/revm_db.rs

@@ -198,6 +198,7 @@ impl RevmState {
                                 )
                             })
                             .collect(),
+                        removed_storage: false,
                     };
                     account_updates.push(new_acc_update);
                 }