refactor(storage): extract shared find_changeset_block_from_index algorithm [1/3] (#20542)

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

@@ -15,7 +15,9 @@ pub use static_file::{
 
 mod state;
 pub use state::{
-    historical::{HistoricalStateProvider, HistoricalStateProviderRef, LowestAvailableBlocks},
+    historical::{
+        HistoricalStateProvider, HistoricalStateProviderRef, HistoryInfo, LowestAvailableBlocks,
+    },
     latest::{LatestStateProvider, LatestStateProviderRef},
     overlay::{OverlayStateProvider, OverlayStateProviderFactory},
 };

crates/storage/provider/src/providers/state/historical.rs

@@ -32,6 +32,62 @@ use reth_trie_db::{
 
 use std::fmt::Debug;
 
+/// Result of a history lookup for an account or storage slot.
+///
+/// Indicates where to find the historical value for a given key at a specific block.
+#[derive(Debug, Eq, PartialEq)]
+pub enum HistoryInfo {
+    /// The key is written to, but only after our block (not yet written at the target block). Or
+    /// it has never been written.
+    NotYetWritten,
+    /// The chunk contains an entry for a write after our block at the given block number.
+    /// The value should be looked up in the changeset at this block.
+    InChangeset(u64),
+    /// The chunk does not contain an entry for a write after our block. This can only
+    /// happen if this is the last chunk, so we need to look in the plain state.
+    InPlainState,
+    /// The key may have been written, but due to pruning we may not have changesets and
+    /// history, so we need to make a plain state lookup.
+    MaybeInPlainState,
+}
+
+impl HistoryInfo {
+    /// Determines where to find the historical value based on computed shard lookup results.
+    ///
+    /// This is a pure function shared by both MDBX and `RocksDB` backends.
+    ///
+    /// # Arguments
+    /// * `found_block` - The block number from the shard lookup
+    /// * `is_before_first_write` - True if the target block is before the first write to this key.
+    ///   This should be computed as: `rank == 0 && found_block != Some(block_number) &&
+    ///   !has_previous_shard` where `has_previous_shard` comes from a lazy `cursor.prev()` check.
+    /// * `lowest_available` - Lowest block where history is available (pruning boundary)
+    pub const fn from_lookup(
+        found_block: Option<u64>,
+        is_before_first_write: bool,
+        lowest_available: Option<BlockNumber>,
+    ) -> Self {
+        if is_before_first_write {
+            if let (Some(_), Some(block_number)) = (lowest_available, found_block) {
+                // The key may have been written, but due to pruning we may not have changesets
+                // and history, so we need to make a changeset lookup.
+                return Self::InChangeset(block_number)
+            }
+            // The key is written to, but only after our block.
+            return Self::NotYetWritten
+        }
+
+        if let Some(block_number) = found_block {
+            // The chunk contains an entry for a write after our block, return it.
+            Self::InChangeset(block_number)
+        } else {
+            // The chunk does not contain an entry for a write after our block. This can only
+            // happen if this is the last chunk and so we need to look in the plain state.
+            Self::InPlainState
+        }
+    }
+}
+
 /// State provider for a given block number which takes a tx reference.
 ///
 /// Historical state provider accesses the state at the start of the provided block number.
@@ -53,14 +109,6 @@ pub struct HistoricalStateProviderRef<'b, Provider> {
     lowest_available_blocks: LowestAvailableBlocks,
 }
 
-#[derive(Debug, Eq, PartialEq)]
-pub enum HistoryInfo {
-    NotYetWritten,
-    InChangeset(u64),
-    InPlainState,
-    MaybeInPlainState,
-}
-
 impl<'b, Provider: DBProvider + BlockNumReader> HistoricalStateProviderRef<'b, Provider> {
     /// Create new `StateProvider` for historical block number
     pub fn new(provider: &'b Provider, block_number: BlockNumber) -> Self {
@@ -166,20 +214,20 @@ impl<'b, Provider: DBProvider + BlockNumReader> HistoricalStateProviderRef<'b, P
     {
         let mut cursor = self.tx().cursor_read::<T>()?;
 
-        // Lookup the history chunk in the history index. If they key does not appear in the
+        // Lookup the history chunk in the history index. If the key does not appear in the
         // index, the first chunk for the next key will be returned so we filter out chunks that
         // have a different key.
-        if let Some(chunk) = cursor.seek(key)?.filter(|(key, _)| key_filter(key)).map(|x| x.1 .0) {
+        if let Some(chunk) = cursor.seek(key)?.filter(|(key, _)| key_filter(key)).map(|x| x.1) {
             // Get the rank of the first entry before or equal to our block.
             let mut rank = chunk.rank(self.block_number);
 
             // Adjust the rank, so that we have the rank of the first entry strictly before our
             // block (not equal to it).
-            if rank.checked_sub(1).and_then(|rank| chunk.select(rank)) == Some(self.block_number) {
-                rank -= 1
-            };
+            if rank.checked_sub(1).and_then(|r| chunk.select(r)) == Some(self.block_number) {
+                rank -= 1;
+            }
 
-            let block_number = chunk.select(rank);
+            let found_block = chunk.select(rank);
 
             // If our block is before the first entry in the index chunk and this first entry
             // doesn't equal to our block, it might be before the first write ever. To check, we
@@ -187,27 +235,15 @@ impl<'b, Provider: DBProvider + BlockNumReader> HistoricalStateProviderRef<'b, P
             // This check is worth it, the `cursor.prev()` check is rarely triggered (the if will
             // short-circuit) and when it passes we save a full seek into the changeset/plain state
             // table.
-            if rank == 0 &&
-                block_number != Some(self.block_number) &&
-                !cursor.prev()?.is_some_and(|(key, _)| key_filter(&key))
-            {
-                if let (Some(_), Some(block_number)) = (lowest_available_block_number, block_number)
-                {
-                    // The key may have been written, but due to pruning we may not have changesets
-                    // and history, so we need to make a changeset lookup.
-                    Ok(HistoryInfo::InChangeset(block_number))
-                } else {
-                    // The key is written to, but only after our block.
-                    Ok(HistoryInfo::NotYetWritten)
-                }
-            } else if let Some(block_number) = block_number {
-                // The chunk contains an entry for a write after our block, return it.
-                Ok(HistoryInfo::InChangeset(block_number))
-            } else {
-                // The chunk does not contain an entry for a write after our block. This can only
-                // happen if this is the last chunk and so we need to look in the plain state.
-                Ok(HistoryInfo::InPlainState)
-            }
+            let is_before_first_write =
+                needs_prev_shard_check(rank, found_block, self.block_number) &&
+                    !cursor.prev()?.is_some_and(|(key, _)| key_filter(&key));
+
+            Ok(HistoryInfo::from_lookup(
+                found_block,
+                is_before_first_write,
+                lowest_available_block_number,
+            ))
         } else if lowest_available_block_number.is_some() {
             // The key may have been written, but due to pruning we may not have changesets and
             // history, so we need to make a plain state lookup.
@@ -524,8 +560,17 @@ impl LowestAvailableBlocks {
     }
 }
 
+/// Checks if a previous shard lookup is needed to determine if we're before the first write.
+///
+/// Returns `true` when `rank == 0` (first entry in shard) and the found block doesn't match
+/// the target block number. In this case, we need to check if there's a previous shard.
+fn needs_prev_shard_check(rank: u64, found_block: Option<u64>, block_number: BlockNumber) -> bool {
+    rank == 0 && found_block != Some(block_number)
+}
+
 #[cfg(test)]
 mod tests {
+    use super::needs_prev_shard_check;
     use crate::{
         providers::state::historical::{HistoryInfo, LowestAvailableBlocks},
         test_utils::create_test_provider_factory,
@@ -827,4 +872,28 @@ mod tests {
             Ok(HistoryInfo::MaybeInPlainState)
         ));
     }
+
+    #[test]
+    fn test_history_info_from_lookup() {
+        // Before first write, no pruning → not yet written
+        assert_eq!(HistoryInfo::from_lookup(Some(10), true, None), HistoryInfo::NotYetWritten);
+        assert_eq!(HistoryInfo::from_lookup(None, true, None), HistoryInfo::NotYetWritten);
+
+        // Before first write WITH pruning → check changeset (pruning may have removed history)
+        assert_eq!(HistoryInfo::from_lookup(Some(10), true, Some(5)), HistoryInfo::InChangeset(10));
+        assert_eq!(HistoryInfo::from_lookup(None, true, Some(5)), HistoryInfo::NotYetWritten);
+
+        // Not before first write → check changeset or plain state
+        assert_eq!(HistoryInfo::from_lookup(Some(10), false, None), HistoryInfo::InChangeset(10));
+        assert_eq!(HistoryInfo::from_lookup(None, false, None), HistoryInfo::InPlainState);
+    }
+
+    #[test]
+    fn test_needs_prev_shard_check() {
+        // Only needs check when rank == 0 and found_block != block_number
+        assert!(needs_prev_shard_check(0, Some(10), 5));
+        assert!(needs_prev_shard_check(0, None, 5));
+        assert!(!needs_prev_shard_check(0, Some(5), 5)); // found_block == block_number
+        assert!(!needs_prev_shard_check(1, Some(10), 5)); // rank > 0
+    }
 }