From d963b0284dd53dd7fbbe1e467e56bc351f0d65e3 Mon Sep 17 00:00:00 2001
From: James <james@prestwi.ch>
Date: Sat, 14 Sep 2024 14:22:17 -0400
Subject: [PATCH] fix: remove dbgs, improve imports

---
 src/evm.rs            |  5 ++---
 src/ext.rs            | 17 +++++++++--------
 src/system/eip2935.rs |  3 ---
 src/test_utils.rs     |  7 ++++++-
 4 files changed, 17 insertions(+), 15 deletions(-)

diff --git a/src/evm.rs b/src/evm.rs
index 57154a6..b622fcf 100644
--- a/src/evm.rs
+++ b/src/evm.rs
@@ -6,12 +6,12 @@ use crate::{
 };
 use alloy_primitives::{Address, Bytes, U256};
 use revm::{
-    db::{states::bundle_state::BundleRetention, BundleState},
+    db::{states::bundle_state::BundleRetention, BundleState, State},
     primitives::{
         AccountInfo, Bytecode, EVMError, EvmState, ExecutionResult, InvalidTransaction,
         ResultAndState, SpecId,
     },
-    Database, DatabaseCommit, DatabaseRef, Evm, State,
+    Database, DatabaseCommit, DatabaseRef, Evm,
 };
 use std::{convert::Infallible, fmt};
 
@@ -892,7 +892,6 @@ impl<'a, Ext, Db: Database> EvmNeedsBlock<'a, Ext, State<Db>> {
     /// See [`State::merge_transitions`] and [`State::take_bundle`].
     pub fn finish(self) -> BundleState {
         let Self { inner: mut evm, .. } = self;
-
         evm.db_mut().merge_transitions(BundleRetention::Reverts);
         let bundle = evm.db_mut().take_bundle();
 
diff --git a/src/ext.rs b/src/ext.rs
index ddecd8d..18ccea4 100644
--- a/src/ext.rs
+++ b/src/ext.rs
@@ -2,7 +2,7 @@ use std::collections::HashMap;
 
 use alloy_primitives::{Address, B256, U256};
 use revm::{
-    primitives::{Account, AccountInfo, Bytecode, EvmStorageSlot},
+    primitives::{Account, AccountInfo, Bytecode, EvmState, EvmStorageSlot},
     Database, DatabaseCommit,
 };
 
@@ -48,28 +48,29 @@ pub trait EvmExtUnchecked<Db: Database> {
     {
         let mut state = HashMap::default();
         let mut old = HashMap::default();
-        for account in changes.into_iter() {
-            let mut acct = self.account(account)?;
-            old.insert(account, acct.info.clone());
+        for addr in changes.into_iter() {
+            let mut acct = self.account(addr)?;
+            old.insert(addr, acct.info.clone());
             f(&mut acct.info);
             acct.mark_touch();
-            state.insert(account, acct);
+            state.insert(addr, acct);
         }
         self.db_mut_ext().commit(state);
         Ok(old)
     }
 
     /// Modify an account with a closure and commit the modified account.
-    fn modify_account<F>(&mut self, address: Address, f: F) -> Result<AccountInfo, Db::Error>
+    fn modify_account<F>(&mut self, addr: Address, f: F) -> Result<AccountInfo, Db::Error>
     where
         F: FnOnce(&mut AccountInfo),
         Db: DatabaseCommit,
     {
-        let mut acct = self.account(address)?;
+        let mut acct = self.account(addr)?;
         let old = acct.info.clone();
         f(&mut acct.info);
         acct.mark_touch();
-        let changes = [(address, acct)].into_iter().collect();
+
+        let changes: EvmState = [(addr, acct)].into_iter().collect();
         self.db_mut_ext().commit(changes);
         Ok(old)
     }
diff --git a/src/system/eip2935.rs b/src/system/eip2935.rs
index 13021c6..11794a1 100644
--- a/src/system/eip2935.rs
+++ b/src/system/eip2935.rs
@@ -46,9 +46,6 @@ impl<'a, Ext, Db: Database + DatabaseCommit> EvmNeedsTx<'a, Ext, Db> {
             .block_hash(prev_block)
             .map_err(EVMError::Database)?;
 
-        dbg!(slot);
-        dbg!(parent_block_hash);
-
         self.try_set_storage_unchecked(HISTORY_STORAGE_ADDRESS, slot, parent_block_hash.into())
             .map_err(EVMError::Database)?;
 
diff --git a/src/test_utils.rs b/src/test_utils.rs
index 6911348..dd70df3 100644
--- a/src/test_utils.rs
+++ b/src/test_utils.rs
@@ -1,7 +1,7 @@
 use crate::{EvmNeedsCfg, Trevm, TrevmBuilder};
 use alloy_primitives::{Address, U256};
 use revm::{
-    db::{CacheDB, EmptyDB, InMemoryDB},
+    db::{CacheDB, EmptyDB, InMemoryDB, State},
     inspector_handle_register,
     inspectors::TracerEip3155,
     primitives::{AccountInfo, Bytecode},
@@ -106,6 +106,11 @@ pub fn test_trevm() -> EvmNeedsCfg<'static, (), CacheDB<EmptyDB>> {
     EvmBuilder::default().with_db(CacheDB::new(EmptyDB::default())).build_trevm()
 }
 
+/// Make a new [`Trevm`] with an in-memory database wrapped in a [`State`].
+pub fn test_state_trevm() -> EvmNeedsCfg<'static, (), State<EmptyDB>> {
+    EvmBuilder::default().with_db(State::builder().with_bundle_update().build()).build_trevm()
+}
+
 /// Make a new [`Trevm`] with an in-memory database and a tracer inspector.
 /// The tracer will print all EVM instructions to stdout.
 pub fn test_trevm_tracing() -> EvmNeedsCfg<'static, TracerEip3155, CacheDB<EmptyDB>> {