refactor(levm): use AccountStatus for managing status of accounts

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,5 +1,4 @@
 use std::collections::BTreeMap;
-use std::collections::HashSet;
 use std::sync::Arc;
 
 use bytes::Bytes;
@@ -10,6 +9,7 @@ use ethrex_common::types::Account;
 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 +29,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 +38,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 +60,6 @@ impl GeneralizedDatabase {
             current_accounts_state: levm_accounts.clone(),
             initial_accounts_state: levm_accounts,
             tx_backup: None,
-            destroyed_accounts: HashSet::new(),
             codes,
         }
     }
@@ -93,7 +87,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 +117,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 +151,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,10 +163,12 @@ 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() {
+            // 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.
-                // 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,
@@ -404,6 +401,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(())