feat(stages): Simplify TD stage (#2344)

crates/consensus/auto-seal/src/lib.rs

@@ -60,7 +60,7 @@ impl Consensus for AutoSealConsensus {
 
     fn validate_header(
         &self,
-        _header: &SealedHeader,
+        _header: &Header,
         _total_difficulty: U256,
     ) -> Result<(), ConsensusError> {
         Ok(())

crates/consensus/beacon/src/beacon_consensus.rs

@@ -2,7 +2,7 @@
 use reth_consensus_common::validation;
 use reth_interfaces::consensus::{Consensus, ConsensusError};
 use reth_primitives::{
-    Chain, ChainSpec, Hardfork, SealedBlock, SealedHeader, EMPTY_OMMER_ROOT, U256,
+    Chain, ChainSpec, Hardfork, Header, SealedBlock, SealedHeader, EMPTY_OMMER_ROOT, U256,
 };
 use std::sync::Arc;
 
@@ -36,7 +36,7 @@ impl Consensus for BeaconConsensus {
 
     fn validate_header(
         &self,
-        header: &SealedHeader,
+        header: &Header,
         total_difficulty: U256,
     ) -> Result<(), ConsensusError> {
         if self.chain_spec.fork(Hardfork::Paris).active_at_ttd(total_difficulty, header.difficulty)
@@ -85,7 +85,7 @@ impl Consensus for BeaconConsensus {
 ///
 /// From yellow paper: extraData: An arbitrary byte array containing data relevant to this block.
 /// This must be 32 bytes or fewer; formally Hx.
-fn validate_header_extradata(header: &SealedHeader) -> Result<(), ConsensusError> {
+fn validate_header_extradata(header: &Header) -> Result<(), ConsensusError> {
     if header.extra_data.len() > 32 {
         Err(ConsensusError::ExtraDataExceedsMax { len: header.extra_data.len() })
     } else {

crates/interfaces/src/consensus.rs

@@ -1,6 +1,6 @@
 use async_trait::async_trait;
 use reth_primitives::{
-    BlockHash, BlockNumber, InvalidTransactionError, SealedBlock, SealedHeader, H256, U256,
+    BlockHash, BlockNumber, Header, InvalidTransactionError, SealedBlock, SealedHeader, H256, U256,
 };
 use std::fmt::Debug;
 
@@ -29,7 +29,7 @@ pub trait Consensus: Debug + Send + Sync {
     /// Some consensus engines may want to do additional checks here.
     fn validate_header(
         &self,
-        header: &SealedHeader,
+        header: &Header,
         total_difficulty: U256,
     ) -> Result<(), ConsensusError>;
 

crates/interfaces/src/test_utils/headers.rs

@@ -320,7 +320,7 @@ impl Consensus for TestConsensus {
 
     fn validate_header(
         &self,
-        header: &SealedHeader,
+        header: &Header,
         total_difficulty: U256,
     ) -> Result<(), ConsensusError> {
         if self.fail_validation() {

crates/stages/src/stages/total_difficulty.rs

@@ -60,7 +60,6 @@ impl<DB: Database> Stage<DB> for TotalDifficultyStage {
 
         // Acquire cursor over total difficulty and headers tables
         let mut cursor_td = tx.cursor_write::<tables::HeaderTD>()?;
-        let mut cursor_canonical = tx.cursor_read::<tables::CanonicalHeaders>()?;
         let mut cursor_headers = tx.cursor_read::<tables::Headers>()?;
 
         // Get latest total difficulty
@@ -72,24 +71,16 @@ impl<DB: Database> Stage<DB> for TotalDifficultyStage {
         let mut td: U256 = last_entry.1.into();
         debug!(target: "sync::stages::total_difficulty", ?td, block_number = last_header_number, "Last total difficulty entry");
 
-        // Acquire canonical walker
-        let walker = cursor_canonical.walk_range(range)?;
-
         // Walk over newly inserted headers, update & insert td
-        for entry in walker {
-            let (number, hash) = entry?;
-            let (_, header) =
-                cursor_headers.seek_exact(number)?.ok_or(ProviderError::Header { number })?;
-            let header = header.seal(hash);
+        for entry in cursor_headers.walk_range(range)? {
+            let (block_number, header) = entry?;
             td += header.difficulty;
 
             self.consensus
                 .validate_header(&header, td)
                 .map_err(|error| StageError::Validation { block: header.number, error })?;
-
-            cursor_td.append(number, td.into())?;
+            cursor_td.append(block_number, td.into())?;
         }
-
         info!(target: "sync::stages::total_difficulty", stage_progress = end_block, is_final_range, "Sync iteration finished");
         Ok(ExecOutput { stage_progress: end_block, done: is_final_range })
     }

crates/storage/db/src/tables/raw.rs

@@ -56,6 +56,12 @@ impl<K: Key> RawKey<K> {
     }
 }
 
+impl<K: Key> From<K> for RawKey<K> {
+    fn from(key: K) -> Self {
+        RawKey::new(key)
+    }
+}
+
 impl AsRef<[u8]> for RawKey<Vec<u8>> {
     fn as_ref(&self) -> &[u8] {
         &self.key