From 1707a2135ea90a7f2560019ec8e3b951113c6cff Mon Sep 17 00:00:00 2001
From: chirag-bgh <chirag2bsingh@gmail.com>
Date: Sat, 22 Jun 2024 03:45:35 +0530
Subject: [PATCH] modify constructor to accept &self in block_on

---
 crates/revm/src/db/alloydb.rs  | 11 ++++-----
 crates/revm/src/db/ethersdb.rs | 42 ++++++++++++++++++----------------
 2 files changed, 27 insertions(+), 26 deletions(-)

diff --git a/crates/revm/src/db/alloydb.rs b/crates/revm/src/db/alloydb.rs
index 19216d59d5..487ba136c5 100644
--- a/crates/revm/src/db/alloydb.rs
+++ b/crates/revm/src/db/alloydb.rs
@@ -44,12 +44,12 @@ impl<T: Transport + Clone, N: Network, P: Provider<T, N>> AlloyDB<T, N, P> {
 
     /// Internal utility function that allows us to block on a future regardless of the runtime flavor.
     #[inline]
-    fn block_on<F>(f: F, handle: &Handle) -> F::Output
+    fn block_on<F>(&self, f: F) -> F::Output
     where
         F: std::future::Future + Send,
         F::Output: Send,
     {
-        tokio::task::block_in_place(move || handle.block_on(f))
+        tokio::task::block_in_place(move || self.handle.block_on(f))
     }
 
     /// Set the block number on which the queries will be based on.
@@ -82,7 +82,7 @@ impl<T: Transport + Clone, N: Network, P: Provider<T, N>> DatabaseRef for AlloyD
             )
         };
 
-        let (nonce, balance, code) = Self::block_on(f, &self.handle);
+        let (nonce, balance, code) = self.block_on(f);
 
         let balance = balance?;
         let code = Bytecode::new_raw(code?.0.into());
@@ -98,11 +98,10 @@ impl<T: Transport + Clone, N: Network, P: Provider<T, N>> DatabaseRef for AlloyD
             return Ok(KECCAK_EMPTY);
         }
 
-        let block = Self::block_on(
+        let block = self.block_on(
             self.provider
                 // SAFETY: We know number <= u64::MAX, so we can safely convert it to u64
                 .get_block_by_number(number.to::<u64>().into(), false),
-            &self.handle,
         )?;
         // SAFETY: If the number is given, the block is supposed to be finalized, so unwrapping is safe.
         Ok(B256::new(*block.unwrap().header.hash.unwrap()))
@@ -118,7 +117,7 @@ impl<T: Transport + Clone, N: Network, P: Provider<T, N>> DatabaseRef for AlloyD
             .provider
             .get_storage_at(address, index)
             .block_id(self.block_number);
-        let slot_val = Self::block_on(f.into_future(), &self.handle)?;
+        let slot_val = self.block_on(f.into_future())?;
         Ok(slot_val)
     }
 }
diff --git a/crates/revm/src/db/ethersdb.rs b/crates/revm/src/db/ethersdb.rs
index 21057ab5ea..d7c277265c 100644
--- a/crates/revm/src/db/ethersdb.rs
+++ b/crates/revm/src/db/ethersdb.rs
@@ -27,29 +27,33 @@ impl<M: Middleware> EthersDB<M> {
             Err(_) => return None,
         };
 
-        let block_number: Option<BlockId> = if block_number.is_some() {
-            block_number
+        if block_number.is_some() {
+            Some(Self {
+                client,
+                block_number,
+                handle,
+            })
         } else {
-            Some(BlockId::from(
-                Self::block_on(client.get_block_number(), &handle).ok()?,
-            ))
-        };
-
-        Some(Self {
-            client,
-            block_number,
-            handle,
-        })
+            let mut instance = Self {
+                client: client.clone(),
+                block_number: None,
+                handle,
+            };
+            instance.block_number = Some(BlockId::from(
+                instance.block_on(client.get_block_number()).ok()?,
+            ));
+            Some(instance)
+        }
     }
 
     /// Internal utility function to call tokio feature and wait for output
     #[inline]
-    fn block_on<F>(f: F, handle: &Handle) -> F::Output
+    fn block_on<F>(&self, f: F) -> F::Output
     where
         F: core::future::Future + Send,
         F::Output: Send,
     {
-        tokio::task::block_in_place(move || handle.block_on(f))
+        tokio::task::block_in_place(move || self.handle.block_on(f))
     }
 
     /// set block number on which upcoming queries will be based
@@ -71,7 +75,7 @@ impl<M: Middleware> DatabaseRef for EthersDB<M> {
             let code = self.client.get_code(add, self.block_number);
             tokio::join!(nonce, balance, code)
         };
-        let (nonce, balance, code) = Self::block_on(f, &self.handle);
+        let (nonce, balance, code) = self.block_on(f);
 
         let balance = U256::from_limbs(balance?.0);
         let nonce = nonce?.as_u64();
@@ -88,10 +92,8 @@ impl<M: Middleware> DatabaseRef for EthersDB<M> {
     fn storage_ref(&self, address: Address, index: U256) -> Result<U256, Self::Error> {
         let add = eH160::from(address.0 .0);
         let index = H256::from(index.to_be_bytes());
-        let slot_value: H256 = Self::block_on(
-            self.client.get_storage_at(add, index, self.block_number),
-            &self.handle,
-        )?;
+        let slot_value: H256 =
+            self.block_on(self.client.get_storage_at(add, index, self.block_number))?;
         Ok(U256::from_be_bytes(slot_value.to_fixed_bytes()))
     }
 
@@ -103,7 +105,7 @@ impl<M: Middleware> DatabaseRef for EthersDB<M> {
         // We know number <= u64::MAX so unwrap is safe
         let number = eU64::from(u64::try_from(number).unwrap());
         let block: Option<Block<TxHash>> =
-            Self::block_on(self.client.get_block(BlockId::from(number)), &self.handle)?;
+            self.block_on(self.client.get_block(BlockId::from(number)))?;
         // If number is given, the block is supposed to be finalized so unwrap is safe too.
         Ok(B256::new(block.unwrap().hash.unwrap().0))
     }