refactor(rpc): simplify state_at functions (#1771)

crates/rpc/rpc/src/eth/api/call.rs

@@ -46,7 +46,7 @@ where
         state_overrides: Option<StateOverride>,
     ) -> EthResult<(ResultAndState, Env)> {
         let (cfg, block_env, at) = self.evm_env_at(at).await?;
-        let state = self.state_at_block_id(at)?.ok_or_else(|| EthApiError::UnknownBlockNumber)?;
+        let state = self.state_at(at)?;
         self.call_with(cfg, block_env, request, state, state_overrides)
     }
 
@@ -86,7 +86,7 @@ where
         at: BlockId,
     ) -> EthResult<U256> {
         let (cfg, block_env, at) = self.evm_env_at(at).await?;
-        let state = self.state_at_block_id(at)?.ok_or_else(|| EthApiError::UnknownBlockNumber)?;
+        let state = self.state_at(at)?;
         self.estimate_gas_with(cfg, block_env, request, state)
     }
 
@@ -272,7 +272,7 @@ where
     ) -> EthResult<AccessList> {
         let block_id = at.unwrap_or(BlockId::Number(BlockNumberOrTag::Latest));
         let (mut cfg, block, at) = self.evm_env_at(block_id).await?;
-        let state = self.state_at_block_id(at)?.ok_or_else(|| EthApiError::UnknownBlockNumber)?;
+        let state = self.state_at(at)?;
 
         // we want to disable this in eth_call, since this is common practice used by other node
         // impls and providers <https://github.com/foundry-rs/foundry/issues/4388>

crates/rpc/rpc/src/eth/api/mod.rs

@@ -19,6 +19,7 @@ mod server;
 mod sign;
 mod state;
 mod transactions;
+use crate::eth::error::{EthApiError, EthResult};
 pub use transactions::{EthTransactions, TransactionSource};
 
 /// Cache limit of block-level fee history for `eth_feeHistory` RPC method.
@@ -100,23 +101,25 @@ where
         self.client().convert_block_number(num)
     }
 
+    /// Returns the state at the given [BlockId] enum.
+    pub(crate) fn state_at_block_id(&self, at: BlockId) -> EthResult<ChainState<'_>> {
+        match at {
+            BlockId::Hash(hash) => Ok(self.state_at_hash(hash.into()).map(ChainState::boxed)?),
+            BlockId::Number(num) => {
+                self.state_at_block_number(num)?.ok_or(EthApiError::UnknownBlockNumber)
+            }
+        }
+    }
+
     /// Returns the state at the given [BlockId] enum or the latest.
     pub(crate) fn state_at_block_id_or_latest(
         &self,
         block_id: Option<BlockId>,
-    ) -> Result<Option<ChainState<'_>>> {
+    ) -> EthResult<ChainState<'_>> {
         if let Some(block_id) = block_id {
             self.state_at_block_id(block_id)
         } else {
-            self.latest_state().map(ChainState::boxed).map(Some)
+            Ok(self.latest_state().map(ChainState::boxed)?)
-        }
-    }
-
-    /// Returns the state at the given [BlockId] enum.
-    pub(crate) fn state_at_block_id(&self, block_id: BlockId) -> Result<Option<ChainState<'_>>> {
-        match block_id {
-            BlockId::Hash(hash) => self.state_at_hash(hash.into()).map(ChainState::boxed).map(Some),
-            BlockId::Number(num) => self.state_at_block_number(num),
         }
     }
 

crates/rpc/rpc/src/eth/api/state.rs

@@ -15,15 +15,13 @@ where
     Client: BlockProvider + StateProviderFactory + EvmEnvProvider + 'static,
 {
     pub(crate) fn get_code(&self, address: Address, block_id: Option<BlockId>) -> EthResult<Bytes> {
-        let state =
+        let state = self.state_at_block_id_or_latest(block_id)?;
-            self.state_at_block_id_or_latest(block_id)?.ok_or(EthApiError::UnknownBlockNumber)?;
         let code = state.account_code(address)?.unwrap_or_default();
         Ok(code.original_bytes().into())
     }
 
     pub(crate) fn balance(&self, address: Address, block_id: Option<BlockId>) -> EthResult<U256> {
-        let state =
+        let state = self.state_at_block_id_or_latest(block_id)?;
-            self.state_at_block_id_or_latest(block_id)?.ok_or(EthApiError::UnknownBlockNumber)?;
         let balance = state.account_balance(address)?.unwrap_or_default();
         Ok(balance)
     }
@@ -33,8 +31,7 @@ where
         address: Address,
         block_id: Option<BlockId>,
     ) -> EthResult<U256> {
-        let state =
+        let state = self.state_at_block_id_or_latest(block_id)?;
-            self.state_at_block_id_or_latest(block_id)?.ok_or(EthApiError::UnknownBlockNumber)?;
         let nonce = U256::from(state.account_nonce(address)?.unwrap_or_default());
         Ok(nonce)
     }
@@ -45,8 +42,7 @@ where
         index: U256,
         block_id: Option<BlockId>,
     ) -> EthResult<H256> {
-        let state =
+        let state = self.state_at_block_id_or_latest(block_id)?;
-            self.state_at_block_id_or_latest(block_id)?.ok_or(EthApiError::UnknownBlockNumber)?;
         let storage_key = H256(index.to_be_bytes());
         let value = state.storage(address, storage_key)?.unwrap_or_default();
         Ok(H256(value.to_be_bytes()))
@@ -65,9 +61,7 @@ where
             return Err(EthApiError::InvalidBlockRange)
         }
 
-        let state = self
+        let state = self.state_at_block_id(block_id)?;
-            .state_at_block_id_or_latest(Some(block_id))?
-            .ok_or(EthApiError::UnknownBlockNumber)?;
 
         let (account_proof, storage_hash, stg_proofs) = state.proof(address, &keys)?;
 

crates/rpc/rpc/src/eth/api/transactions.rs

@@ -17,6 +17,9 @@ use revm::primitives::{BlockEnv, CfgEnv};
 /// Commonly used transaction related functions for the [EthApi] type in the `eth_` namespace
 #[async_trait::async_trait]
 pub trait EthTransactions: Send + Sync {
+    /// Returns the state at the given [BlockId]
+    fn state_at(&self, at: BlockId) -> EthResult<ChainState<'_>>;
+
     /// Executes the closure with the state that corresponds to the given [BlockId].
     fn with_state_at<F, T>(&self, _at: BlockId, _f: F) -> EthResult<T>
     where
@@ -49,11 +52,15 @@ where
     Client: BlockProvider + StateProviderFactory + EvmEnvProvider + 'static,
     Network: Send + Sync + 'static,
 {
+    fn state_at(&self, at: BlockId) -> EthResult<ChainState<'_>> {
+        self.state_at_block_id(at)
+    }
+
     fn with_state_at<F, T>(&self, at: BlockId, f: F) -> EthResult<T>
     where
         F: FnOnce(ChainState<'_>) -> EthResult<T>,
     {
-        let state = self.state_at_block_id(at)?.ok_or(EthApiError::UnknownBlockNumber)?;
+        let state = self.state_at(at)?;
         f(state)
     }