diff --git a/Cargo.lock b/Cargo.lock index 029aa0c8a02..fee0e33b890 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3741,6 +3741,7 @@ dependencies = [ "async-trait", "auto_impl", "eyre", + "futures", "futures-util", "reth-db-api", "reth-exex", @@ -3750,7 +3751,9 @@ dependencies = [ "reth-provider", "reth-trie", "serde", + "test-case", "thiserror 2.0.16", + "tokio", ] [[package]] diff --git a/crates/exex/external-proofs/Cargo.toml b/crates/exex/external-proofs/Cargo.toml index e9c76888f5b..248330fb338 100644 --- a/crates/exex/external-proofs/Cargo.toml +++ b/crates/exex/external-proofs/Cargo.toml @@ -32,3 +32,18 @@ eyre.workspace = true futures-util.workspace = true serde.workspace = true thiserror.workspace = true +tokio.workspace = true + +# dev dependencies +[dev-dependencies] +tokio = { workspace = true, features = ["test-util", "rt-multi-thread", "macros"] } +test-case.workspace = true +futures.workspace = true + +[features] +test-utils = [ + "reth-primitives-traits/test-utils", + "reth-db-api/test-utils", + "reth-provider/test-utils", + "reth-trie/test-utils", +] diff --git a/crates/exex/external-proofs/src/in_memory.rs b/crates/exex/external-proofs/src/in_memory.rs new file mode 100644 index 00000000000..869e997e38a --- /dev/null +++ b/crates/exex/external-proofs/src/in_memory.rs @@ -0,0 +1,673 @@ +//! In-memory implementation of `ExternalStorage` for testing purposes. +//! +//! This module provides a complete in-memory implementation of the [`ExternalStorage`] trait +//! that can be used for testing and development. The implementation uses tokio async `RwLock` +//! for thread-safe concurrent access and stores all data in memory using `BTreeMap` collections. + +#![allow(dead_code, unreachable_pub)] + +use alloy_primitives::{map::HashMap, B256, U256}; +use async_trait::async_trait; +use reth_primitives_traits::Account; +use reth_trie::{updates::TrieUpdates, BranchNodeCompact, HashedPostState, Nibbles}; +use std::{collections::BTreeMap, sync::Arc}; +use tokio::sync::RwLock; + +use super::storage::{ + BlockStateDiff, ExternalHashedCursor, ExternalStorage, ExternalStorageError, + ExternalStorageResult, ExternalTrieCursor, +}; + +/// In-memory implementation of `ExternalStorage` for testing purposes +#[derive(Debug, Clone)] +pub struct InMemoryExternalStorage { + /// Shared state across all instances + inner: Arc>, +} + +#[derive(Debug, Default)] +struct InMemoryStorageInner { + /// Account trie branches: (`block_number`, path) -> `branch_node` + account_branches: BTreeMap<(u64, Nibbles), Option>, + + /// Storage trie branches: (`block_number`, `hashed_address`, path) -> `branch_node` + storage_branches: BTreeMap<(u64, B256, Nibbles), Option>, + + /// Hashed accounts: (`block_number`, `hashed_address`) -> account + hashed_accounts: BTreeMap<(u64, B256), Option>, + + /// Hashed storages: (`block_number`, `hashed_address`, `hashed_slot`) -> value + hashed_storages: BTreeMap<(u64, B256, B256), U256>, + + /// Trie updates by block number + trie_updates: BTreeMap, + + /// Post state by block number + post_states: BTreeMap, + + /// Earliest block number and hash + earliest_block: Option<(u64, B256)>, +} + +impl InMemoryStorageInner { + fn store_trie_updates(&mut self, block_number: u64, block_state_diff: BlockStateDiff) { + // Store account branch nodes + for (path, branch) in block_state_diff.trie_updates.account_nodes_ref() { + self.account_branches.insert((block_number, *path), Some(branch.clone())); + } + + // Store removed account nodes + let account_removals = block_state_diff + .trie_updates + .removed_nodes_ref() + .iter() + .filter_map(|n| { + (!block_state_diff.trie_updates.account_nodes_ref().contains_key(n)) + .then_some((n, None)) + }) + .collect::>(); + + for (path, branch) in account_removals { + self.account_branches.insert((block_number, *path), branch); + } + + // Store storage branch nodes and removals + for (address, storage_trie_updates) in block_state_diff.trie_updates.storage_tries_ref() { + // Store storage branch nodes + for (path, branch) in storage_trie_updates.storage_nodes_ref() { + self.storage_branches.insert((block_number, *address, *path), Some(branch.clone())); + } + + // Store removed storage nodes + let storage_removals = storage_trie_updates + .removed_nodes_ref() + .iter() + .filter_map(|n| { + (!storage_trie_updates.storage_nodes_ref().contains_key(n)).then_some((n, None)) + }) + .collect::>(); + + for (path, branch) in storage_removals { + self.storage_branches.insert((block_number, *address, *path), branch); + } + } + + for (address, account) in &block_state_diff.post_state.accounts { + self.hashed_accounts.insert((block_number, *address), *account); + } + + for (hashed_address, storage) in &block_state_diff.post_state.storages { + // Handle wiped storage: iterate all existing values and mark them as deleted + // This is an expensive operation and should never happen for blocks going forward. + if storage.wiped { + // Collect latest values for each slot up to the current block + let mut slot_to_latest: std::collections::BTreeMap = + std::collections::BTreeMap::new(); + + for ((block, address, slot), value) in &self.hashed_storages { + if *block < block_number && *address == *hashed_address { + if let Some((existing_block, _)) = slot_to_latest.get(slot) { + if *block > *existing_block { + slot_to_latest.insert(*slot, (*block, *value)); + } + } else { + slot_to_latest.insert(*slot, (*block, *value)); + } + } + } + + // Store zero values for all non-zero slots to mark them as deleted + for (slot, (_, value)) in slot_to_latest { + if !value.is_zero() { + self.hashed_storages + .insert((block_number, *hashed_address, slot), U256::ZERO); + } + } + } else { + for (slot, value) in &storage.storage { + self.hashed_storages.insert((block_number, *hashed_address, *slot), *value); + } + } + } + + self.trie_updates.insert(block_number, block_state_diff.trie_updates.clone()); + self.post_states.insert(block_number, block_state_diff.post_state.clone()); + } +} + +impl Default for InMemoryExternalStorage { + fn default() -> Self { + Self::new() + } +} + +impl InMemoryExternalStorage { + /// Create a new in-memory external storage instance + pub fn new() -> Self { + Self { inner: Arc::new(RwLock::new(InMemoryStorageInner::default())) } + } +} + +/// In-memory implementation of `ExternalTrieCursor` +#[derive(Debug)] +pub struct InMemoryTrieCursor { + /// Current position in the iteration (-1 means not positioned yet) + position: isize, + /// Sorted entries that match the query parameters + entries: Vec<(Nibbles, BranchNodeCompact)>, +} + +impl InMemoryTrieCursor { + fn new( + storage: &InMemoryStorageInner, + hashed_address: Option, + max_block_number: u64, + ) -> Self { + // Common logic: collect latest values for each path + let mut path_to_latest: std::collections::BTreeMap< + Nibbles, + (u64, Option), + > = std::collections::BTreeMap::new(); + + let mut collected_entries: Vec<(Nibbles, BranchNodeCompact)> = + if let Some(addr) = hashed_address { + // Storage trie cursor + for ((block, address, path), branch) in &storage.storage_branches { + if *block <= max_block_number && *address == addr { + if let Some((existing_block, _)) = path_to_latest.get(path) { + if *block > *existing_block { + path_to_latest.insert(*path, (*block, branch.clone())); + } + } else { + path_to_latest.insert(*path, (*block, branch.clone())); + } + } + } + + path_to_latest + .into_iter() + .filter_map(|(path, (_, branch))| branch.map(|b| (path, b))) + .collect() + } else { + // Account trie cursor + for ((block, path), branch) in &storage.account_branches { + if *block <= max_block_number { + if let Some((existing_block, _)) = path_to_latest.get(path) { + if *block > *existing_block { + path_to_latest.insert(*path, (*block, branch.clone())); + } + } else { + path_to_latest.insert(*path, (*block, branch.clone())); + } + } + } + + path_to_latest + .into_iter() + .filter_map(|(path, (_, branch))| branch.map(|b| (path, b))) + .collect() + }; + + // Sort by path for consistent ordering + collected_entries.sort_by(|(a, _), (b, _)| a.cmp(b)); + + Self { position: -1, entries: collected_entries } + } +} + +impl ExternalTrieCursor for InMemoryTrieCursor { + fn seek_exact( + &mut self, + path: Nibbles, + ) -> ExternalStorageResult> { + if let Some(pos) = self.entries.iter().position(|(p, _)| *p == path) { + self.position = pos as isize; + Ok(Some(self.entries[pos].clone())) + } else { + Ok(None) + } + } + + fn seek( + &mut self, + path: Nibbles, + ) -> ExternalStorageResult> { + if let Some(pos) = self.entries.iter().position(|(p, _)| *p >= path) { + self.position = pos as isize; + Ok(Some(self.entries[pos].clone())) + } else { + Ok(None) + } + } + + fn next(&mut self) -> ExternalStorageResult> { + self.position += 1; + if self.position >= 0 && (self.position as usize) < self.entries.len() { + Ok(Some(self.entries[self.position as usize].clone())) + } else { + Ok(None) + } + } + + fn current(&mut self) -> ExternalStorageResult> { + if self.position >= 0 && (self.position as usize) < self.entries.len() { + Ok(Some(self.entries[self.position as usize].0)) + } else { + Ok(None) + } + } +} + +/// In-memory implementation of `ExternalHashedCursor` for storage slots +#[derive(Debug)] +pub struct InMemoryStorageCursor { + /// Current position in the iteration (-1 means not positioned yet) + position: isize, + /// Sorted entries that match the query parameters + entries: Vec<(B256, U256)>, +} + +impl InMemoryStorageCursor { + fn new(storage: &InMemoryStorageInner, hashed_address: B256, max_block_number: u64) -> Self { + // Collect latest values for each slot + let mut slot_to_latest: std::collections::BTreeMap = + std::collections::BTreeMap::new(); + + for ((block, address, slot), value) in &storage.hashed_storages { + if *block <= max_block_number && *address == hashed_address { + if let Some((existing_block, _)) = slot_to_latest.get(slot) { + if *block > *existing_block { + slot_to_latest.insert(*slot, (*block, *value)); + } + } else { + slot_to_latest.insert(*slot, (*block, *value)); + } + } + } + + // Filter out zero values - they represent deleted/empty storage slots + let mut entries: Vec<(B256, U256)> = slot_to_latest + .into_iter() + .filter_map( + |(slot, (_, value))| { + if value.is_zero() { + None + } else { + Some((slot, value)) + } + }, + ) + .collect(); + + entries.sort_by_key(|(slot, _)| *slot); + + Self { position: -1, entries } + } +} + +impl ExternalHashedCursor for InMemoryStorageCursor { + type Value = U256; + + fn seek(&mut self, key: B256) -> ExternalStorageResult> { + if let Some(pos) = self.entries.iter().position(|(k, _)| *k >= key) { + self.position = pos as isize; + Ok(Some(self.entries[pos])) + } else { + Ok(None) + } + } + + fn next(&mut self) -> ExternalStorageResult> { + self.position += 1; + if self.position >= 0 && (self.position as usize) < self.entries.len() { + Ok(Some(self.entries[self.position as usize])) + } else { + Ok(None) + } + } +} + +/// In-memory implementation of `ExternalHashedCursor` for accounts +#[derive(Debug)] +pub struct InMemoryAccountCursor { + /// Current position in the iteration (-1 means not positioned yet) + position: isize, + /// Sorted entries that match the query parameters + entries: Vec<(B256, Account)>, +} + +impl InMemoryAccountCursor { + fn new(storage: &InMemoryStorageInner, max_block_number: u64) -> Self { + // Collect latest accounts for each address + let mut addr_to_latest: std::collections::BTreeMap)> = + std::collections::BTreeMap::new(); + + for ((block, address), account) in &storage.hashed_accounts { + if *block <= max_block_number { + if let Some((existing_block, _)) = addr_to_latest.get(address) { + if *block > *existing_block { + addr_to_latest.insert(*address, (*block, *account)); + } + } else { + addr_to_latest.insert(*address, (*block, *account)); + } + } + } + + let mut entries: Vec<(B256, Account)> = addr_to_latest + .into_iter() + .filter_map(|(address, (_, account))| account.map(|acc| (address, acc))) + .collect(); + + entries.sort_by_key(|(address, _)| *address); + + Self { position: -1, entries } + } +} + +impl ExternalHashedCursor for InMemoryAccountCursor { + type Value = Account; + + fn seek(&mut self, key: B256) -> ExternalStorageResult> { + if let Some(pos) = self.entries.iter().position(|(k, _)| *k >= key) { + self.position = pos as isize; + Ok(Some(self.entries[pos])) + } else { + Ok(None) + } + } + + fn next(&mut self) -> ExternalStorageResult> { + self.position += 1; + if self.position >= 0 && (self.position as usize) < self.entries.len() { + Ok(Some(self.entries[self.position as usize])) + } else { + Ok(None) + } + } +} + +#[async_trait] +impl ExternalStorage for InMemoryExternalStorage { + type TrieCursor = InMemoryTrieCursor; + type StorageCursor = InMemoryStorageCursor; + type AccountHashedCursor = InMemoryAccountCursor; + + async fn store_account_branches( + &self, + block_number: u64, + updates: Vec<(Nibbles, Option)>, + ) -> ExternalStorageResult<()> { + let mut inner = self.inner.write().await; + + for (path, branch) in updates { + inner.account_branches.insert((block_number, path), branch); + } + + Ok(()) + } + + async fn store_storage_branches( + &self, + block_number: u64, + hashed_address: B256, + items: Vec<(Nibbles, Option)>, + ) -> ExternalStorageResult<()> { + let mut inner = self.inner.write().await; + + for (path, branch) in items { + inner.storage_branches.insert((block_number, hashed_address, path), branch); + } + + Ok(()) + } + + async fn store_hashed_accounts( + &self, + accounts: Vec<(B256, Option)>, + block_number: u64, + ) -> ExternalStorageResult<()> { + let mut inner = self.inner.write().await; + + for (address, account) in accounts { + inner.hashed_accounts.insert((block_number, address), account); + } + + Ok(()) + } + + async fn store_hashed_storages( + &self, + hashed_address: B256, + storages: Vec<(B256, U256)>, + block_number: u64, + ) -> ExternalStorageResult<()> { + let mut inner = self.inner.write().await; + + for (slot, value) in storages { + inner.hashed_storages.insert((block_number, hashed_address, slot), value); + } + + Ok(()) + } + + async fn get_earliest_block_number(&self) -> ExternalStorageResult> { + let inner = self.inner.read().await; + Ok(inner.earliest_block) + } + + async fn get_latest_block_number(&self) -> ExternalStorageResult> { + let inner = self.inner.read().await; + // Find the latest block number from trie_updates + let latest_block = inner.trie_updates.keys().max().copied(); + if let Some(block) = latest_block { + // We don't have a hash stored, so return a default + Ok(Some((block, B256::ZERO))) + } else { + Ok(None) + } + } + + fn trie_cursor( + &self, + hashed_address: Option, + max_block_number: u64, + ) -> ExternalStorageResult { + // For synchronous methods, we need to try_read() and handle potential blocking + let inner = self + .inner + .try_read() + .map_err(|_| ExternalStorageError::Other(eyre::eyre!("Failed to acquire read lock")))?; + Ok(InMemoryTrieCursor::new(&inner, hashed_address, max_block_number)) + } + + fn storage_hashed_cursor( + &self, + hashed_address: B256, + max_block_number: u64, + ) -> ExternalStorageResult { + let inner = self + .inner + .try_read() + .map_err(|_| ExternalStorageError::Other(eyre::eyre!("Failed to acquire read lock")))?; + Ok(InMemoryStorageCursor::new(&inner, hashed_address, max_block_number)) + } + + fn account_hashed_cursor( + &self, + max_block_number: u64, + ) -> ExternalStorageResult { + let inner = self + .inner + .try_read() + .map_err(|_| ExternalStorageError::Other(eyre::eyre!("Failed to acquire read lock")))?; + Ok(InMemoryAccountCursor::new(&inner, max_block_number)) + } + + async fn store_trie_updates( + &self, + block_number: u64, + block_state_diff: BlockStateDiff, + ) -> ExternalStorageResult<()> { + let mut inner = self.inner.write().await; + + inner.store_trie_updates(block_number, block_state_diff); + + Ok(()) + } + + async fn fetch_trie_updates(&self, block_number: u64) -> ExternalStorageResult { + let inner = self.inner.read().await; + + let trie_updates = inner.trie_updates.get(&block_number).cloned().unwrap_or_default(); + let post_state = inner.post_states.get(&block_number).cloned().unwrap_or_default(); + + Ok(BlockStateDiff { trie_updates, post_state }) + } + + async fn prune_earliest_state( + &self, + new_earliest_block_number: u64, + diff: BlockStateDiff, + ) -> ExternalStorageResult<()> { + let mut inner = self.inner.write().await; + + let branches_diff = diff.trie_updates; + let leaves_diff = diff.post_state; + + // Apply branch updates to the earliest state (block 0) + for (path, branch) in &branches_diff.account_nodes { + inner.account_branches.insert((0, *path), Some(branch.clone())); + } + + // Remove pruned account branches + for path in &branches_diff.removed_nodes { + inner.account_branches.remove(&(0, *path)); + } + + // Apply storage trie updates + for (hashed_address, storage_updates) in &branches_diff.storage_tries { + for (path, branch) in &storage_updates.storage_nodes { + inner.storage_branches.insert((0, *hashed_address, *path), Some(branch.clone())); + } + + for path in &storage_updates.removed_nodes { + inner.storage_branches.remove(&(0, *hashed_address, *path)); + } + } + + // Apply account updates + for (hashed_address, account) in &leaves_diff.accounts { + inner.hashed_accounts.insert((0, *hashed_address), *account); + } + + // Apply storage updates + for (hashed_address, storage) in &leaves_diff.storages { + for (slot, value) in &storage.storage { + inner.hashed_storages.insert((0, *hashed_address, *slot), *value); + } + } + + // Update earliest block number if we have one + if let Some((_, hash)) = inner.earliest_block { + inner.earliest_block = Some((new_earliest_block_number, hash)); + } + + // Remove all data for blocks before new_earliest_block_number (except block 0) + inner + .account_branches + .retain(|(block, _), _| *block == 0 || *block >= new_earliest_block_number); + inner + .storage_branches + .retain(|(block, _, _), _| *block == 0 || *block >= new_earliest_block_number); + inner + .hashed_accounts + .retain(|(block, _), _| *block == 0 || *block >= new_earliest_block_number); + inner + .hashed_storages + .retain(|(block, _, _), _| *block == 0 || *block >= new_earliest_block_number); + inner.trie_updates.retain(|block, _| *block >= new_earliest_block_number); + inner.post_states.retain(|block, _| *block >= new_earliest_block_number); + + Ok(()) + } + + async fn replace_updates( + &self, + latest_common_block_number: u64, + blocks_to_add: HashMap, + ) -> ExternalStorageResult<()> { + let mut inner = self.inner.write().await; + + // Remove all updates after latest_common_block_number + inner.trie_updates.retain(|block, _| *block <= latest_common_block_number); + inner.post_states.retain(|block, _| *block <= latest_common_block_number); + inner.account_branches.retain(|(block, _), _| *block <= latest_common_block_number); + inner.storage_branches.retain(|(block, _, _), _| *block <= latest_common_block_number); + inner.hashed_accounts.retain(|(block, _), _| *block <= latest_common_block_number); + inner.hashed_storages.retain(|(block, _, _), _| *block <= latest_common_block_number); + + for (block_number, block_state_diff) in blocks_to_add { + inner.store_trie_updates(block_number, block_state_diff); + } + + Ok(()) + } + + async fn set_earliest_block_number( + &self, + block_number: u64, + hash: B256, + ) -> ExternalStorageResult<()> { + let mut inner = self.inner.write().await; + inner.earliest_block = Some((block_number, hash)); + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use alloy_primitives::U256; + use reth_primitives_traits::Account; + + #[tokio::test] + async fn test_in_memory_storage_basic_operations() -> Result<(), ExternalStorageError> { + let storage = InMemoryExternalStorage::new(); + + // Test setting earliest block + let block_hash = B256::random(); + storage.set_earliest_block_number(1, block_hash).await?; + let earliest = storage.get_earliest_block_number().await?; + assert_eq!(earliest, Some((1, block_hash))); + + // Test storing and retrieving accounts + let account = Account { nonce: 1, balance: U256::from(100), bytecode_hash: None }; + let hashed_address = B256::random(); + + storage.store_hashed_accounts(vec![(hashed_address, Some(account))], 2).await?; + + let _cursor = storage.account_hashed_cursor(10)?; + // Note: cursor testing would require more complex setup with proper seek/next operations + + Ok(()) + } + + #[tokio::test] + async fn test_trie_updates_storage() -> Result<(), ExternalStorageError> { + let storage = InMemoryExternalStorage::new(); + + let trie_updates = TrieUpdates::default(); + let post_state = HashedPostState::default(); + let block_state_diff = + BlockStateDiff { trie_updates: trie_updates.clone(), post_state: post_state.clone() }; + + storage.store_trie_updates(5, block_state_diff).await?; + + let retrieved_diff = storage.fetch_trie_updates(5).await?; + assert_eq!(retrieved_diff.trie_updates, trie_updates); + assert_eq!(retrieved_diff.post_state, post_state); + + Ok(()) + } +} diff --git a/crates/exex/external-proofs/src/lib.rs b/crates/exex/external-proofs/src/lib.rs index 6768373634f..5e37a662ae6 100644 --- a/crates/exex/external-proofs/src/lib.rs +++ b/crates/exex/external-proofs/src/lib.rs @@ -7,7 +7,10 @@ use reth_provider::StateReader; use reth_exex::{ExExContext, ExExEvent}; -mod storage; +pub mod in_memory; +pub mod storage; +#[cfg(test)] +mod storage_tests; /// Saves and serves trie nodes to make proofs faster. This handles the process of /// saving the current state, new blocks as they're added, and serving proof RPCs diff --git a/crates/exex/external-proofs/src/storage.rs b/crates/exex/external-proofs/src/storage.rs index 980f42feb01..71e5a5fdff1 100644 --- a/crates/exex/external-proofs/src/storage.rs +++ b/crates/exex/external-proofs/src/storage.rs @@ -1,4 +1,6 @@ -#![expect(dead_code, unreachable_pub)] +//! Traits for external storage for trie nodes. + +#![allow(dead_code, unreachable_pub)] use alloy_primitives::{map::HashMap, B256, U256}; use async_trait::async_trait; use auto_impl::auto_impl; @@ -57,7 +59,9 @@ pub trait ExternalHashedCursor: Send + Sync { /// Diff of trie updates and post state for a block. #[derive(Debug, Clone)] pub struct BlockStateDiff { + /// Trie updates for branch nodes pub trie_updates: TrieUpdates, + /// Post state for leaf nodes (accounts and storage) pub post_state: HashedPostState, } @@ -68,8 +72,13 @@ pub struct BlockStateDiff { #[async_trait] #[auto_impl(Arc)] pub trait ExternalStorage: Send + Sync + Debug { + /// Cursor for iterating over trie branches. type TrieCursor: ExternalTrieCursor; + + /// Cursor for iterating over storage leaves. type StorageCursor: ExternalHashedCursor; + + /// Cursor for iterating over account leaves. type AccountHashedCursor: ExternalHashedCursor; /// Store a batch of account trie branches. Used for saving existing state. For live state diff --git a/crates/exex/external-proofs/src/storage_tests.rs b/crates/exex/external-proofs/src/storage_tests.rs new file mode 100644 index 00000000000..f34d6a324fd --- /dev/null +++ b/crates/exex/external-proofs/src/storage_tests.rs @@ -0,0 +1,1625 @@ +//! Common test suite for `ExternalStorage` implementations. + +#[cfg(test)] +mod tests { + use crate::{ + in_memory::InMemoryExternalStorage, + storage::{ + BlockStateDiff, ExternalHashedCursor, ExternalStorage, ExternalStorageError, + ExternalTrieCursor, + }, + }; + use alloy_primitives::{map::HashMap, B256, U256}; + use reth_primitives_traits::Account; + use reth_trie::{updates::TrieUpdates, BranchNodeCompact, HashedPostState, Nibbles, TrieMask}; + use std::sync::Arc; + use test_case::test_case; + + /// Helper to create a simple test branch node + fn create_test_branch() -> BranchNodeCompact { + let mut state_mask = TrieMask::default(); + state_mask.set_bit(0); + state_mask.set_bit(1); + + BranchNodeCompact { + state_mask, + tree_mask: TrieMask::default(), + hash_mask: TrieMask::default(), + hashes: Arc::new(vec![]), + root_hash: None, + } + } + + /// Helper to create a variant test branch node for comparison tests + fn create_test_branch_variant() -> BranchNodeCompact { + let mut state_mask = TrieMask::default(); + state_mask.set_bit(5); + state_mask.set_bit(6); + + BranchNodeCompact { + state_mask, + tree_mask: TrieMask::default(), + hash_mask: TrieMask::default(), + hashes: Arc::new(vec![]), + root_hash: None, + } + } + + /// Helper to create nibbles from a vector of u8 values + fn nibbles_from(vec: Vec) -> Nibbles { + Nibbles::from_nibbles_unchecked(vec) + } + + /// Helper to create a test account + fn create_test_account() -> Account { + Account { + nonce: 42, + balance: U256::from(1000000), + bytecode_hash: Some(B256::repeat_byte(0xBB)), + } + } + + /// Helper to create a test account with custom values + fn create_test_account_with_values(nonce: u64, balance: u64, code_hash_byte: u8) -> Account { + Account { + nonce, + balance: U256::from(balance), + bytecode_hash: Some(B256::repeat_byte(code_hash_byte)), + } + } + + /// Test basic storage and retrieval of earliest block number + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_earliest_block_operations( + storage: S, + ) -> Result<(), ExternalStorageError> { + // Initially should be None + let earliest = storage.get_earliest_block_number().await?; + assert!(earliest.is_none()); + + // Set earliest block + let block_hash = B256::repeat_byte(0x42); + storage.set_earliest_block_number(100, block_hash).await?; + + // Should retrieve the same values + let earliest = storage.get_earliest_block_number().await?; + assert_eq!(earliest, Some((100, block_hash))); + + Ok(()) + } + + /// Test storing and retrieving trie updates + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_trie_updates_operations( + storage: S, + ) -> Result<(), ExternalStorageError> { + let block_number = 50; + let trie_updates = TrieUpdates::default(); + let post_state = HashedPostState::default(); + let block_state_diff = + BlockStateDiff { trie_updates: trie_updates.clone(), post_state: post_state.clone() }; + + // Store trie updates + storage.store_trie_updates(block_number, block_state_diff).await?; + + // Retrieve and verify + let retrieved_diff = storage.fetch_trie_updates(block_number).await?; + assert_eq!(retrieved_diff.trie_updates, trie_updates); + assert_eq!(retrieved_diff.post_state, post_state); + + Ok(()) + } + + // ============================================================================= + // 1. Basic Cursor Operations + // ============================================================================= + + /// Test cursor operations on empty trie + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_cursor_empty_trie( + storage: S, + ) -> Result<(), ExternalStorageError> { + let mut cursor = storage.trie_cursor(None, 100)?; + + // All operations should return None on empty trie + assert!(cursor.seek_exact(Nibbles::default())?.is_none()); + assert!(cursor.seek(Nibbles::default())?.is_none()); + assert!(cursor.next()?.is_none()); + assert!(cursor.current()?.is_none()); + + Ok(()) + } + + /// Test cursor operations with single entry + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_cursor_single_entry( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1, 2, 3]); + let branch = create_test_branch(); + + // Store single entry + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + + // Test seek_exact + let result = cursor.seek_exact(path)?.unwrap(); + assert_eq!(result.0, path); + + // Test current position + assert_eq!(cursor.current()?.unwrap(), path); + + // Test next from end should return None + assert!(cursor.next()?.is_none()); + + Ok(()) + } + + /// Test cursor operations with multiple entries + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_cursor_multiple_entries( + storage: S, + ) -> Result<(), ExternalStorageError> { + let paths = vec![ + nibbles_from(vec![1]), + nibbles_from(vec![1, 2]), + nibbles_from(vec![2]), + nibbles_from(vec![2, 3]), + ]; + let branch = create_test_branch(); + + // Store multiple entries + for path in &paths { + storage.store_account_branches(50, vec![(*path, Some(branch.clone()))]).await?; + } + + let mut cursor = storage.trie_cursor(None, 100)?; + + // Test that we can iterate through all entries + let mut found_paths = Vec::new(); + while let Some((path, _)) = cursor.next()? { + found_paths.push(path); + } + + assert_eq!(found_paths.len(), 4); + // Paths should be in lexicographic order + for i in 0..paths.len() { + assert_eq!(found_paths[i], paths[i]); + } + + Ok(()) + } + + // ============================================================================= + // 2. Seek Operations + // ============================================================================= + + /// Test `seek_exact` with existing path + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_seek_exact_existing_path( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1, 2, 3]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + let result = cursor.seek_exact(path)?.unwrap(); + assert_eq!(result.0, path); + + Ok(()) + } + + /// Test `seek_exact` with non-existing path + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_seek_exact_non_existing_path( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1, 2, 3]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + let non_existing = nibbles_from(vec![4, 5, 6]); + assert!(cursor.seek_exact(non_existing)?.is_none()); + + Ok(()) + } + + /// Test `seek_exact` with empty path + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_seek_exact_empty_path( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + let result = cursor.seek_exact(Nibbles::default())?.unwrap(); + assert_eq!(result.0, Nibbles::default()); + + Ok(()) + } + + /// Test seek to existing path + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_seek_to_existing_path( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1, 2, 3]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + let result = cursor.seek(path)?.unwrap(); + assert_eq!(result.0, path); + + Ok(()) + } + + /// Test seek between existing nodes + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_seek_between_existing_nodes( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path1 = nibbles_from(vec![1]); + let path2 = nibbles_from(vec![3]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path1, Some(branch.clone()))]).await?; + storage.store_account_branches(50, vec![(path2, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + // Seek to path between 1 and 3, should return path 3 + let seek_path = nibbles_from(vec![2]); + let result = cursor.seek(seek_path)?.unwrap(); + assert_eq!(result.0, path2); + + Ok(()) + } + + /// Test seek after all nodes + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_seek_after_all_nodes( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + // Seek to path after all nodes + let seek_path = nibbles_from(vec![9]); + assert!(cursor.seek(seek_path)?.is_none()); + + Ok(()) + } + + /// Test seek before all nodes + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_seek_before_all_nodes( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![5]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + // Seek to path before all nodes, should return first node + let seek_path = nibbles_from(vec![1]); + let result = cursor.seek(seek_path)?.unwrap(); + assert_eq!(result.0, path); + + Ok(()) + } + + // ============================================================================= + // 3. Navigation Tests + // ============================================================================= + + /// Test next without prior seek + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_next_without_prior_seek( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1, 2]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + // next() without prior seek should start from beginning + let result = cursor.next()?.unwrap(); + assert_eq!(result.0, path); + + Ok(()) + } + + /// Test next after seek + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_next_after_seek( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path1 = nibbles_from(vec![1]); + let path2 = nibbles_from(vec![2]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path1, Some(branch.clone()))]).await?; + storage.store_account_branches(50, vec![(path2, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + cursor.seek(path1)?; + + // next() should return second node + let result = cursor.next()?.unwrap(); + assert_eq!(result.0, path2); + + Ok(()) + } + + /// Test next at end of trie + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_next_at_end_of_trie( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + cursor.seek(path)?; + + // next() at end should return None + assert!(cursor.next()?.is_none()); + + Ok(()) + } + + /// Test multiple consecutive next calls + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_multiple_consecutive_next( + storage: S, + ) -> Result<(), ExternalStorageError> { + let paths = vec![nibbles_from(vec![1]), nibbles_from(vec![2]), nibbles_from(vec![3])]; + let branch = create_test_branch(); + + for path in &paths { + storage.store_account_branches(50, vec![(*path, Some(branch.clone()))]).await?; + } + + let mut cursor = storage.trie_cursor(None, 100)?; + + // Iterate through all with consecutive next() calls + for expected_path in &paths { + let result = cursor.next()?.unwrap(); + assert_eq!(result.0, *expected_path); + } + + // Final next() should return None + assert!(cursor.next()?.is_none()); + + Ok(()) + } + + /// Test current after operations + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_current_after_operations( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path1 = nibbles_from(vec![1]); + let path2 = nibbles_from(vec![2]); + let branch = create_test_branch(); + + storage.store_account_branches(50, vec![(path1, Some(branch.clone()))]).await?; + storage.store_account_branches(50, vec![(path2, Some(branch.clone()))]).await?; + + let mut cursor = storage.trie_cursor(None, 100)?; + + // Current should be None initially + assert!(cursor.current()?.is_none()); + + // After seek, current should track position + cursor.seek(path1)?; + assert_eq!(cursor.current()?.unwrap(), path1); + + // After next, current should update + cursor.next()?; + assert_eq!(cursor.current()?.unwrap(), path2); + + Ok(()) + } + + /// Test current with no prior operations + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_current_no_prior_operations( + storage: S, + ) -> Result<(), ExternalStorageError> { + let mut cursor = storage.trie_cursor(None, 100)?; + + // Current should be None when no operations performed + assert!(cursor.current()?.is_none()); + + Ok(()) + } + + // ============================================================================= + // 4. Block Number Filtering + // ============================================================================= + + /// Test same path with different blocks + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_same_path_different_blocks( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1, 2]); + let branch1 = create_test_branch(); + let branch2 = create_test_branch_variant(); + + // Store same path at different blocks + storage.store_account_branches(50, vec![(path, Some(branch1.clone()))]).await?; + storage.store_account_branches(100, vec![(path, Some(branch2.clone()))]).await?; + + // Cursor with max_block_number=75 should see only block 50 data + let mut cursor75 = storage.trie_cursor(None, 75)?; + let result75 = cursor75.seek_exact(path)?.unwrap(); + assert_eq!(result75.0, path); + + // Cursor with max_block_number=150 should see block 100 data (latest) + let mut cursor150 = storage.trie_cursor(None, 150)?; + let result150 = cursor150.seek_exact(path)?.unwrap(); + assert_eq!(result150.0, path); + + Ok(()) + } + + /// Test deleted branch nodes + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_deleted_branch_nodes( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1, 2]); + let branch = create_test_branch(); + + // Store branch node, then delete it (store None) + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + storage.store_account_branches(100, vec![(path, None)]).await?; + + // Cursor before deletion should see the node + let mut cursor75 = storage.trie_cursor(None, 75)?; + assert!(cursor75.seek_exact(path)?.is_some()); + + // Cursor after deletion should not see the node + let mut cursor150 = storage.trie_cursor(None, 150)?; + assert!(cursor150.seek_exact(path)?.is_none()); + + Ok(()) + } + + // ============================================================================= + // 5. Hashed Address Filtering + // ============================================================================= + + /// Test account-specific cursor + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_account_specific_cursor( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1, 2]); + let addr1 = B256::repeat_byte(0x01); + let addr2 = B256::repeat_byte(0x02); + let branch = create_test_branch(); + + // Store same path for different accounts (using storage branches) + storage.store_storage_branches(50, addr1, vec![(path, Some(branch.clone()))]).await?; + storage.store_storage_branches(50, addr2, vec![(path, Some(branch.clone()))]).await?; + + // Cursor for addr1 should only see addr1 data + let mut cursor1 = storage.trie_cursor(Some(addr1), 100)?; + let result1 = cursor1.seek_exact(path)?.unwrap(); + assert_eq!(result1.0, path); + + // Cursor for addr2 should only see addr2 data + let mut cursor2 = storage.trie_cursor(Some(addr2), 100)?; + let result2 = cursor2.seek_exact(path)?.unwrap(); + assert_eq!(result2.0, path); + + // Cursor for addr1 should not see addr2 data when iterating + let mut cursor1_iter = storage.trie_cursor(Some(addr1), 100)?; + let mut found_count = 0; + while cursor1_iter.next()?.is_some() { + found_count += 1; + } + assert_eq!(found_count, 1); // Should only see one entry (for addr1) + + Ok(()) + } + + /// Test state trie cursor + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_state_trie_cursor( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path = nibbles_from(vec![1, 2]); + let addr = B256::repeat_byte(0x01); + let branch = create_test_branch(); + + // Store data for account trie and state trie + storage.store_storage_branches(50, addr, vec![(path, Some(branch.clone()))]).await?; + storage.store_account_branches(50, vec![(path, Some(branch.clone()))]).await?; + + // State trie cursor (None address) should only see state trie data + let mut state_cursor = storage.trie_cursor(None, 100)?; + let result = state_cursor.seek_exact(path)?.unwrap(); + assert_eq!(result.0, path); + + // Verify state cursor doesn't see account data when iterating + let mut state_cursor_iter = storage.trie_cursor(None, 100)?; + let mut found_count = 0; + while state_cursor_iter.next()?.is_some() { + found_count += 1; + } + + assert_eq!(found_count, 1); // Should only see state trie entry + + Ok(()) + } + + /// Test mixed account and state data + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_mixed_account_state_data( + storage: S, + ) -> Result<(), ExternalStorageError> { + let path1 = nibbles_from(vec![1]); + let path2 = nibbles_from(vec![2]); + let addr = B256::repeat_byte(0x01); + let branch = create_test_branch(); + + // Store mixed account and state trie data + storage.store_storage_branches(50, addr, vec![(path1, Some(branch.clone()))]).await?; + storage.store_account_branches(50, vec![(path2, Some(branch.clone()))]).await?; + + // Account cursor should only see account data + let mut account_cursor = storage.trie_cursor(Some(addr), 100)?; + let mut account_paths = Vec::new(); + while let Some((path, _)) = account_cursor.next()? { + account_paths.push(path); + } + assert_eq!(account_paths.len(), 1); + assert_eq!(account_paths[0], path1); + + // State cursor should only see state data + let mut state_cursor = storage.trie_cursor(None, 100)?; + let mut state_paths = Vec::new(); + while let Some((path, _)) = state_cursor.next()? { + state_paths.push(path); + } + assert_eq!(state_paths.len(), 1); + assert_eq!(state_paths[0], path2); + + Ok(()) + } + + // ============================================================================= + // 6. Path Ordering Tests + // ============================================================================= + + /// Test lexicographic ordering + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_lexicographic_ordering( + storage: S, + ) -> Result<(), ExternalStorageError> { + let paths = vec![ + nibbles_from(vec![3, 1]), + nibbles_from(vec![1, 2]), + nibbles_from(vec![2]), + nibbles_from(vec![1]), + ]; + let branch = create_test_branch(); + + // Store paths in random order + for path in &paths { + storage.store_account_branches(50, vec![(*path, Some(branch.clone()))]).await?; + } + + let mut cursor = storage.trie_cursor(None, 100)?; + let mut found_paths = Vec::new(); + while let Some((path, _)) = cursor.next()? { + found_paths.push(path); + } + + // Should be returned in lexicographic order: [1], [1,2], [2], [3,1] + let expected_order = vec![ + nibbles_from(vec![1]), + nibbles_from(vec![1, 2]), + nibbles_from(vec![2]), + nibbles_from(vec![3, 1]), + ]; + + assert_eq!(found_paths, expected_order); + + Ok(()) + } + + /// Test path prefix scenarios + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_path_prefix_scenarios( + storage: S, + ) -> Result<(), ExternalStorageError> { + let paths = vec![ + nibbles_from(vec![1]), // Prefix of next + nibbles_from(vec![1, 2]), // Extends first + nibbles_from(vec![1, 2, 3]), // Extends second + ]; + let branch = create_test_branch(); + + for path in &paths { + storage.store_account_branches(50, vec![(*path, Some(branch.clone()))]).await?; + } + + let mut cursor = storage.trie_cursor(None, 100)?; + + // Seek to prefix should find exact match + let result = cursor.seek_exact(paths[0])?.unwrap(); + assert_eq!(result.0, paths[0]); + + // Next should go to next path, not skip prefixed paths + let result = cursor.next()?.unwrap(); + assert_eq!(result.0, paths[1]); + + let result = cursor.next()?.unwrap(); + assert_eq!(result.0, paths[2]); + + Ok(()) + } + + /// Test complex nibble combinations + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_complex_nibble_combinations( + storage: S, + ) -> Result<(), ExternalStorageError> { + // Test various nibble patterns including edge values + let paths = vec![ + nibbles_from(vec![0]), + nibbles_from(vec![0, 15]), + nibbles_from(vec![15]), + nibbles_from(vec![15, 0]), + nibbles_from(vec![7, 8, 9]), + ]; + let branch = create_test_branch(); + + for path in &paths { + storage.store_account_branches(50, vec![(*path, Some(branch.clone()))]).await?; + } + + let mut cursor = storage.trie_cursor(None, 100)?; + let mut found_paths = Vec::new(); + while let Some((path, _)) = cursor.next()? { + found_paths.push(path); + } + + // All paths should be found and in correct order + assert_eq!(found_paths.len(), 5); + + // Verify specific ordering for edge cases + assert_eq!(found_paths[0], nibbles_from(vec![0])); + assert_eq!(found_paths[1], nibbles_from(vec![0, 15])); + assert_eq!(found_paths[4], nibbles_from(vec![15, 0])); + + Ok(()) + } + + // ============================================================================= + // 7. Leaf Node Tests (Hashed Accounts and Storage) + // ============================================================================= + + /// Test store and retrieve single account + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_store_and_retrieve_single_account( + storage: S, + ) -> Result<(), ExternalStorageError> { + let account_key = B256::repeat_byte(0x01); + let account = create_test_account(); + + // Store account + storage.store_hashed_accounts(vec![(account_key, Some(account))], 50).await?; + + // Retrieve via cursor + let mut cursor = storage.account_hashed_cursor(100)?; + let result = cursor.seek(account_key)?.unwrap(); + + assert_eq!(result.0, account_key); + assert_eq!(result.1.nonce, account.nonce); + assert_eq!(result.1.balance, account.balance); + assert_eq!(result.1.bytecode_hash, account.bytecode_hash); + + Ok(()) + } + + /// Test account cursor navigation + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_account_cursor_navigation( + storage: S, + ) -> Result<(), ExternalStorageError> { + let accounts = [ + (B256::repeat_byte(0x01), create_test_account()), + (B256::repeat_byte(0x03), create_test_account()), + (B256::repeat_byte(0x05), create_test_account()), + ]; + + // Store accounts + let accounts_to_store: Vec<_> = accounts.iter().map(|(k, v)| (*k, Some(*v))).collect(); + storage.store_hashed_accounts(accounts_to_store, 50).await?; + + let mut cursor = storage.account_hashed_cursor(100)?; + + // Test seeking to exact key + let result = cursor.seek(accounts[1].0)?.unwrap(); + assert_eq!(result.0, accounts[1].0); + + // Test seeking to key that doesn't exist (should return next greater) + let seek_key = B256::repeat_byte(0x02); + let result = cursor.seek(seek_key)?.unwrap(); + assert_eq!(result.0, accounts[1].0); // Should find 0x03 + + // Test next() navigation + let result = cursor.next()?.unwrap(); + assert_eq!(result.0, accounts[2].0); // Should find 0x05 + + // Test next() at end + assert!(cursor.next()?.is_none()); + + Ok(()) + } + + /// Test account block versioning + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_account_block_versioning( + storage: S, + ) -> Result<(), ExternalStorageError> { + let account_key = B256::repeat_byte(0x01); + let account_v1 = create_test_account_with_values(1, 100, 0xBB); + let account_v2 = create_test_account_with_values(2, 200, 0xDD); + + // Store account at different blocks + storage.store_hashed_accounts(vec![(account_key, Some(account_v1))], 50).await?; + storage.store_hashed_accounts(vec![(account_key, Some(account_v2))], 100).await?; + + // Cursor with max_block_number=75 should see v1 + let mut cursor75 = storage.account_hashed_cursor(75)?; + let result75 = cursor75.seek(account_key)?.unwrap(); + assert_eq!(result75.1.nonce, account_v1.nonce); + assert_eq!(result75.1.balance, account_v1.balance); + + // Cursor with max_block_number=150 should see v2 + let mut cursor150 = storage.account_hashed_cursor(150)?; + let result150 = cursor150.seek(account_key)?.unwrap(); + assert_eq!(result150.1.nonce, account_v2.nonce); + assert_eq!(result150.1.balance, account_v2.balance); + + Ok(()) + } + + /// Test store and retrieve storage + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_store_and_retrieve_storage( + storage: S, + ) -> Result<(), ExternalStorageError> { + let hashed_address = B256::repeat_byte(0x01); + let storage_slots = vec![ + (B256::repeat_byte(0x10), U256::from(100)), + (B256::repeat_byte(0x20), U256::from(200)), + (B256::repeat_byte(0x30), U256::from(300)), + ]; + + // Store storage slots + storage.store_hashed_storages(hashed_address, storage_slots.clone(), 50).await?; + + // Retrieve via cursor + let mut cursor = storage.storage_hashed_cursor(hashed_address, 100)?; + + // Test seeking to each slot + for (key, expected_value) in &storage_slots { + let result = cursor.seek(*key)?.unwrap(); + assert_eq!(result.0, *key); + assert_eq!(result.1, *expected_value); + } + + Ok(()) + } + + /// Test storage cursor navigation + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_storage_cursor_navigation( + storage: S, + ) -> Result<(), ExternalStorageError> { + let hashed_address = B256::repeat_byte(0x01); + let storage_slots = vec![ + (B256::repeat_byte(0x10), U256::from(100)), + (B256::repeat_byte(0x30), U256::from(300)), + (B256::repeat_byte(0x50), U256::from(500)), + ]; + + storage.store_hashed_storages(hashed_address, storage_slots.clone(), 50).await?; + + let mut cursor = storage.storage_hashed_cursor(hashed_address, 100)?; + + // Start from beginning with next() + let mut found_slots = Vec::new(); + while let Some((key, value)) = cursor.next()? { + found_slots.push((key, value)); + } + + assert_eq!(found_slots.len(), 3); + assert_eq!(found_slots[0], storage_slots[0]); + assert_eq!(found_slots[1], storage_slots[1]); + assert_eq!(found_slots[2], storage_slots[2]); + + Ok(()) + } + + /// Test storage account isolation + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_storage_account_isolation( + storage: S, + ) -> Result<(), ExternalStorageError> { + let address1 = B256::repeat_byte(0x01); + let address2 = B256::repeat_byte(0x02); + let storage_key = B256::repeat_byte(0x10); + + // Store same storage key for different accounts + storage.store_hashed_storages(address1, vec![(storage_key, U256::from(100))], 50).await?; + storage.store_hashed_storages(address2, vec![(storage_key, U256::from(200))], 50).await?; + + // Verify each account sees only its own storage + let mut cursor1 = storage.storage_hashed_cursor(address1, 100)?; + let result1 = cursor1.seek(storage_key)?.unwrap(); + assert_eq!(result1.1, U256::from(100)); + + let mut cursor2 = storage.storage_hashed_cursor(address2, 100)?; + let result2 = cursor2.seek(storage_key)?.unwrap(); + assert_eq!(result2.1, U256::from(200)); + + // Verify cursor1 doesn't see address2's storage + let mut cursor1_iter = storage.storage_hashed_cursor(address1, 100)?; + let mut count = 0; + while cursor1_iter.next()?.is_some() { + count += 1; + } + assert_eq!(count, 1); // Should only see one entry + + Ok(()) + } + + /// Test storage block versioning + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_storage_block_versioning( + storage: S, + ) -> Result<(), ExternalStorageError> { + let hashed_address = B256::repeat_byte(0x01); + let storage_key = B256::repeat_byte(0x10); + + // Store storage at different blocks + storage + .store_hashed_storages(hashed_address, vec![(storage_key, U256::from(100))], 50) + .await?; + storage + .store_hashed_storages(hashed_address, vec![(storage_key, U256::from(200))], 100) + .await?; + + // Cursor with max_block_number=75 should see old value + let mut cursor75 = storage.storage_hashed_cursor(hashed_address, 75)?; + let result75 = cursor75.seek(storage_key)?.unwrap(); + assert_eq!(result75.1, U256::from(100)); + + // Cursor with max_block_number=150 should see new value + let mut cursor150 = storage.storage_hashed_cursor(hashed_address, 150)?; + let result150 = cursor150.seek(storage_key)?.unwrap(); + assert_eq!(result150.1, U256::from(200)); + + Ok(()) + } + + /// Test storage zero value deletion + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_storage_zero_value_deletion( + storage: S, + ) -> Result<(), ExternalStorageError> { + let hashed_address = B256::repeat_byte(0x01); + let storage_key = B256::repeat_byte(0x10); + + // Store non-zero value + storage + .store_hashed_storages(hashed_address, vec![(storage_key, U256::from(100))], 50) + .await?; + + // "Delete" by storing zero value + storage.store_hashed_storages(hashed_address, vec![(storage_key, U256::ZERO)], 100).await?; + + // Cursor before deletion should see the value + let mut cursor75 = storage.storage_hashed_cursor(hashed_address, 75)?; + let result75 = cursor75.seek(storage_key)?.unwrap(); + assert_eq!(result75.1, U256::from(100)); + + // Cursor after deletion should NOT see the entry (zero values are skipped) + let mut cursor150 = storage.storage_hashed_cursor(hashed_address, 150)?; + let result150 = cursor150.seek(storage_key)?; + assert!(result150.is_none(), "Zero values should be skipped/deleted"); + + Ok(()) + } + + /// Test that zero values are skipped during iteration + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_storage_cursor_skips_zero_values( + storage: S, + ) -> Result<(), ExternalStorageError> { + let hashed_address = B256::repeat_byte(0x01); + + // Create a mix of non-zero and zero value storage slots + let storage_slots = vec![ + (B256::repeat_byte(0x10), U256::from(100)), // Non-zero + (B256::repeat_byte(0x20), U256::ZERO), // Zero value - should be skipped + (B256::repeat_byte(0x30), U256::from(300)), // Non-zero + (B256::repeat_byte(0x40), U256::ZERO), // Zero value - should be skipped + (B256::repeat_byte(0x50), U256::from(500)), // Non-zero + ]; + + // Store all slots + storage.store_hashed_storages(hashed_address, storage_slots.clone(), 50).await?; + + // Create cursor and iterate through all entries + let mut cursor = storage.storage_hashed_cursor(hashed_address, 100)?; + let mut found_slots = Vec::new(); + while let Some((key, value)) = cursor.next()? { + found_slots.push((key, value)); + } + + // Should only find 3 non-zero values + assert_eq!(found_slots.len(), 3, "Zero values should be skipped during iteration"); + + // Verify the non-zero values are the ones we stored + assert_eq!(found_slots[0], (B256::repeat_byte(0x10), U256::from(100))); + assert_eq!(found_slots[1], (B256::repeat_byte(0x30), U256::from(300))); + assert_eq!(found_slots[2], (B256::repeat_byte(0x50), U256::from(500))); + + // Verify seeking to a zero-value slot returns None or skips to next non-zero + let mut seek_cursor = storage.storage_hashed_cursor(hashed_address, 100)?; + let seek_result = seek_cursor.seek(B256::repeat_byte(0x20))?; + + // Should either return None or skip to the next non-zero value (0x30) + if let Some((key, value)) = seek_result { + assert_eq!( + key, + B256::repeat_byte(0x30), + "Should skip zero value and find next non-zero" + ); + assert_eq!(value, U256::from(300)); + } + + Ok(()) + } + + /// Test empty cursors + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_empty_cursors( + storage: S, + ) -> Result<(), ExternalStorageError> { + // Test empty account cursor + let mut account_cursor = storage.account_hashed_cursor(100)?; + assert!(account_cursor.seek(B256::repeat_byte(0x01))?.is_none()); + assert!(account_cursor.next()?.is_none()); + + // Test empty storage cursor + let mut storage_cursor = storage.storage_hashed_cursor(B256::repeat_byte(0x01), 100)?; + assert!(storage_cursor.seek(B256::repeat_byte(0x10))?.is_none()); + assert!(storage_cursor.next()?.is_none()); + + Ok(()) + } + + /// Test cursor boundary conditions + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_cursor_boundary_conditions( + storage: S, + ) -> Result<(), ExternalStorageError> { + let account_key = B256::repeat_byte(0x80); // Middle value + let account = create_test_account(); + + storage.store_hashed_accounts(vec![(account_key, Some(account))], 50).await?; + + let mut cursor = storage.account_hashed_cursor(100)?; + + // Seek to minimum key should find our account + let result = cursor.seek(B256::ZERO)?.unwrap(); + assert_eq!(result.0, account_key); + + // Seek to maximum key should find nothing + assert!(cursor.seek(B256::repeat_byte(0xFF))?.is_none()); + + // Seek to key just before our account should find our account + let just_before = B256::repeat_byte(0x7F); + let result = cursor.seek(just_before)?.unwrap(); + assert_eq!(result.0, account_key); + + Ok(()) + } + + /// Test large batch operations + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_large_batch_operations( + storage: S, + ) -> Result<(), ExternalStorageError> { + // Create large batch of accounts + let mut accounts = Vec::new(); + for i in 0..100 { + let key = B256::from([i as u8; 32]); + let account = create_test_account_with_values(i, i * 1000, (i + 1) as u8); + accounts.push((key, Some(account))); + } + + // Store in batch + storage.store_hashed_accounts(accounts.clone(), 50).await?; + + // Verify all accounts can be retrieved + let mut cursor = storage.account_hashed_cursor(100)?; + let mut found_count = 0; + while cursor.next()?.is_some() { + found_count += 1; + } + assert_eq!(found_count, 100); + + // Test specific account retrieval + let test_key = B256::from([42u8; 32]); + let result = cursor.seek(test_key)?.unwrap(); + assert_eq!(result.0, test_key); + assert_eq!(result.1.nonce, 42); + + Ok(()) + } + + /// Test wiped storage in `HashedPostState` + /// + /// When `store_trie_updates` receives a `HashedPostState` with wiped=true for a storage entry, + /// it should iterate all existing values for that address and create deletion entries for them. + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_store_trie_updates_with_wiped_storage( + storage: S, + ) -> Result<(), ExternalStorageError> { + use reth_trie::HashedStorage; + + let hashed_address = B256::repeat_byte(0x01); + + // First, store some storage values at block 50 + let storage_slots = vec![ + (B256::repeat_byte(0x10), U256::from(100)), + (B256::repeat_byte(0x20), U256::from(200)), + (B256::repeat_byte(0x30), U256::from(300)), + (B256::repeat_byte(0x40), U256::from(400)), + ]; + + storage.store_hashed_storages(hashed_address, storage_slots.clone(), 50).await?; + + // Verify all values are present at block 75 + let mut cursor75 = storage.storage_hashed_cursor(hashed_address, 75)?; + let mut found_slots = Vec::new(); + while let Some((key, value)) = cursor75.next()? { + found_slots.push((key, value)); + } + assert_eq!(found_slots.len(), 4, "All storage slots should be present before wipe"); + assert_eq!(found_slots[0], (B256::repeat_byte(0x10), U256::from(100))); + assert_eq!(found_slots[1], (B256::repeat_byte(0x20), U256::from(200))); + assert_eq!(found_slots[2], (B256::repeat_byte(0x30), U256::from(300))); + assert_eq!(found_slots[3], (B256::repeat_byte(0x40), U256::from(400))); + + // Now create a HashedPostState with wiped=true for this address at block 100 + let mut post_state = HashedPostState::default(); + let wiped_storage = HashedStorage::new(true); // wiped=true, empty storage map + post_state.storages.insert(hashed_address, wiped_storage); + + let block_state_diff = BlockStateDiff { trie_updates: TrieUpdates::default(), post_state }; + + // Store the wiped state + storage.store_trie_updates(100, block_state_diff).await?; + + // After wiping, cursor at block 150 should see NO storage values + let mut cursor150 = storage.storage_hashed_cursor(hashed_address, 150)?; + let mut found_slots_after_wipe = Vec::new(); + while let Some((key, value)) = cursor150.next()? { + found_slots_after_wipe.push((key, value)); + } + + assert_eq!( + found_slots_after_wipe.len(), + 0, + "All storage slots should be deleted after wipe. Found: {:?}", + found_slots_after_wipe + ); + + // Verify individual seeks also return None + for (slot, _) in &storage_slots { + let mut seek_cursor = storage.storage_hashed_cursor(hashed_address, 150)?; + let result = seek_cursor.seek(*slot)?; + assert!( + result.is_none() || result.unwrap().0 != *slot, + "Storage slot {:?} should be deleted after wipe", + slot + ); + } + + // Verify cursor at block 75 (before wipe) still sees all values + let mut cursor75_after = storage.storage_hashed_cursor(hashed_address, 75)?; + let mut found_slots_before_wipe = Vec::new(); + while let Some((key, value)) = cursor75_after.next()? { + found_slots_before_wipe.push((key, value)); + } + assert_eq!( + found_slots_before_wipe.len(), + 4, + "All storage slots should still be present when querying before wipe block" + ); + + Ok(()) + } + + /// Test that `store_trie_updates` properly stores branch nodes, leaf nodes, and removals + /// + /// This test verifies that all data stored via `store_trie_updates` can be read back + /// through the cursor APIs. + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_store_trie_updates_comprehensive( + storage: S, + ) -> Result<(), ExternalStorageError> { + use reth_trie::{updates::StorageTrieUpdates, HashedStorage}; + + let block_number = 100; + + // Create comprehensive trie updates with branches, leaves, and removals + let mut trie_updates = TrieUpdates::default(); + + // Add account branch nodes + let account_path1 = nibbles_from(vec![1, 2, 3]); + let account_path2 = nibbles_from(vec![4, 5, 6]); + let account_branch1 = create_test_branch(); + let account_branch2 = create_test_branch_variant(); + + trie_updates.account_nodes.insert(account_path1, account_branch1.clone()); + trie_updates.account_nodes.insert(account_path2, account_branch2.clone()); + + // Add removed account nodes + let removed_account_path = nibbles_from(vec![7, 8, 9]); + trie_updates.removed_nodes.insert(removed_account_path); + + // Add storage branch nodes for an address + let hashed_address = B256::repeat_byte(0x42); + let storage_path1 = nibbles_from(vec![1, 1]); + let storage_path2 = nibbles_from(vec![2, 2]); + let storage_branch = create_test_branch(); + + let mut storage_trie = StorageTrieUpdates::default(); + storage_trie.storage_nodes.insert(storage_path1, storage_branch.clone()); + storage_trie.storage_nodes.insert(storage_path2, storage_branch.clone()); + + // Add removed storage node + let removed_storage_path = nibbles_from(vec![3, 3]); + storage_trie.removed_nodes.insert(removed_storage_path); + + trie_updates.insert_storage_updates(hashed_address, storage_trie); + + // Create post state with accounts and storage + let mut post_state = HashedPostState::default(); + + // Add accounts + let account1_addr = B256::repeat_byte(0x10); + let account2_addr = B256::repeat_byte(0x20); + let account1 = create_test_account_with_values(1, 1000, 0xAA); + let account2 = create_test_account_with_values(2, 2000, 0xBB); + + post_state.accounts.insert(account1_addr, Some(account1)); + post_state.accounts.insert(account2_addr, Some(account2)); + + // Add deleted account + let deleted_account_addr = B256::repeat_byte(0x30); + post_state.accounts.insert(deleted_account_addr, None); + + // Add storage for an address + let storage_addr = B256::repeat_byte(0x50); + let mut hashed_storage = HashedStorage::new(false); + hashed_storage.storage.insert(B256::repeat_byte(0x01), U256::from(111)); + hashed_storage.storage.insert(B256::repeat_byte(0x02), U256::from(222)); + hashed_storage.storage.insert(B256::repeat_byte(0x03), U256::ZERO); // Deleted storage + post_state.storages.insert(storage_addr, hashed_storage); + + let block_state_diff = BlockStateDiff { trie_updates, post_state }; + + // Store the updates + storage.store_trie_updates(block_number, block_state_diff).await?; + + // ========== Verify Account Branch Nodes ========== + let mut account_trie_cursor = storage.trie_cursor(None, block_number + 10)?; + + // Should find the added branches + let result1 = account_trie_cursor.seek_exact(account_path1)?; + assert!(result1.is_some(), "Account branch node 1 should be found"); + assert_eq!(result1.unwrap().0, account_path1); + + let result2 = account_trie_cursor.seek_exact(account_path2)?; + assert!(result2.is_some(), "Account branch node 2 should be found"); + assert_eq!(result2.unwrap().0, account_path2); + + // Removed node should not be found + let removed_result = account_trie_cursor.seek_exact(removed_account_path)?; + assert!(removed_result.is_none(), "Removed account node should not be found"); + + // ========== Verify Storage Branch Nodes ========== + let mut storage_trie_cursor = + storage.trie_cursor(Some(hashed_address), block_number + 10)?; + + let storage_result1 = storage_trie_cursor.seek_exact(storage_path1)?; + assert!(storage_result1.is_some(), "Storage branch node 1 should be found"); + + let storage_result2 = storage_trie_cursor.seek_exact(storage_path2)?; + assert!(storage_result2.is_some(), "Storage branch node 2 should be found"); + + // Removed storage node should not be found + let removed_storage_result = storage_trie_cursor.seek_exact(removed_storage_path)?; + assert!(removed_storage_result.is_none(), "Removed storage node should not be found"); + + // ========== Verify Account Leaves ========== + let mut account_cursor = storage.account_hashed_cursor(block_number + 10)?; + + let acc1_result = account_cursor.seek(account1_addr)?; + assert!(acc1_result.is_some(), "Account 1 should be found"); + assert_eq!(acc1_result.unwrap().0, account1_addr); + assert_eq!(acc1_result.unwrap().1.nonce, 1); + assert_eq!(acc1_result.unwrap().1.balance, U256::from(1000)); + + let acc2_result = account_cursor.seek(account2_addr)?; + assert!(acc2_result.is_some(), "Account 2 should be found"); + assert_eq!(acc2_result.unwrap().1.nonce, 2); + + // Deleted account should not be found + let deleted_acc_result = account_cursor.seek(deleted_account_addr)?; + assert!( + deleted_acc_result.is_none() || deleted_acc_result.unwrap().0 != deleted_account_addr, + "Deleted account should not be found" + ); + + // ========== Verify Storage Leaves ========== + let mut storage_cursor = storage.storage_hashed_cursor(storage_addr, block_number + 10)?; + + let slot1_result = storage_cursor.seek(B256::repeat_byte(0x01))?; + assert!(slot1_result.is_some(), "Storage slot 1 should be found"); + assert_eq!(slot1_result.unwrap().1, U256::from(111)); + + let slot2_result = storage_cursor.seek(B256::repeat_byte(0x02))?; + assert!(slot2_result.is_some(), "Storage slot 2 should be found"); + assert_eq!(slot2_result.unwrap().1, U256::from(222)); + + // Zero-valued storage should not be found (deleted) + let slot3_result = storage_cursor.seek(B256::repeat_byte(0x03))?; + assert!( + slot3_result.is_none() || slot3_result.unwrap().0 != B256::repeat_byte(0x03), + "Zero-valued storage slot should not be found" + ); + + // ========== Verify fetch_trie_updates can retrieve the data ========== + let fetched_diff = storage.fetch_trie_updates(block_number).await?; + + // Check that trie updates are stored + assert_eq!( + fetched_diff.trie_updates.account_nodes_ref().len(), + 2, + "Should have 2 account nodes" + ); + assert_eq!( + fetched_diff.trie_updates.storage_tries_ref().len(), + 1, + "Should have 1 storage trie" + ); + + // Check that post state is stored + assert_eq!( + fetched_diff.post_state.accounts.len(), + 3, + "Should have 3 accounts (including deleted)" + ); + assert_eq!(fetched_diff.post_state.storages.len(), 1, "Should have 1 storage entry"); + + Ok(()) + } + + /// Test that `replace_updates` properly applies hashed/trie storage updates to the DB + /// + /// This test verifies the bug fix where `replace_updates` was only storing `trie_updates` + /// and `post_states` directly without populating the internal data structures + /// (`hashed_accounts`, `hashed_storages`, `account_branches`, `storage_branches`). + #[test_case(InMemoryExternalStorage::new(); "InMemory")] + #[tokio::test] + async fn test_replace_updates_applies_all_updates( + storage: S, + ) -> Result<(), ExternalStorageError> { + use reth_trie::{updates::StorageTrieUpdates, HashedStorage}; + + // ========== Setup: Store initial state at blocks 50, 100, 101 ========== + let initial_account_addr = B256::repeat_byte(0x10); + let initial_account = create_test_account_with_values(1, 1000, 0xAA); + + let initial_storage_addr = B256::repeat_byte(0x20); + let initial_storage_slot = B256::repeat_byte(0x01); + let initial_storage_value = U256::from(100); + + let initial_branch_path = nibbles_from(vec![1, 2, 3]); + let initial_branch = create_test_branch(); + + // Store initial data at block 50 + let mut initial_trie_updates_50 = TrieUpdates::default(); + initial_trie_updates_50.account_nodes.insert(initial_branch_path, initial_branch.clone()); + + let mut initial_post_state_50 = HashedPostState::default(); + initial_post_state_50.accounts.insert(initial_account_addr, Some(initial_account)); + + let initial_diff_50 = BlockStateDiff { + trie_updates: initial_trie_updates_50, + post_state: initial_post_state_50, + }; + storage.store_trie_updates(50, initial_diff_50).await?; + + // Store data at block 100 (common block) + let mut initial_trie_updates_100 = TrieUpdates::default(); + let common_branch_path = nibbles_from(vec![4, 5, 6]); + initial_trie_updates_100.account_nodes.insert(common_branch_path, initial_branch.clone()); + + let mut initial_post_state_100 = HashedPostState::default(); + let mut initial_storage_100 = HashedStorage::new(false); + initial_storage_100.storage.insert(initial_storage_slot, initial_storage_value); + initial_post_state_100.storages.insert(initial_storage_addr, initial_storage_100); + + let initial_diff_100 = BlockStateDiff { + trie_updates: initial_trie_updates_100, + post_state: initial_post_state_100, + }; + storage.store_trie_updates(100, initial_diff_100).await?; + + // Store data at block 101 (will be replaced) + let mut initial_trie_updates_101 = TrieUpdates::default(); + let old_branch_path = nibbles_from(vec![7, 8, 9]); + initial_trie_updates_101.account_nodes.insert(old_branch_path, initial_branch.clone()); + + let mut initial_post_state_101 = HashedPostState::default(); + let old_account_addr = B256::repeat_byte(0x30); + let old_account = create_test_account_with_values(99, 9999, 0xFF); + initial_post_state_101.accounts.insert(old_account_addr, Some(old_account)); + + let initial_diff_101 = BlockStateDiff { + trie_updates: initial_trie_updates_101, + post_state: initial_post_state_101, + }; + storage.store_trie_updates(101, initial_diff_101).await?; + + // ========== Verify initial state exists ========== + // Verify block 50 data exists + let mut cursor_initial = storage.trie_cursor(None, 75)?; + assert!( + cursor_initial.seek_exact(initial_branch_path)?.is_some(), + "Initial branch should exist before replace" + ); + + // Verify block 101 old data exists + let mut cursor_old = storage.trie_cursor(None, 150)?; + assert!( + cursor_old.seek_exact(old_branch_path)?.is_some(), + "Old branch at block 101 should exist before replace" + ); + + let mut account_cursor_old = storage.account_hashed_cursor(150)?; + assert!( + account_cursor_old.seek(old_account_addr)?.is_some(), + "Old account at block 101 should exist before replace" + ); + + // ========== Call replace_updates to replace blocks after 100 ========== + let mut blocks_to_add: HashMap = HashMap::default(); + + // New data for block 101 + let new_account_addr = B256::repeat_byte(0x40); + let new_account = create_test_account_with_values(5, 5000, 0xCC); + + let new_storage_addr = B256::repeat_byte(0x50); + let new_storage_slot = B256::repeat_byte(0x02); + let new_storage_value = U256::from(999); + + let new_branch_path = nibbles_from(vec![10, 11, 12]); + let new_branch = create_test_branch_variant(); + + let storage_branch_path = nibbles_from(vec![5, 5]); + let storage_hashed_addr = B256::repeat_byte(0x60); + + let mut new_trie_updates = TrieUpdates::default(); + new_trie_updates.account_nodes.insert(new_branch_path, new_branch.clone()); + + // Add storage trie updates + let mut storage_trie = StorageTrieUpdates::default(); + storage_trie.storage_nodes.insert(storage_branch_path, new_branch.clone()); + new_trie_updates.insert_storage_updates(storage_hashed_addr, storage_trie); + + let mut new_post_state = HashedPostState::default(); + new_post_state.accounts.insert(new_account_addr, Some(new_account)); + + let mut new_storage = HashedStorage::new(false); + new_storage.storage.insert(new_storage_slot, new_storage_value); + new_post_state.storages.insert(new_storage_addr, new_storage); + + blocks_to_add.insert( + 101, + BlockStateDiff { trie_updates: new_trie_updates, post_state: new_post_state }, + ); + + // New data for block 102 + let block_102_account_addr = B256::repeat_byte(0x70); + let block_102_account = create_test_account_with_values(10, 10000, 0xDD); + + let mut trie_updates_102 = TrieUpdates::default(); + let block_102_branch_path = nibbles_from(vec![15, 14, 13]); + trie_updates_102.account_nodes.insert(block_102_branch_path, new_branch.clone()); + + let mut post_state_102 = HashedPostState::default(); + post_state_102.accounts.insert(block_102_account_addr, Some(block_102_account)); + + blocks_to_add.insert( + 102, + BlockStateDiff { trie_updates: trie_updates_102, post_state: post_state_102 }, + ); + + // Execute replace_updates + storage.replace_updates(100, blocks_to_add).await?; + + // ========== Verify that data up to block 100 still exists ========== + let mut cursor_50 = storage.trie_cursor(None, 75)?; + assert!( + cursor_50.seek_exact(initial_branch_path)?.is_some(), + "Block 50 branch should still exist after replace" + ); + + let mut cursor_100 = storage.trie_cursor(None, 100)?; + assert!( + cursor_100.seek_exact(common_branch_path)?.is_some(), + "Block 100 branch should still exist after replace" + ); + + let mut storage_cursor_100 = storage.storage_hashed_cursor(initial_storage_addr, 100)?; + let result_100 = storage_cursor_100.seek(initial_storage_slot)?; + assert!(result_100.is_some(), "Block 100 storage should still exist after replace"); + assert_eq!( + result_100.unwrap().1, + initial_storage_value, + "Block 100 storage value should be unchanged" + ); + + // ========== Verify that old data after block 100 is gone ========== + let mut cursor_old_gone = storage.trie_cursor(None, 150)?; + assert!( + cursor_old_gone.seek_exact(old_branch_path)?.is_none(), + "Old branch at block 101 should be removed after replace" + ); + + let mut account_cursor_old_gone = storage.account_hashed_cursor(150)?; + let old_acc_result = account_cursor_old_gone.seek(old_account_addr)?; + assert!( + old_acc_result.is_none() || old_acc_result.unwrap().0 != old_account_addr, + "Old account at block 101 should be removed after replace" + ); + + // ========== Verify new data is properly accessible via cursors ========== + + // Verify new account branch nodes + let mut trie_cursor = storage.trie_cursor(None, 150)?; + let branch_result = trie_cursor.seek_exact(new_branch_path)?; + assert!(branch_result.is_some(), "New account branch should be accessible via cursor"); + assert_eq!(branch_result.unwrap().0, new_branch_path); + + // Verify new storage branch nodes + let mut storage_trie_cursor = storage.trie_cursor(Some(storage_hashed_addr), 150)?; + let storage_branch_result = storage_trie_cursor.seek_exact(storage_branch_path)?; + assert!( + storage_branch_result.is_some(), + "New storage branch should be accessible via cursor" + ); + assert_eq!(storage_branch_result.unwrap().0, storage_branch_path); + + // Verify new hashed accounts + let mut account_cursor = storage.account_hashed_cursor(150)?; + let account_result = account_cursor.seek(new_account_addr)?; + assert!(account_result.is_some(), "New account should be accessible via cursor"); + assert_eq!(account_result.as_ref().unwrap().0, new_account_addr); + assert_eq!(account_result.as_ref().unwrap().1.nonce, new_account.nonce); + assert_eq!(account_result.as_ref().unwrap().1.balance, new_account.balance); + assert_eq!(account_result.as_ref().unwrap().1.bytecode_hash, new_account.bytecode_hash); + + // Verify new hashed storages + let mut storage_cursor = storage.storage_hashed_cursor(new_storage_addr, 150)?; + let storage_result = storage_cursor.seek(new_storage_slot)?; + assert!(storage_result.is_some(), "New storage should be accessible via cursor"); + assert_eq!(storage_result.as_ref().unwrap().0, new_storage_slot); + assert_eq!(storage_result.as_ref().unwrap().1, new_storage_value); + + // Verify block 102 data + let mut trie_cursor_102 = storage.trie_cursor(None, 150)?; + let branch_result_102 = trie_cursor_102.seek_exact(block_102_branch_path)?; + assert!(branch_result_102.is_some(), "Block 102 branch should be accessible"); + assert_eq!(branch_result_102.unwrap().0, block_102_branch_path); + + let mut account_cursor_102 = storage.account_hashed_cursor(150)?; + let account_result_102 = account_cursor_102.seek(block_102_account_addr)?; + assert!(account_result_102.is_some(), "Block 102 account should be accessible"); + assert_eq!(account_result_102.as_ref().unwrap().0, block_102_account_addr); + assert_eq!(account_result_102.as_ref().unwrap().1.nonce, block_102_account.nonce); + + // Verify fetch_trie_updates returns the new data + let fetched_101 = storage.fetch_trie_updates(101).await?; + assert_eq!( + fetched_101.trie_updates.account_nodes_ref().len(), + 1, + "Should have 1 account branch node at block 101" + ); + assert!( + fetched_101.trie_updates.account_nodes_ref().contains_key(&new_branch_path), + "New branch path should be in trie_updates" + ); + assert_eq!(fetched_101.post_state.accounts.len(), 1, "Should have 1 account at block 101"); + assert!( + fetched_101.post_state.accounts.contains_key(&new_account_addr), + "New account should be in post_state" + ); + + Ok(()) + } +}