From 45e8cab4d072c3257d493221c1c167aca22588c9 Mon Sep 17 00:00:00 2001
From: Alex Miao <alex.miao@alchemy.com>
Date: Mon, 18 Sep 2023 11:03:54 -0700
Subject: [PATCH] refactor(entrypoint): replace get_deposit calls with
 balance_of

---
 crates/provider/src/ethers/entry_point.rs     | 25 ++++++++-----------
 crates/provider/src/traits/entry_point.rs     | 10 ++++----
 crates/rundler/src/builder/bundle_proposer.rs |  9 ++++---
 crates/rundler/src/common/precheck.rs         |  2 +-
 4 files changed, 23 insertions(+), 23 deletions(-)

diff --git a/crates/provider/src/ethers/entry_point.rs b/crates/provider/src/ethers/entry_point.rs
index 55a85f4ff..4731ee730 100644
--- a/crates/provider/src/ethers/entry_point.rs
+++ b/crates/provider/src/ethers/entry_point.rs
@@ -6,8 +6,8 @@ use ethers::{
     contract::{ContractError, FunctionCall},
     providers::{spoof, Middleware, RawCall},
     types::{
-        transaction::eip2718::TypedTransaction, Address, Bytes, Eip1559TransactionRequest, H256,
-        U256,
+        transaction::eip2718::TypedTransaction, Address, BlockId, Bytes, Eip1559TransactionRequest,
+        H256, U256,
     },
 };
 use rundler_types::{
@@ -72,23 +72,20 @@ where
         Err(error)?
     }
 
-    async fn balance_of(&self, address: Address) -> anyhow::Result<U256> {
-        self.balance_of(address)
+    async fn balance_of(
+        &self,
+        address: Address,
+        block_id: Option<BlockId>,
+    ) -> anyhow::Result<U256> {
+        block_id
+            .map_or(self.balance_of(address), |bid| {
+                self.balance_of(address).block(bid)
+            })
             .call()
             .await
             .context("entry point should return balance")
     }
 
-    async fn get_deposit(&self, address: Address, block_hash: H256) -> anyhow::Result<U256> {
-        let deposit_info = self
-            .get_deposit_info(address)
-            .block(block_hash)
-            .call()
-            .await
-            .context("entry point should return deposit info")?;
-        Ok(deposit_info.deposit.into())
-    }
-
     async fn call_spoofed_simulate_op(
         &self,
         op: UserOperation,
diff --git a/crates/provider/src/traits/entry_point.rs b/crates/provider/src/traits/entry_point.rs
index fe8e686e9..07792a624 100644
--- a/crates/provider/src/traits/entry_point.rs
+++ b/crates/provider/src/traits/entry_point.rs
@@ -1,4 +1,6 @@
-use ethers::types::{spoof, transaction::eip2718::TypedTransaction, Address, Bytes, H256, U256};
+use ethers::types::{
+    spoof, transaction::eip2718::TypedTransaction, Address, BlockId, Bytes, H256, U256,
+};
 #[cfg(feature = "test-utils")]
 use mockall::automock;
 use rundler_types::{
@@ -35,10 +37,8 @@ pub trait EntryPoint: Send + Sync + 'static {
     ) -> anyhow::Result<HandleOpsOut>;
 
     /// Get the balance of an address
-    async fn balance_of(&self, address: Address) -> anyhow::Result<U256>;
-
-    /// Get the deposit value of an address on the entry point contract
-    async fn get_deposit(&self, address: Address, block_hash: H256) -> anyhow::Result<U256>;
+    async fn balance_of(&self, address: Address, block_id: Option<BlockId>)
+        -> anyhow::Result<U256>;
 
     /// Call the entry point contract's `simulateValidation` function
     async fn simulate_validation(
diff --git a/crates/rundler/src/builder/bundle_proposer.rs b/crates/rundler/src/builder/bundle_proposer.rs
index 462ed7bd0..c92a6cc8f 100644
--- a/crates/rundler/src/builder/bundle_proposer.rs
+++ b/crates/rundler/src/builder/bundle_proposer.rs
@@ -6,7 +6,7 @@ use std::{
 };
 
 use anyhow::Context;
-use ethers::types::{Address, Bytes, H256, U256};
+use ethers::types::{Address, BlockId, Bytes, H256, U256};
 use futures::future;
 use linked_hash_map::LinkedHashMap;
 #[cfg(test)]
@@ -421,7 +421,10 @@ where
         block_hash: H256,
     ) -> anyhow::Result<HashMap<Address, U256>> {
         let futures = addreses.map(|address| async move {
-            let deposit = self.entry_point.get_deposit(address, block_hash).await?;
+            let deposit = self
+                .entry_point
+                .balance_of(address, Some(BlockId::Hash(block_hash)))
+                .await?;
             Ok::<_, anyhow::Error>((address, deposit))
         });
         let addresses_and_deposits = future::try_join_all(futures)
@@ -1222,7 +1225,7 @@ mod tests {
         }
         for deposit in mock_paymaster_deposits {
             entry_point
-                .expect_get_deposit()
+                .expect_balance_of()
                 .times(..=1)
                 .return_once(move |_, _| Ok(deposit));
         }
diff --git a/crates/rundler/src/common/precheck.rs b/crates/rundler/src/common/precheck.rs
index d95316d02..caf6fae38 100644
--- a/crates/rundler/src/common/precheck.rs
+++ b/crates/rundler/src/common/precheck.rs
@@ -274,7 +274,7 @@ impl<P: Provider, E: EntryPoint> PrecheckerImpl<P, E> {
             None => op.sender,
         };
         self.entry_point
-            .balance_of(payer)
+            .balance_of(payer, None)
             .await
             .context("precheck should get payer balance")
     }