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/account.rs

@@ -53,6 +53,19 @@ impl From<GenesisAccount> for LevmAccount {
 }
 
 impl LevmAccount {
+    pub fn mark_destroyed(&mut self) {
+        self.status = AccountStatus::Destroyed;
+    }
+
+    pub fn mark_modified(&mut self) {
+        if self.status == AccountStatus::Unmodified {
+            self.status = AccountStatus::Modified;
+        }
+        if self.status == AccountStatus::Destroyed {
+            self.status = AccountStatus::DestroyedModified;
+        }
+    }
+
     pub fn has_nonce(&self) -> bool {
         self.info.nonce != 0
     }
@@ -69,11 +82,6 @@ impl LevmAccount {
         self.info.is_empty()
     }
 
-    /// Updates the account status.
-    pub fn update_status(&mut self, status: AccountStatus) {
-        self.status = status;
-    }
-
     /// Checks if the account is unmodified.
     pub fn is_unmodified(&self) -> bool {
         matches!(self.status, AccountStatus::Unmodified)
@@ -83,13 +91,13 @@ impl LevmAccount {
 #[derive(Clone, Default, Debug, PartialEq, Eq, Serialize, Deserialize)]
 pub enum AccountStatus {
     #[default]
+    /// Account was only read and not mutated at all.
     Unmodified,
+    /// Account accessed mutably, doesn't necessarily mean that its state has changed though but it could
     Modified,
     /// Contract executed a SELFDESTRUCT
     Destroyed,
-    /// Contract created via external transaction or CREATE/CREATE2
-    Created,
-    /// Contract has been destroyed and then re-created, usually with CREATE2
+    /// Contract has been destroyed and then modified
     /// This is a particular state because we'll still have in the Database the storage (trie) values but they are actually invalid.
-    DestroyedCreated,
+    DestroyedModified,
 }

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

@@ -1,15 +1,16 @@
 use std::collections::BTreeMap;
-use std::collections::HashSet;
 use std::sync::Arc;
 
 use bytes::Bytes;
 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;
+use crate::account::AccountStatus;
 use crate::account::LevmAccount;
 use crate::call_frame::CallFrameBackup;
 use crate::errors::InternalError;
@@ -29,10 +30,6 @@ pub struct GeneralizedDatabase {
     pub initial_accounts_state: CacheDB,
     pub codes: BTreeMap<H256, Bytes>,
     pub tx_backup: Option<CallFrameBackup>,
-    /// For keeping track of all destroyed accounts during block execution.
-    /// Used in get_state_transitions for edge case in which account is destroyed and re-created afterwards
-    /// In that scenario we want to remove the previous storage of the account but we still want the account to exist.
-    pub destroyed_accounts: HashSet<Address>,
 }
 
 impl GeneralizedDatabase {
@@ -42,7 +39,6 @@ impl GeneralizedDatabase {
             current_accounts_state: CacheDB::new(),
             initial_accounts_state: CacheDB::new(),
             tx_backup: None,
-            destroyed_accounts: HashSet::new(),
             codes: BTreeMap::new(),
         }
     }
@@ -65,7 +61,6 @@ impl GeneralizedDatabase {
             current_accounts_state: levm_accounts.clone(),
             initial_accounts_state: levm_accounts,
             tx_backup: None,
-            destroyed_accounts: HashSet::new(),
             codes,
         }
     }
@@ -93,7 +88,9 @@ impl GeneralizedDatabase {
     /// Gets mutable reference of an account
     /// Warning: Use directly only if outside of the EVM, otherwise use `vm.get_account_mut` because it contemplates call frame backups.
     pub fn get_account_mut(&mut self, address: Address) -> Result<&mut LevmAccount, InternalError> {
-        self.load_account(address)
+        let acc = self.load_account(address)?;
+        acc.mark_modified();
+        Ok(acc)
     }
 
     /// Gets code immutably given the code hash.
@@ -121,11 +118,6 @@ impl GeneralizedDatabase {
         address: Address,
         key: H256,
     ) -> Result<U256, InternalError> {
-        // If the account was destroyed then we cannot rely on the DB to obtain its previous value
-        // This is critical when executing blocks in batches, as an account may be destroyed and created within the same batch
-        if self.destroyed_accounts.contains(&address) {
-            return Ok(Default::default());
-        }
         let value = self.store.get_storage_value(address, key)?;
         // Account must already be in initial_accounts_state
         match self.initial_accounts_state.get_mut(&address) {
@@ -160,6 +152,10 @@ impl GeneralizedDatabase {
     pub fn get_state_transitions(&mut self) -> Result<Vec<AccountUpdate>, VMError> {
         let mut account_updates: Vec<AccountUpdate> = vec![];
         for (address, new_state_account) in self.current_accounts_state.iter() {
+            if new_state_account.is_unmodified() {
+                // Skip processing account that we know wasn't mutably accessed during execution
+                continue;
+            }
             // In case the account is not in immutable_cache (rare) we search for it in the actual database.
             let initial_state_account =
                 self.initial_accounts_state
@@ -168,29 +164,6 @@ impl GeneralizedDatabase {
                         "Failed to get account {address} from immutable cache",
                     ))))?;
 
-            // Edge case: Account was destroyed and created again afterwards with CREATE2.
-            if self.destroyed_accounts.contains(address) && !new_state_account.is_empty() {
-                // Push to account updates the removal of the account and then push the new state of the account.
-                // This is for clearing the account's storage when it was selfdestructed in the first place.
-                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;
 
@@ -216,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);
@@ -235,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 post-Merge client, empty accounts have already been pruned from the trie 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;
             }
@@ -250,6 +234,7 @@ impl GeneralizedDatabase {
                 info,
                 code: code.cloned(),
                 added_storage,
+                removed_storage,
             };
 
             account_updates.push(account_update);
@@ -404,6 +389,10 @@ impl<'a> VM<'a> {
             if let Some(value) = account.storage.get(&key) {
                 return Ok(*value);
             }
+            // If the account was destroyed and then created then we cannot rely on the DB to obtain storage values
+            if account.status == AccountStatus::DestroyedModified {
+                return Ok(U256::zero());
+            }
         } else {
             // When requesting storage of an account we should've previously requested and cached the account
             return Err(InternalError::AccountNotFound);

crates/vm/levm/src/hooks/default_hook.rs

@@ -234,7 +234,7 @@ pub fn delete_self_destruct_accounts(vm: &mut VM<'_>) -> Result<(), VMError> {
             .backup_account_info(*address, account_to_remove)?;
 
         *account_to_remove = LevmAccount::default();
-        vm.db.destroyed_accounts.insert(*address);
+        account_to_remove.mark_destroyed();
     }
 
     Ok(())

tooling/ef_tests/state/runner/revm_db.rs

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