Add in-memory block TD lookup support (#10052)

Co-authored-by: Dan Cline <6798349+Rjected@users.noreply.github.com>

crates/chain-state/src/in_memory.rs

@@ -95,6 +95,11 @@ impl InMemoryState {
         self.numbers.read().get(&number).and_then(|hash| self.blocks.read().get(hash).cloned())
     }
 
+    /// Returns the hash for a specific block number
+    pub(crate) fn hash_by_number(&self, number: u64) -> Option<B256> {
+        self.numbers.read().get(&number).cloned()
+    }
+
     /// Returns the current chain head state.
     pub(crate) fn head_state(&self) -> Option<Arc<BlockState>> {
         self.numbers
@@ -202,7 +207,12 @@ impl CanonicalInMemoryState {
         Self { inner: Arc::new(inner) }
     }
 
-    /// Returns in the header corresponding to the given hash.
+    /// Returns the block hash corresponding to the given number
+    pub fn hash_by_number(&self, number: u64) -> Option<B256> {
+        self.inner.in_memory_state.hash_by_number(number)
+    }
+
+    /// Returns the header corresponding to the given hash.
     pub fn header_by_hash(&self, hash: B256) -> Option<SealedHeader> {
         self.state_by_hash(hash).map(|block| block.block().block.header.clone())
     }

crates/storage/provider/src/providers/blockchain_provider.rs

@@ -184,7 +184,22 @@ where
     }
 
     fn header_td_by_number(&self, number: BlockNumber) -> ProviderResult<Option<U256>> {
-        self.database.header_td_by_number(number)
+        // If the TD is recorded on disk, we can just return that
+        if let Some(td) = self.database.header_td_by_number(number)? {
+            Ok(Some(td))
+        } else if self.canonical_in_memory_state.hash_by_number(number).is_some() {
+            // Otherwise, if the block exists in memory, we should return a TD for it.
+            //
+            // The canonical in memory state should only store post-merge blocks. Post-merge blocks
+            // have zero difficulty. This means we can use the total difficulty for the last
+            // persisted block number.
+            let last_persisted_block_number = self.database.last_block_number()?;
+            self.database.header_td_by_number(last_persisted_block_number)
+        } else {
+            // If the block does not exist in memory, and does not exist on-disk, we should not
+            // return a TD for it.
+            Ok(None)
+        }
     }
 
     fn headers_range(&self, range: impl RangeBounds<BlockNumber>) -> ProviderResult<Vec<Header>> {