fix(l2): fixes and cleanup of to_execution_db() (#2482)

**Motivation**

Prover's execution was failing because of wrong data in the ExecutionDB.
There were also some cases where data was missing but the current tests
didn't catch it.

**Description**

- fixes saving final storage values instead of initial ones.
- fixes saving only touched storage values, instead of read ones too.
- removes unused `new_store`
- simplifies code

Author: xqft

crates/vm/backends/levm/db.rs

@@ -13,8 +13,7 @@ use std::sync::{Arc, Mutex};
 #[derive(Clone)]
 pub struct DatabaseLogger {
     pub block_hashes_accessed: Arc<Mutex<HashMap<u64, CoreH256>>>,
-    pub accounts_accessed: Arc<Mutex<Vec<CoreAddress>>>,
-    pub storage_accessed: Arc<Mutex<HashMap<(CoreAddress, CoreH256), CoreU256>>>,
+    pub state_accessed: Arc<Mutex<HashMap<CoreAddress, Vec<CoreH256>>>>,
     pub code_accessed: Arc<Mutex<Vec<CoreH256>>>,
     pub store: Arc<dyn LevmDatabase>,
 }
@@ -23,8 +22,7 @@ impl DatabaseLogger {
     pub fn new(store: Arc<dyn LevmDatabase>) -> Self {
         Self {
             block_hashes_accessed: Arc::new(Mutex::new(HashMap::new())),
-            accounts_accessed: Arc::new(Mutex::new(Vec::new())),
-            storage_accessed: Arc::new(Mutex::new(HashMap::new())),
+            state_accessed: Arc::new(Mutex::new(HashMap::new())),
             code_accessed: Arc::new(Mutex::new(vec![])),
             store,
         }
@@ -36,12 +34,12 @@ impl LevmDatabase for DatabaseLogger {
         &self,
         address: CoreAddress,
     ) -> Result<ethrex_levm::AccountInfo, DatabaseError> {
-        let acc_info = self.store.get_account_info(address)?;
-        self.accounts_accessed
+        self.state_accessed
             .lock()
             .map_err(|_| DatabaseError::Custom("Could not lock mutex".to_string()))?
-            .push(address);
-        Ok(acc_info)
+            .entry(address)
+            .or_default();
+        self.store.get_account_info(address)
     }
 
     fn account_exists(&self, address: CoreAddress) -> bool {
@@ -53,12 +51,13 @@ impl LevmDatabase for DatabaseLogger {
         address: CoreAddress,
         key: CoreH256,
     ) -> Result<CoreU256, DatabaseError> {
-        let slot = self.store.get_storage_slot(address, key)?;
-        self.storage_accessed
+        self.state_accessed
             .lock()
             .map_err(|_| DatabaseError::Custom("Could not lock mutex".to_string()))?
-            .insert((address, key), slot);
-        Ok(slot)
+            .entry(address)
+            .and_modify(|keys| keys.push(key))
+            .or_insert(vec![key]);
+        self.store.get_storage_slot(address, key)
     }
 
     fn get_block_hash(&self, block_number: u64) -> Result<Option<CoreH256>, DatabaseError> {

crates/vm/backends/levm/mod.rs

@@ -404,119 +404,74 @@ impl LEVM {
             .map_err(Box::new)?
             .account_updates;
 
-        // index read and touched account addresses and storage keys
-        let account_accessed = logger_ref
-            .accounts_accessed
+        // index accessed account addresses and storage keys
+        let state_accessed = logger_ref
+            .state_accessed
             .lock()
             .map_err(|_| {
                 ExecutionDBError::Store(StoreError::Custom("Could not lock mutex".to_string()))
             })?
             .clone();
 
-        let mut index = Vec::with_capacity(account_accessed.len());
-        for address in account_accessed.iter() {
-            let storage_keys = logger_ref
-                .storage_accessed
-                .lock()
-                .map_err(|_| {
-                    ExecutionDBError::Store(StoreError::Custom("Could not lock mutex".to_string()))
-                })?
-                .iter()
-                .filter_map(
-                    |((addr, key), _)| {
-                        if *addr == *address {
-                            Some(key)
-                        } else {
-                            None
-                        }
-                    },
-                )
-                .map(|key| H256::from_slice(&key.to_fixed_bytes()))
-                .collect::<Vec<_>>();
-            index.push((address, storage_keys));
-        }
+        // fetch all read/written accounts from store
+        let accounts = state_accessed
+            .keys()
+            .chain(execution_updates.iter().map(|update| &update.address))
+            .filter_map(|address| {
+                store
+                    .get_account_info_by_hash(parent_hash, *address)
+                    .transpose()
+                    .map(|account| Ok((*address, account?)))
+            })
+            .collect::<Result<HashMap<_, _>, ExecutionDBError>>()?;
 
-        // fetch all read/written values from store
-        let cache_accounts = account_accessed.iter().filter_map(|address| {
-            // filter new accounts (accounts that didn't exist before) assuming our store is
-            // correct (based on the success of the pre-execution).
-            if let Ok(Some(info)) = db.store.get_account_info_by_hash(parent_hash, *address) {
-                Some((address, info))
+        // fetch all read/written code from store
+        let code_accessed = logger_ref
+            .code_accessed
+            .lock()
+            .map_err(|_| {
+                ExecutionDBError::Store(StoreError::Custom("Could not lock mutex".to_string()))
+            })?
+            .clone();
+        let code = accounts
+            .values()
+            .map(|account| account.code_hash)
+            .chain(code_accessed.into_iter())
+            .filter_map(|hash| {
+                store
+                    .get_account_code(hash)
+                    .transpose()
+                    .map(|account| Ok((hash, account?)))
+            })
+            .collect::<Result<HashMap<_, _>, ExecutionDBError>>()?;
+
+        // fetch all read/written storage from store
+        let added_storage = execution_updates.iter().filter_map(|update| {
+            if !update.added_storage.is_empty() {
+                let keys = update.added_storage.keys().cloned().collect::<Vec<_>>();
+                Some((update.address, keys))
             } else {
                 None
             }
         });
-        let accounts = cache_accounts
+        let storage = state_accessed
             .clone()
-            .map(|(address, _)| {
-                // return error if account is missing
-                let account = match db.store.get_account_info_by_hash(parent_hash, *address) {
-                    Ok(Some(some)) => Ok(some),
-                    Err(err) => Err(ExecutionDBError::Database(err)),
-                    Ok(None) => unreachable!(), // we are filtering out accounts that are not present
-                                                // in the store
-                };
-                Ok((*address, account?))
+            .into_iter()
+            .chain(added_storage)
+            .map(|(address, keys)| {
+                let keys: Result<HashMap<_, _>, ExecutionDBError> = keys
+                    .iter()
+                    .filter_map(|key| {
+                        store
+                            .get_storage_at_hash(parent_hash, address, *key)
+                            .transpose()
+                            .map(|value| Ok((*key, value?)))
+                    })
+                    .collect();
+                Ok((address, keys?))
             })
             .collect::<Result<HashMap<_, _>, ExecutionDBError>>()?;
-        let mut code: HashMap<H256, Bytes> = execution_updates
-            .clone()
-            .iter()
-            .map(|update| {
-                // return error if code is missing
-                let hash = update.info.clone().unwrap_or_default().code_hash;
-                Ok((
-                    hash,
-                    db.store
-                        .get_account_code(hash)?
-                        .ok_or(ExecutionDBError::NewMissingCode(hash))?,
-                ))
-            })
-            .collect::<Result<_, ExecutionDBError>>()?;
-
-        let mut code_accessed = HashMap::new();
-
-        {
-            let all_code_accessed = logger_ref.code_accessed.lock().map_err(|_| {
-                ExecutionDBError::Store(StoreError::Custom("Could not lock mutex".to_string()))
-            })?;
-            for code_hash in all_code_accessed.iter() {
-                code_accessed.insert(
-                    *code_hash,
-                    store
-                        .get_account_code(*code_hash)?
-                        .ok_or(ExecutionDBError::CodeNotFound(code_hash.0.into()))?
-                        .clone(),
-                );
-            }
-            code.extend(code_accessed);
-        }
 
-        let storage = execution_updates
-            .iter()
-            .map(|update| {
-                // return error if storage is missing
-                Ok((
-                    update.address,
-                    update
-                        .added_storage
-                        .keys()
-                        .map(|key| {
-                            let key = H256::from(key.to_fixed_bytes());
-                            let value = store
-                                .get_storage_at_hash(
-                                    block.header.compute_block_hash(),
-                                    update.address,
-                                    key,
-                                )
-                                .map_err(ExecutionDBError::Store)?
-                                .ok_or(ExecutionDBError::NewMissingStorage(update.address, key))?;
-                            Ok((key, value))
-                        })
-                        .collect::<Result<HashMap<_, _>, ExecutionDBError>>()?,
-                ))
-            })
-            .collect::<Result<HashMap<_, _>, ExecutionDBError>>()?;
         let block_hashes = logger_ref
             .block_hashes_accessed
             .lock()
@@ -528,23 +483,14 @@ impl LEVM {
             .map(|(num, hash)| (num, H256::from(hash.0)))
             .collect();
 
-        let new_store = store.clone();
-        new_store
-            .apply_account_updates(block.hash(), &execution_updates)
-            .await?;
-
         // get account proofs
-        let state_trie = new_store
+        let state_trie = store
             .state_trie(block.hash())?
             .ok_or(ExecutionDBError::NewMissingStateTrie(parent_hash))?;
         let parent_state_trie = store
             .state_trie(parent_hash)?
             .ok_or(ExecutionDBError::NewMissingStateTrie(parent_hash))?;
-        let hashed_addresses: Vec<_> = index
-            .clone()
-            .into_iter()
-            .map(|(address, _)| hash_address(address))
-            .collect();
+        let hashed_addresses: Vec<_> = state_accessed.keys().map(hash_address).collect();
         let initial_state_proofs = parent_state_trie.get_proofs(&hashed_addresses)?;
         let final_state_proofs: Vec<_> = hashed_addresses
             .iter()
@@ -563,14 +509,14 @@ impl LEVM {
         // get storage proofs
         let mut storage_proofs = HashMap::new();
         let mut final_storage_proofs = HashMap::new();
-        for (address, storage_keys) in index {
-            let Some(parent_storage_trie) = store.storage_trie(parent_hash, *address)? else {
+        for (address, storage_keys) in state_accessed {
+            let Some(parent_storage_trie) = store.storage_trie(parent_hash, address)? else {
                 // the storage of this account was empty or the account is newly created, either
                 // way the storage trie was initially empty so there aren't any proofs to add.
                 continue;
             };
-            let storage_trie = new_store.storage_trie(block.hash(), *address)?.ok_or(
-                ExecutionDBError::NewMissingStorageTrie(block.hash(), *address),
+            let storage_trie = store.storage_trie(block.hash(), address)?.ok_or(
+                ExecutionDBError::NewMissingStorageTrie(block.hash(), address),
             )?;
             let paths = storage_keys.iter().map(hash_key).collect::<Vec<_>>();
 
@@ -594,7 +540,7 @@ impl LEVM {
                 [initial_proofs.1, potential_child_nodes].concat(),
             );
 
-            storage_proofs.insert(*address, proofs);
+            storage_proofs.insert(address, proofs);
             final_storage_proofs.insert(address, final_proofs);
         }