Improve EthersDB

crates/revm/src/db/ethersdb.rs

@@ -1,10 +1,44 @@
-use crate::primitives::{AccountInfo, Address, Bytecode, B256, KECCAK_EMPTY, U256};
-use crate::{Database, DatabaseRef};
+use std::sync::Arc;
+
 use ethers_core::types::{BlockId, H160 as eH160, H256, U64 as eU64};
 use ethers_providers::Middleware;
-use std::sync::Arc;
 use tokio::runtime::{Builder, Handle, RuntimeFlavor};
 
+use crate::primitives::{AccountInfo, Address, Bytecode, B256, KECCAK_EMPTY, U256};
+use crate::{Database, DatabaseRef};
+
+/// internal utility function to call tokio feature and wait for output
+#[inline]
+fn block_on<F>(f: F) -> F::Output
+where
+    F: core::future::Future + Send,
+    F::Output: Send,
+{
+    match Handle::try_current() {
+        Ok(handle) => match handle.runtime_flavor() {
+            // This essentially equals to tokio::task::spawn_blocking because tokio doesn't
+            // allow current_thread runtime to block_in_place
+            RuntimeFlavor::CurrentThread => std::thread::scope(move |s| {
+                s.spawn(move || {
+                    Builder::new_current_thread()
+                        .enable_all()
+                        .build()
+                        .unwrap()
+                        .block_on(f)
+                })
+                .join()
+                .unwrap()
+            }),
+            _ => tokio::task::block_in_place(move || handle.block_on(f)),
+        },
+        Err(_) => Builder::new_current_thread()
+            .enable_all()
+            .build()
+            .unwrap()
+            .block_on(f),
+    }
+}
+
 #[derive(Debug, Clone)]
 pub struct EthersDB<M: Middleware> {
     client: Arc<M>,
@@ -14,52 +48,21 @@ pub struct EthersDB<M: Middleware> {
 impl<M: Middleware> EthersDB<M> {
     /// create ethers db connector inputs are url and block on what we are basing our database (None for latest)
     pub fn new(client: Arc<M>, block_number: Option<BlockId>) -> Option<Self> {
-        let client = client;
-        let mut out = Self {
-            client,
-            block_number: None,
-        };
-
-        out.block_number = if block_number.is_some() {
+        let block_number: Option<BlockId> = if block_number.is_some() {
             block_number
         } else {
-            Some(BlockId::from(
-                out.block_on(out.client.get_block_number()).ok()?,
-            ))
+            Some(BlockId::from(block_on(client.get_block_number()).ok()?))
         };
 
-        Some(out)
+        Some(Self {
+            client,
+            block_number,
+        })
     }
 
-    /// internal utility function to call tokio feature and wait for output
-    fn block_on<F>(&self, f: F) -> F::Output
-    where
-        F: core::future::Future + Send,
-        F::Output: Send,
-    {
-        match Handle::try_current() {
-            Ok(handle) => match handle.runtime_flavor() {
-                // This essentially equals to tokio::task::spawn_blocking because tokio doesn't
-                // allow current_thread runtime to block_in_place
-                RuntimeFlavor::CurrentThread => std::thread::scope(move |s| {
-                    s.spawn(move || {
-                        Builder::new_current_thread()
-                            .enable_all()
-                            .build()
-                            .unwrap()
-                            .block_on(f)
-                    })
-                    .join()
-                    .unwrap()
-                }),
-                _ => tokio::task::block_in_place(move || handle.block_on(f)),
-            },
-            Err(_) => Builder::new_current_thread()
-                .enable_all()
-                .build()
-                .unwrap()
-                .block_on(f),
-        }
+    /// set block number on which upcoming queries will be based
+    pub fn set_block_number(&mut self, block_number: BlockId) {
+        self.block_number = Some(block_number);
     }
 }
 
@@ -75,7 +78,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);
+        let (nonce, balance, code) = block_on(f);
 
         let balance = U256::from_limbs(balance?.0);
         let nonce = nonce?.as_u64();
@@ -99,19 +102,19 @@ impl<M: Middleware> DatabaseRef for EthersDB<M> {
                 .await?;
             Ok(U256::from_be_bytes(storage.to_fixed_bytes()))
         };
-        self.block_on(f)
+        block_on(f)
     }
 
     fn block_hash_ref(&self, number: U256) -> Result<B256, Self::Error> {
         // saturate usize
         if number > U256::from(u64::MAX) {
             return Ok(KECCAK_EMPTY);
         }
-        // We known number <= u64::MAX so unwrap is safe
+        // We know number <= u64::MAX so unwrap is safe
         let number = eU64::from(u64::try_from(number).unwrap());
         let f = async { self.client.get_block(BlockId::from(number)).await };
         // If number is given, the block is supposed to be finalized so unwrap is safe too.
-        Ok(B256::new(self.block_on(f)?.unwrap().hash.unwrap().0))
+        Ok(B256::new(block_on(f)?.unwrap().hash.unwrap().0))
     }
 }