feat(rpc/trace): dry

Signed-off-by: jsvisa <[email protected]>
paradigmxyz · Jun 16, 2024 · 0a58b40 · 0a58b40
1 parent fa63334
commit 0a58b40
Showing 1 changed file with 51 additions and 65 deletions.
diff --git a/crates/rpc/rpc/src/trace.rs b/crates/rpc/rpc/src/trace.rs
@@ -7,7 +7,9 @@ use crate::eth::{
 use async_trait::async_trait;
 use jsonrpsee::core::RpcResult as Result;
 use reth_consensus_common::calc::{base_block_reward, block_reward, ommer_reward};
-use reth_primitives::{revm::env::tx_env_with_recovered, BlockId, Bytes, SealedHeader, B256, U256};
+use reth_primitives::{
+    revm::env::tx_env_with_recovered, BlockId, Bytes, SealedBlock, SealedHeader, B256, U256,
+};
 use reth_provider::{BlockReader, ChainSpecProvider, EvmEnvProvider, StateProviderFactory};
 use reth_revm::database::StateProviderDatabase;
 use reth_rpc_api::TraceApiServer;
@@ -315,42 +317,12 @@ where
             .collect::<Vec<_>>();
 
         // add reward traces
-        let mut reward_blocks = Vec::with_capacity(blocks.len());
-        for block_id in &blocks {
-            reward_blocks.push(self.inner.eth_api.block_by_id(block_id.number.into()));
-        }
+        let reward_blocks =
+            blocks.iter().map(|block| self.inner.eth_api.block_by_id(block.number.into()));
 
         for block in (futures::future::try_join_all(reward_blocks).await?).into_iter().flatten() {
-            if let Some(header_td) = self.provider().header_td(&block.header.hash())? {
-                if let Some(base_block_reward) = base_block_reward(
-                    self.provider().chain_spec().as_ref(),
-                    block.header.number,
-                    block.header.difficulty,
-                    header_td,
-                ) {
-                    let block_reward = block_reward(base_block_reward, block.ommers.len());
-                    all_traces.push(reward_trace(
-                        &block.header,
-                        RewardAction {
-                            author: block.header.beneficiary,
-                            reward_type: RewardType::Block,
-                            value: U256::from(block_reward),
-                        },
-                    ));
-
-                    for uncle in &block.ommers {
-                        let uncle_reward =
-                            ommer_reward(base_block_reward, block.header.number, uncle.number);
-                        all_traces.push(reward_trace(
-                            &block.header,
-                            RewardAction {
-                                author: uncle.beneficiary,
-                                reward_type: RewardType::Uncle,
-                                value: U256::from(uncle_reward),
-                            },
-                        ));
-                    }
-                }
+            if let Some(reward_traces) = self.extract_block_reward_traces(block).await? {
+                all_traces.extend(reward_traces)
             }
         }
 
@@ -402,36 +374,8 @@ where
             maybe_traces.map(|traces| traces.into_iter().flatten().collect::<Vec<_>>());
 
         if let (Some(block), Some(traces)) = (maybe_block, maybe_traces.as_mut()) {
-            if let Some(header_td) = self.provider().header_td(&block.header.hash())? {
-                if let Some(base_block_reward) = base_block_reward(
-                    self.provider().chain_spec().as_ref(),
-                    block.header.number,
-                    block.header.difficulty,
-                    header_td,
-                ) {
-                    let block_reward = block_reward(base_block_reward, block.ommers.len());
-                    traces.push(reward_trace(
-                        &block.header,
-                        RewardAction {
-                            author: block.header.beneficiary,
-                            reward_type: RewardType::Block,
-                            value: U256::from(block_reward),
-                        },
-                    ));
-
-                    for uncle in &block.ommers {
-                        let uncle_reward =
-                            ommer_reward(base_block_reward, block.header.number, uncle.number);
-                        traces.push(reward_trace(
-                            &block.header,
-                            RewardAction {
-                                author: uncle.beneficiary,
-                                reward_type: RewardType::Uncle,
-                                value: U256::from(uncle_reward),
-                            },
-                        ));
-                    }
-                }
+            if let Some(reward_traces) = self.extract_block_reward_traces(block).await? {
+                traces.extend(reward_traces);
             }
         }
 
@@ -525,6 +469,48 @@ where
             transactions,
         }))
     }
+
+    /// Extracts the reward traces for the given block.
+    async fn extract_block_reward_traces(
+        &self,
+        block: SealedBlock,
+    ) -> EthResult<Option<Vec<LocalizedTransactionTrace>>> {
+        if let Some(header_td) = self.provider().header_td(&block.header.hash())? {
+            let mut traces = Vec::new();
+            if let Some(base_block_reward) = base_block_reward(
+                self.provider().chain_spec().as_ref(),
+                block.header.number,
+                block.header.difficulty,
+                header_td,
+            ) {
+                let block_reward = block_reward(base_block_reward, block.ommers.len());
+                traces.push(reward_trace(
+                    &block.header,
+                    RewardAction {
+                        author: block.header.beneficiary,
+                        reward_type: RewardType::Block,
+                        value: U256::from(block_reward),
+                    },
+                ));
+
+                for uncle in &block.ommers {
+                    let uncle_reward =
+                        ommer_reward(base_block_reward, block.header.number, uncle.number);
+                    traces.push(reward_trace(
+                        &block.header,
+                        RewardAction {
+                            author: uncle.beneficiary,
+                            reward_type: RewardType::Uncle,
+                            value: U256::from(uncle_reward),
+                        },
+                    ));
+                }
+            }
+            Ok(Some(traces))
+        } else {
+            Ok(None)
+        }
+    }
 }
 
 #[async_trait]