diff --git a/benches/crud_benchmarks.rs b/benches/crud_benchmarks.rs index 450601d..d12d902 100644 --- a/benches/crud_benchmarks.rs +++ b/benches/crud_benchmarks.rs @@ -20,10 +20,17 @@ use alloy_primitives::{StorageKey, StorageValue, U256}; use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY}; use criterion::{criterion_group, criterion_main, BenchmarkId, Criterion}; use rand::prelude::*; -use std::{fs, io, path::Path, time::Duration}; +use std::{ + fs, io, + path::Path, + sync::{Arc, Barrier}, + thread, + time::Duration, +}; use tempdir::TempDir; use triedb::{ account::Account, + overlay::{OverlayStateMut, OverlayValue}, path::{AddressPath, StoragePath}, Database, }; @@ -82,6 +89,65 @@ fn bench_account_reads(c: &mut Criterion) { }, ); }); + group.bench_function(BenchmarkId::new("eoa_reads_parallel", BATCH_SIZE), |b| { + b.iter_with_setup( + || { + let db_path = dir.path().join(&file_name); + let db = Arc::new(Database::open(db_path.clone()).unwrap()); + + // Spawn 4 reader threads + let thread_count = 4; + let setup_barrier = Arc::new(Barrier::new(thread_count + 1)); + let test_barrier = Arc::new(Barrier::new(thread_count + 1)); + let chunk_size = addresses.len() / thread_count; + let mut handles = Vec::new(); + + for i in 0..thread_count { + let start_idx = i * chunk_size; + let end_idx = if i == thread_count { + // Last thread handles any remaining addresses + addresses.len() + } else { + (i + 1) * chunk_size + }; + + let thread_addresses = addresses[start_idx..end_idx].to_vec(); + let db_clone = Arc::clone(&db); + let setup_barrier = Arc::clone(&setup_barrier); + let test_barrier = Arc::clone(&test_barrier); + + let handle = thread::spawn(move || { + setup_barrier.wait(); + test_barrier.wait(); + // Each thread creates its own RO transaction + let mut tx = db_clone.begin_ro().unwrap(); + + // Read its chunk of addresses + for addr in thread_addresses { + let a = tx.get_account(&addr).unwrap(); + assert!(a.is_some()); + } + + // Commit the transaction + tx.commit().unwrap(); + }); + + handles.push(handle); + } + + setup_barrier.wait(); + + (handles, test_barrier) + }, + |(handles, test_barrier)| { + test_barrier.wait(); + // Wait for all threads to complete + for handle in handles { + handle.join().unwrap(); + } + }, + ); + }); group.finish(); } @@ -191,7 +257,7 @@ fn bench_account_updates(c: &mut Criterion) { b.iter_with_setup( || { let db_path = dir.path().join(&file_name); - Database::open(db_path.clone()).unwrap() + Database::open(db_path).unwrap() }, |db| { let mut tx = db.begin_rw().unwrap(); @@ -538,6 +604,51 @@ fn bench_storage_deletes(c: &mut Criterion) { group.finish(); } +fn bench_state_root_with_overlay(c: &mut Criterion) { + let mut group = c.benchmark_group("state_root_with_overlay"); + let base_dir = get_base_database( + DEFAULT_SETUP_DB_EOA_SIZE, + DEFAULT_SETUP_DB_CONTRACT_SIZE, + DEFAULT_SETUP_DB_STORAGE_PER_CONTRACT, + ); + let dir = TempDir::new("triedb_bench_state_root_with_overlay").unwrap(); + let file_name = base_dir.main_file_name.clone(); + copy_files(&base_dir, dir.path()).unwrap(); + + let mut rng = StdRng::seed_from_u64(SEED_CONTRACT); + let addresses: Vec = + (0..BATCH_SIZE).map(|_| generate_random_address(&mut rng)).collect(); + + let mut account_overlay_mut = OverlayStateMut::new(); + addresses.iter().enumerate().for_each(|(i, addr)| { + let new_account = + Account::new(i as u64, U256::from(i as u64), EMPTY_ROOT_HASH, KECCAK_EMPTY); + account_overlay_mut.insert(addr.clone().into(), Some(OverlayValue::Account(new_account))); + }); + let account_overlay = account_overlay_mut.freeze(); + + group.throughput(criterion::Throughput::Elements(BATCH_SIZE as u64)); + group.measurement_time(Duration::from_secs(30)); + group.bench_function(BenchmarkId::new("state_root_with_account_overlay", BATCH_SIZE), |b| { + b.iter_with_setup( + || { + let db_path = dir.path().join(&file_name); + Database::open(db_path).unwrap() + }, + |db| { + let tx = db.begin_ro().unwrap(); + + // Compute the root hash with the overlay + let _root_result = tx.compute_root_with_overlay(account_overlay.clone()).unwrap(); + + tx.commit().unwrap(); + }, + ); + }); + + group.finish(); +} + criterion_group!( benches, bench_account_reads, @@ -550,5 +661,6 @@ criterion_group!( bench_storage_inserts, bench_storage_updates, bench_storage_deletes, + bench_state_root_with_overlay, ); criterion_main!(benches); diff --git a/src/account.rs b/src/account.rs index e4937ba..b127d58 100644 --- a/src/account.rs +++ b/src/account.rs @@ -1,9 +1,10 @@ use alloy_primitives::{B256, U256}; +use alloy_rlp::{MaxEncodedLen, RlpEncodable}; use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY}; use proptest::prelude::*; use proptest_derive::Arbitrary; -#[derive(Debug, Clone, PartialEq, Eq, Default, Arbitrary)] +#[derive(Debug, Clone, PartialEq, Eq, Default, Arbitrary, RlpEncodable)] pub struct Account { pub nonce: u64, pub balance: U256, @@ -19,6 +20,16 @@ impl Account { } } +/// This is the maximum possible RLP-encoded length of an account. +/// +/// This value is derived from the maximum possible length of an account, which is the largest +/// case. An account is encoded as a list of 4 elements, with 3 of these represnting 32 byte values +/// and the nonce being an 8 byte value. Each element has 1 extra byte of encoding overhead. +/// The list also has 2 bytes of encoding overhead. The total length is `2 + 3*33 + 9 = 110`. +const MAX_RLP_ENCODED_LEN: usize = 2 + 3 * 33 + 9; + +unsafe impl MaxEncodedLen for Account {} + #[cfg(test)] mod tests { use super::*; diff --git a/src/lib.rs b/src/lib.rs index 022db22..98e9c5c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -17,6 +17,7 @@ pub mod location; pub mod meta; pub mod metrics; pub mod node; +pub mod overlay; pub mod page; pub mod path; pub mod pointer; diff --git a/src/node.rs b/src/node.rs index 86e7d52..a0ccd59 100644 --- a/src/node.rs +++ b/src/node.rs @@ -3,7 +3,7 @@ use crate::{ pointer::Pointer, storage::value::{self, Value}, }; -use alloy_primitives::{hex, StorageValue, B256, U256}; +use alloy_primitives::{StorageValue, B256, U256}; use alloy_rlp::{ decode_exact, encode_fixed_size, length_of_length, BufMut, Encodable, Header, MaxEncodedLen, EMPTY_STRING_CODE, @@ -17,11 +17,6 @@ use proptest::{arbitrary, strategy, strategy::Strategy}; use proptest_derive::Arbitrary; use std::cmp::{max, min}; -// This is equivalent to RlpNode::word_rlp(&EMPTY_ROOT_HASH), and is used to encode the storage root -// of an account with no storage. -const EMPTY_ROOT_RLP: [u8; 33] = - hex!("0xa056e81f171bcc55a6ff8345e692c0f86e5b48e01b996cadc001622fb5e363b421"); - const MAX_PREFIX_LENGTH: usize = 64; /// A node in the trie. @@ -152,6 +147,10 @@ impl Node { &self.0.kind } + pub fn into_kind(self) -> NodeKind { + self.0.kind + } + /// Returns whether the [Node] type supports children. pub const fn has_children(&self) -> bool { matches!(self.kind(), NodeKind::Branch { .. } | NodeKind::AccountLeaf { .. }) @@ -567,17 +566,17 @@ impl Encodable for Node { let mut buf = [0u8; 110]; // max RLP length for an account: 2 bytes for list length, 9 for nonce, 33 for // balance, 33 for storage root, 33 for code hash let mut value_rlp = buf.as_mut(); - let storage_root_rlp = - storage_root.as_ref().map(|p| p.rlp().as_slice()).unwrap_or(&EMPTY_ROOT_RLP); - let len = 2 + nonce_rlp.len() + balance_rlp.len() + storage_root_rlp.len() + 33; - value_rlp.put_u8(0xf8); - value_rlp.put_u8((len - 2) as u8); - value_rlp.put_slice(nonce_rlp); - value_rlp.put_slice(balance_rlp); - value_rlp.put_slice(storage_root_rlp); - value_rlp.put_u8(0xa0); - value_rlp.put_slice(code_hash.as_slice()); - LeafNodeRef { key: prefix, value: &buf[..len] }.encode(out); + let storage_root_hash = storage_root + .as_ref() + .map_or(EMPTY_ROOT_HASH, |p| p.rlp().as_hash().unwrap_or(EMPTY_ROOT_HASH)); + let account_rlp_length = encode_account_leaf( + nonce_rlp, + balance_rlp, + code_hash, + &storage_root_hash, + &mut value_rlp, + ); + LeafNodeRef { key: prefix, value: &buf[..account_rlp_length] }.encode(out); } NodeKind::Branch { ref children } => { if prefix.is_empty() { @@ -634,6 +633,28 @@ impl Encodable for Node { } } +#[inline] +pub fn encode_account_leaf( + nonce_rlp: &ArrayVec, + balance_rlp: &ArrayVec, + code_hash: &B256, + storage_root: &B256, + out: &mut dyn BufMut, +) -> usize { + let len = 2 + nonce_rlp.len() + balance_rlp.len() + 33 * 2; + out.put_u8(0xf8); + out.put_u8((len - 2) as u8); + out.put_slice(nonce_rlp); + out.put_slice(balance_rlp); + out.put_u8(0xa0); + out.put_slice(storage_root.as_slice()); + out.put_u8(0xa0); + out.put_slice(code_hash.as_slice()); + + len +} + +#[inline] pub fn encode_branch(children: &[Option], out: &mut dyn BufMut) -> usize { // first encode the header let mut payload_length = 1; diff --git a/src/overlay.rs b/src/overlay.rs new file mode 100644 index 0000000..6ec6cc7 --- /dev/null +++ b/src/overlay.rs @@ -0,0 +1,890 @@ +use crate::{account::Account, node::TrieValue}; +use alloy_primitives::{StorageValue, B256, U256}; +use alloy_trie::Nibbles; +use std::{cmp::min, sync::Arc}; + +#[derive(Debug, Clone)] +pub enum OverlayValue { + Account(Account), + Storage(StorageValue), + Hash(B256), +} + +impl TryFrom for TrieValue { + type Error = &'static str; + + fn try_from(value: OverlayValue) -> Result { + match value { + OverlayValue::Account(account) => Ok(TrieValue::Account(account)), + OverlayValue::Storage(storage) => Ok(TrieValue::Storage(storage)), + OverlayValue::Hash(_) => Err("Cannot convert Hash overlay value to TrieValue"), + } + } +} + +/// Mutable overlay state that accumulates changes during transaction building. +/// Changes are stored unsorted for fast insertion, then sorted when frozen. +#[derive(Debug, Clone, Default)] +pub struct OverlayStateMut { + changes: Vec<(Nibbles, Option)>, +} + +impl OverlayStateMut { + /// Creates a new empty mutable overlay state. + pub fn new() -> Self { + Self { changes: Vec::new() } + } + + /// Creates a new mutable overlay state with the specified capacity. + pub fn with_capacity(capacity: usize) -> Self { + Self { changes: Vec::with_capacity(capacity) } + } + + /// Inserts a change into the overlay state. + /// Multiple changes to the same path will keep the latest value. + pub fn insert(&mut self, path: Nibbles, value: Option) { + // For now, just append. We could optimize by checking for duplicates, + // but the freeze operation will handle deduplication anyway. + + match value { + Some(OverlayValue::Storage(U256::ZERO)) => { + // If the value is zero, this is actually a tombstone + self.changes.push((path, None)); + } + _ => { + self.changes.push((path, value)); + } + } + } + + /// Returns the number of changes in the overlay. + pub fn len(&self) -> usize { + self.changes.len() + } + + /// Returns true if the overlay is empty. + pub fn is_empty(&self) -> bool { + self.changes.is_empty() + } + + /// Freezes the mutable overlay into an immutable sorted overlay. + /// This sorts the changes and deduplicates by path, keeping the last value for each path. + pub fn freeze(mut self) -> OverlayState { + if self.changes.is_empty() { + return OverlayState { data: Arc::new([]), start_idx: 0, end_idx: 0, prefix_offset: 0 }; + } + + // Sort by path + self.changes.sort_by(|a, b| a.0.cmp(&b.0)); + + // Deduplicate by path, keeping the last occurrence of each path + let mut deduped = Vec::with_capacity(self.changes.len()); + let mut last_path: Option<&Nibbles> = None; + + for (path, value) in &self.changes { + if last_path != Some(path) { + deduped.push((path.clone(), value.clone())); + last_path = Some(path); + } else { + // Update the last entry with the newer value + if let Some(last_entry) = deduped.last_mut() { + last_entry.1 = value.clone(); + } + } + } + + let data: Arc<[(Nibbles, Option)]> = Arc::from(deduped.into_boxed_slice()); + let len = data.len(); + OverlayState { data, start_idx: 0, end_idx: len, prefix_offset: 0 } + } +} + +/// Immutable overlay state with sorted changes for efficient querying and zero-copy sub-slicing. +/// This is created by freezing an OverlayStateMut and allows for efficient prefix operations. +/// +/// The overlay uses Arc-based storage for zero-copy sub-slicing and thread-safe sharing. +/// Sub-slices share the same underlying data and only track different bounds within it. +/// +/// # Thread Safety +/// +/// OverlayState is thread-safe and can be safely shared across threads. Multiple threads +/// can work with non-overlapping sub-slices simultaneously without any synchronization. +/// The Arc<[...]> provides thread-safe reference counting for the underlying data. +#[derive(Debug, Clone)] +pub struct OverlayState { + data: Arc<[(Nibbles, Option)]>, + start_idx: usize, + end_idx: usize, + prefix_offset: usize, +} + +impl OverlayState { + /// Creates an empty overlay state. + pub fn empty() -> Self { + Self { data: Arc::new([]), start_idx: 0, end_idx: 0, prefix_offset: 0 } + } + + #[cfg(test)] + pub fn data(&self) -> &[(Nibbles, Option)] { + &self.data + } + + pub fn get(&self, index: usize) -> Option<(&[u8], &Option)> { + let slice = self.effective_slice(); + if index < slice.len() { + Some((&slice[index].0[self.prefix_offset..], &slice[index].1)) + } else { + None + } + } + + #[inline] + pub fn first(&self) -> Option<(&[u8], &Option)> { + if self.is_empty() { + None + } else { + let (path, value) = &self.data[self.start_idx]; + Some((&path[self.prefix_offset..], value)) + } + } + + /// Returns the effective slice of data for this overlay, respecting bounds. + pub fn effective_slice(&self) -> &[(Nibbles, Option)] { + &self.data[self.start_idx..self.end_idx] + } + + /// Returns the number of changes in the overlay. + pub fn len(&self) -> usize { + self.end_idx - self.start_idx + } + + /// Returns true if the overlay is empty. + pub fn is_empty(&self) -> bool { + self.start_idx >= self.end_idx + } + + /// Looks up a specific path in the overlay. + /// Returns Some(value) if found, None if not in overlay. + /// The path is adjusted by the prefix_offset before lookup. + pub fn lookup(&self, path: &[u8]) -> Option<&Option> { + let slice = self.effective_slice(); + match slice.binary_search_by(|(p, _)| self.compare_with_offset(p, path)) { + Ok(index) => Some(&slice[index].1), + Err(_) => None, + } + } + + /// Creates a new OverlayState with a prefix offset applied. + /// This is used for storage trie traversal where we want to strip the account prefix (64 + /// nibbles) from all paths, effectively converting 128-nibble storage paths to 64-nibble + /// storage-relative paths. + /// + /// # Panics + /// + /// Panics if the prefix offset would break the sort order. This happens when paths don't share + /// a common prefix up to the offset length, which would cause the offset-adjusted paths to be + /// out of order. + pub fn with_prefix_offset(&self, offset: usize) -> OverlayState { + let total_offset = self.prefix_offset + offset; + + // Validate that all paths share the same prefix up to total_offset + // and that offset-adjusted paths maintain sort order + if self.len() > 1 { + let slice = self.effective_slice(); + let first_path = &slice[0].0; + let last_path = &slice[slice.len() - 1].0; + + // Check that both first and last paths are long enough + if first_path.len() < total_offset || last_path.len() < total_offset { + panic!( + "with_prefix_offset: offset {} would exceed path length (first: {}, last: {})", + total_offset, + first_path.len(), + last_path.len() + ); + } + + // Check that first and last paths share the same prefix up to total_offset + let first_prefix = first_path.slice(..total_offset); + let last_prefix = last_path.slice(..total_offset); + if first_prefix != last_prefix { + panic!("with_prefix_offset: paths don't share common prefix up to offset {total_offset}\nfirst: {first_path:?}\nlast: {last_path:?}"); + } + } + + OverlayState { + data: Arc::clone(&self.data), + start_idx: self.start_idx, + end_idx: self.end_idx, + prefix_offset: total_offset, + } + } + + /// Returns the first index of the effective slice with the given nibbles prefix + fn find_start_index_with_prefix( + &self, + effective_slice: &[(Nibbles, Option)], + prefix: &[u8], + ) -> usize { + effective_slice + .binary_search_by(|(p, _)| self.compare_with_offset(p, prefix)) + .unwrap_or_else(|i| i) + } + + /// Returns the first index of the effective slice immediately following the given prefix. + /// This is used to find the end of a prefix range in the overlay. + fn find_end_index_with_prefix( + &self, + effective_slice: &[(Nibbles, Option)], + prefix: &[u8], + ) -> usize { + effective_slice.partition_point(|(stored_path, _)| { + if stored_path.len() < self.prefix_offset { + true + } else { + let adjusted_path = &stored_path[self.prefix_offset..]; + &adjusted_path[..min(adjusted_path.len(), prefix.len())] <= prefix + } + }) + } + + /// Helper method to compare a stored path with a target path, applying prefix_offset. + /// Returns the comparison result of the offset-adjusted stored path vs target path. + fn compare_with_offset(&self, stored_path: &[u8], target_path: &[u8]) -> std::cmp::Ordering { + if stored_path.len() < self.prefix_offset { + std::cmp::Ordering::Less + } else { + let adjusted_path = &stored_path[self.prefix_offset..]; + adjusted_path.cmp(target_path) + } + } + + /// Creates a zero-copy sub-slice of the overlay from the given range. + /// This is useful for recursive traversal where we want to pass a subset + /// of changes to child operations without copying data. + /// + /// # Performance + /// + /// This operation is O(1) and creates no allocations. The returned OverlayState + /// shares the same underlying Arc<[...]> data and only tracks different bounds. + pub fn sub_slice(&self, start: usize, end: usize) -> OverlayState { + let current_len = self.len(); + let end = end.min(current_len); + let start = start.min(end); + + OverlayState { + data: Arc::clone(&self.data), + start_idx: self.start_idx + start, + end_idx: self.start_idx + end, + prefix_offset: self.prefix_offset, + } + } + + /// Partitions the overlay into three slices: + /// - before: all changes before the prefix + /// - with_prefix: all changes with the prefix + /// - after: all changes after the prefix + pub fn sub_slice_by_prefix(&self, prefix: &[u8]) -> (OverlayState, OverlayState, OverlayState) { + let slice = self.effective_slice(); + + let start_index = self.find_start_index_with_prefix(slice, prefix); + let end_index = self.find_end_index_with_prefix(slice, prefix); + + let before = self.sub_slice(0, start_index); + let with_prefix = self.sub_slice(start_index, end_index); + let after = self.sub_slice(end_index, slice.len()); + (before, with_prefix, after) + } + + /// Creates a zero-copy sub-slice containing only changes that affect the subtree + /// rooted at the given path prefix. This is used during recursive trie traversal + /// to filter overlay changes relevant to each subtree. + pub fn sub_slice_for_prefix(&self, prefix: &[u8]) -> OverlayState { + if self.is_empty() { + return OverlayState::empty(); + } + + let slice = self.effective_slice(); + + let start_index = self.find_start_index_with_prefix(slice, prefix); + let end_index = self.find_end_index_with_prefix(slice, prefix); + + self.sub_slice(start_index, end_index) + } + + /// Returns an iterator over all changes in the overlay, respecting slice bounds. + /// The paths are adjusted by the prefix_offset. + pub fn iter(&self) -> impl Iterator)> { + self.effective_slice().iter().filter_map(move |(path, value)| { + if path.len() > self.prefix_offset { + let adjusted_path = &path.as_slice()[self.prefix_offset..]; + Some((adjusted_path, value)) + } else { + None + } + }) + } + + pub fn contains_prefix_of(&self, path: &[u8]) -> bool { + // Check for exact match or any path that starts with the given path after applying offset + self.iter().any(|(overlay_path, _)| path.starts_with(overlay_path)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::account::Account; + use alloy_primitives::U256; + use alloy_trie::{Nibbles, EMPTY_ROOT_HASH, KECCAK_EMPTY}; + + fn test_account() -> Account { + Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY) + } + + #[test] + fn test_mutable_overlay_basic_operations() { + let mut overlay = OverlayStateMut::new(); + assert!(overlay.is_empty()); + assert_eq!(overlay.len(), 0); + + let path1 = Nibbles::from_nibbles([1, 2, 3, 4]); + let path2 = Nibbles::from_nibbles([5, 6, 7, 8]); + let account = test_account(); + + overlay.insert(path1.clone(), Some(OverlayValue::Account(account.clone()))); + overlay.insert(path2.clone(), None); // Tombstone + + assert!(!overlay.is_empty()); + assert_eq!(overlay.len(), 2); + } + + #[test] + fn test_freeze_and_lookup() { + let mut mutable = OverlayStateMut::new(); + let path1 = Nibbles::from_nibbles([1, 2, 3, 4]); + let path2 = Nibbles::from_nibbles([5, 6, 7, 8]); + let account = test_account(); + + mutable.insert(path1.clone(), Some(OverlayValue::Account(account.clone()))); + mutable.insert(path2.clone(), None); + + let frozen = mutable.freeze(); + + assert_eq!(frozen.len(), 2); + + // Test lookup + let result1 = frozen.lookup(&path1); + assert!(result1.is_some()); + assert!(result1.unwrap().is_some()); + + let result2 = frozen.lookup(&path2); + assert!(result2.is_some()); + assert!(result2.unwrap().is_none()); // Tombstone + + let path3 = Nibbles::from_nibbles([9, 10, 11, 12]); + let result3 = frozen.lookup(&path3); + assert!(result3.is_none()); + } + + #[test] + fn test_deduplication_on_freeze() { + let mut mutable = OverlayStateMut::new(); + let path = Nibbles::from_nibbles([1, 2, 3, 4]); + let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // Insert multiple values for the same path + mutable.insert(path.clone(), Some(OverlayValue::Account(account1))); + mutable.insert(path.clone(), Some(OverlayValue::Account(account2.clone()))); + mutable.insert(path.clone(), None); // Tombstone should win + + let frozen = mutable.freeze(); + + assert_eq!(frozen.len(), 1); + let result = frozen.lookup(&path); + assert!(result.is_some()); + assert!(result.unwrap().is_none()); // Should be the tombstone + } + + #[test] + fn test_prefix_range() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + // Add some paths with common prefixes + mutable.insert( + Nibbles::from_nibbles([1, 2, 3, 4]), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles([1, 2, 5, 6]), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles([1, 3, 7, 8]), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles([2, 3, 4, 5]), + Some(OverlayValue::Account(account.clone())), + ); + + let frozen = mutable.freeze(); + + // Find entries with prefix [1, 2] + let prefix = Nibbles::from_nibbles([1, 2]); + let subset = frozen.sub_slice_for_prefix(&prefix); + + assert_eq!(subset.len(), 2); // Should match first two entries + + // Find entries with prefix [1] + let prefix = Nibbles::from_nibbles([1]); + let subset = frozen.sub_slice_for_prefix(&prefix); + + assert_eq!(subset.len(), 3); // Should match first three entries + } + + #[test] + fn test_contains_prefix_of() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + mutable.insert( + Nibbles::from_nibbles([1, 2, 3, 4]), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles([1, 2, 5, 6]), + Some(OverlayValue::Account(account.clone())), + ); + + let frozen = mutable.freeze(); + + // Test exact match + assert!(frozen.contains_prefix_of(&[1, 2, 3, 4])); + assert!(frozen.contains_prefix_of(&[1, 2, 5, 6])); + + // Test child path + assert!(frozen.contains_prefix_of(&[1, 2, 3, 4, 5])); + assert!(frozen.contains_prefix_of(&[1, 2, 5, 6, 15, 14, 13, 12])); + + // Test parent path + assert!(!frozen.contains_prefix_of(&[1, 2])); + assert!(!frozen.contains_prefix_of(&[1, 2, 3])); + assert!(!frozen.contains_prefix_of(&[1, 2, 5])); + + // Test unrelated path + assert!(!frozen.contains_prefix_of(&[7, 8, 9, 10])); + } + + #[test] + fn test_empty_overlay() { + let empty_mutable = OverlayStateMut::new(); + let frozen = empty_mutable.freeze(); + + assert!(frozen.is_empty()); + assert_eq!(frozen.len(), 0); + + let path = Nibbles::from_nibbles([1, 2, 3, 4]); + assert!(frozen.lookup(&path).is_none()); + assert!(!frozen.contains_prefix_of(&path)); + + let subset = frozen.sub_slice_for_prefix(&path); + assert!(subset.is_empty()); + } + + #[test] + fn test_prefix_offset_functionality() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + // Add paths that simulate account (64 nibbles) + storage slot (64 nibbles) = 128 nibbles + // total Account path: [1, 2, 3, 4] (simulating first 4 nibbles of 64-nibble account + // path) Storage paths extend the account path with storage slot nibbles + let account_prefix = vec![1, 2, 3, 4]; + + // Storage path 1: account + [5, 6] + let mut storage_path1 = account_prefix.clone(); + storage_path1.extend([5, 6]); + + // Storage path 2: account + [7, 8] + let mut storage_path2 = account_prefix.clone(); + storage_path2.extend([7, 8]); + + // Storage path 3: same account + [13, 14] (changed from different account to same account) + let mut storage_path3 = account_prefix.clone(); + storage_path3.extend([13, 14]); + + mutable.insert( + Nibbles::from_nibbles(storage_path1.clone()), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles(storage_path2.clone()), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles(storage_path3.clone()), + Some(OverlayValue::Account(account.clone())), + ); + + let frozen = mutable.freeze(); + + // Test with prefix_offset = 4 (simulating stripping 4-nibble account prefix for 2-nibble + // storage keys) + let storage_overlay = frozen.with_prefix_offset(4); + + // Test lookup with offset + let storage_key1 = Nibbles::from_nibbles([5, 6]); // Should find storage_path1 + let storage_key2 = Nibbles::from_nibbles([7, 8]); // Should find storage_path2 + let storage_key3 = Nibbles::from_nibbles([13, 14]); // Should find storage_path3 + let storage_key_missing = Nibbles::from_nibbles([15, 15]); // Should not find + + assert!(storage_overlay.lookup(&storage_key1).is_some()); + assert!(storage_overlay.lookup(&storage_key2).is_some()); + assert!(storage_overlay.lookup(&storage_key3).is_some()); + assert!(storage_overlay.lookup(&storage_key_missing).is_none()); + + // Test iter with offset - should return offset-adjusted paths + let iter_results: Vec<_> = storage_overlay.iter().collect(); + assert_eq!(iter_results.len(), 3); + + // The paths should be the storage-relative parts (after prefix_offset) + let expected_paths = [ + Nibbles::from_nibbles([5, 6]), + Nibbles::from_nibbles([7, 8]), + Nibbles::from_nibbles([13, 14]), + ]; + + for (i, (path, _)) in iter_results.iter().enumerate() { + assert_eq!(*path, expected_paths[i].as_slice()); + } + } + + #[test] + fn test_prefix_offset_sub_slice_methods() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + // Create overlay with account-prefixed storage paths for a SINGLE account + let _account_prefix = [1, 0, 0, 0]; // 4-nibble account prefix + + let paths = [ + [1, 0, 0, 0, 2, 0], // account + storage [2, 0] + [1, 0, 0, 0, 5, 0], // account + storage [5, 0] + [1, 0, 0, 0, 8, 0], // account + storage [8, 0] + [1, 0, 0, 0, 9, 0], // account + storage [9, 0] + ]; + + for path in &paths { + mutable + .insert(Nibbles::from_nibbles(*path), Some(OverlayValue::Account(account.clone()))); + } + + let frozen = mutable.freeze(); + let storage_overlay = frozen.with_prefix_offset(4); // Strip 4-nibble account prefix + + // Test sub_slice_before_prefix + let split_point = [5, 0]; // Storage key [5, 0] + + let (before, with_prefix, after) = storage_overlay.sub_slice_by_prefix(&split_point); + + // Before should contain storage keys [2, 0] (< [5, 0]) + assert_eq!(before.len(), 1); + let before_paths: Vec<_> = before.iter().map(|(path, _)| path).collect(); + assert!(before_paths.contains(&[2, 0].as_slice())); + + // With prefix should contain storage keys [5, 0] + assert_eq!(with_prefix.len(), 1); + let with_prefix_paths: Vec<_> = with_prefix.iter().map(|(path, _)| path).collect(); + assert!(with_prefix_paths.contains(&[5, 0].as_slice())); + + // After should contain storage keys [8, 0] and [9, 0] (strictly > [5, 0]) + assert_eq!(after.len(), 2); + let after_paths: Vec<_> = after.iter().map(|(path, _)| path).collect(); + assert!(!after_paths.contains(&[5, 0].as_slice())); // Strictly after, so [5, 0] not included + assert!(after_paths.contains(&[8, 0].as_slice())); + assert!(after_paths.contains(&[9, 0].as_slice())); + } + + #[test] + fn test_prefix_offset_find_prefix_range() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + // Create storage overlays for two different accounts + let account1_prefix = [1, 0, 0, 0]; // Account 1: 4 nibbles + let _account2_prefix = [2, 0, 0, 0]; // Account 2: 4 nibbles + + // Storage for account 1 + mutable.insert( + Nibbles::from_nibbles([1, 0, 0, 0, 5, 5]), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles([1, 0, 0, 0, 5, 6]), + Some(OverlayValue::Account(account.clone())), + ); + + // Storage for account 2 + mutable.insert( + Nibbles::from_nibbles([2, 0, 0, 0, 7, 7]), + Some(OverlayValue::Account(account.clone())), + ); + + let frozen = mutable.freeze(); + + // Test finding storage for account 1 using find_prefix_range on original overlay + let account1_storage = frozen.sub_slice_for_prefix(&account1_prefix); + assert_eq!(account1_storage.len(), 2); + + // Now test with prefix_offset - should convert to storage-relative paths + let storage_overlay = account1_storage.with_prefix_offset(4); + let storage_with_prefix5 = storage_overlay.sub_slice_for_prefix(&[5]); + + assert_eq!(storage_with_prefix5.len(), 2); + assert!(storage_with_prefix5.iter().any(|(path, _)| path == [5, 5].as_slice())); + assert!(storage_with_prefix5.iter().any(|(path, _)| path == [5, 6].as_slice())); + } + + #[test] + fn test_zero_copy_sub_slicing() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + // Add multiple entries + for i in 0..10 { + mutable.insert( + Nibbles::from_nibbles([i, 0, 0, 0]), + Some(OverlayValue::Account(account.clone())), + ); + } + + let frozen = mutable.freeze(); + assert_eq!(frozen.len(), 10); + + // Create sub-slices + let first_half = frozen.sub_slice(0, 5); + let second_half = frozen.sub_slice(5, 10); + let middle = frozen.sub_slice(2, 8); + + // Verify lengths + assert_eq!(first_half.len(), 5); + assert_eq!(second_half.len(), 5); + assert_eq!(middle.len(), 6); + + // Verify they share the same Arc (same pointer) + assert!(std::ptr::eq(frozen.data.as_ptr(), first_half.data.as_ptr())); + assert!(std::ptr::eq(frozen.data.as_ptr(), second_half.data.as_ptr())); + assert!(std::ptr::eq(frozen.data.as_ptr(), middle.data.as_ptr())); + + // Test recursive sub-slicing + let sub_sub = middle.sub_slice(1, 4); + assert_eq!(sub_sub.len(), 3); + assert!(std::ptr::eq(frozen.data.as_ptr(), sub_sub.data.as_ptr())); + + // Verify bounds are correct + assert_eq!(sub_sub.start_idx, 3); // middle starts at 2, +1 = 3 + assert_eq!(sub_sub.end_idx, 6); // middle starts at 2, +4 = 6 + } + + #[test] + fn test_thread_safety_with_arc() { + use std::thread; + + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + // Add entries + for i in 0..100 { + mutable.insert( + Nibbles::from_nibbles([i % 16, (i / 16) % 16, 0, 0]), + Some(OverlayValue::Account(account.clone())), + ); + } + + let frozen = mutable.freeze(); + let shared_overlay = Arc::new(frozen); + + // Spawn multiple threads that work with different sub-slices + let handles: Vec<_> = (0..4) + .map(|i| { + let overlay = Arc::clone(&shared_overlay); + thread::spawn(move || { + let start = i * 25; + let end = (i + 1) * 25; + let sub_slice = overlay.sub_slice(start, end); + + // Each thread verifies its slice + assert_eq!(sub_slice.len(), 25); + + // Count entries (just to do some work) + let count = sub_slice.iter().count(); + assert_eq!(count, 25); + + count + }) + }) + .collect(); + + // Wait for all threads and collect results + let results: Vec<_> = handles.into_iter().map(|h| h.join().unwrap()).collect(); + assert_eq!(results, vec![25, 25, 25, 25]); + } + + #[test] + fn test_sub_slice_by_prefix() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + // Create a diverse set of paths to test prefix partitioning + // Use vectors instead of arrays to avoid size mismatch issues + let test_paths = [ + // Before prefix [1, 2] + (vec![0, 5, 6, 7], Some(OverlayValue::Account(account.clone()))), + (vec![1, 0, 0, 0], Some(OverlayValue::Account(account.clone()))), + (vec![1, 1, 9, 9], Some(OverlayValue::Account(account.clone()))), + // Exact prefix [1, 2] and extensions + (vec![1, 2], Some(OverlayValue::Account(account.clone()))), + (vec![1, 2, 3, 4], Some(OverlayValue::Account(account.clone()))), + (vec![1, 2, 5, 6], Some(OverlayValue::Account(account.clone()))), + (vec![1, 2, 7, 8], None), // Tombstone with prefix + // After prefix [1, 2] + (vec![1, 3, 0, 0], Some(OverlayValue::Account(account.clone()))), + (vec![2, 0, 0, 0], Some(OverlayValue::Account(account.clone()))), + (vec![5, 5, 5, 5], Some(OverlayValue::Account(account.clone()))), + ]; + + for (path_nibbles, value) in test_paths.iter() { + let path = Nibbles::from_nibbles(path_nibbles.clone()); + mutable.insert(path, value.clone()); + } + + let frozen = mutable.freeze(); + let prefix = [1, 2]; + + // Test sub_slice_by_prefix + let (before, with_prefix, after) = frozen.sub_slice_by_prefix(&prefix); + + // Verify before slice - should contain paths < [1, 2] + assert_eq!(before.len(), 3); + let before_paths: Vec<_> = before.iter().map(|(path, _)| path).collect(); + assert!(before_paths.contains(&[0, 5, 6, 7].as_slice())); + assert!(before_paths.contains(&[1, 0, 0, 0].as_slice())); + assert!(before_paths.contains(&[1, 1, 9, 9].as_slice())); + + // Verify with_prefix slice - should contain paths that start with [1, 2] + assert_eq!(with_prefix.len(), 4); + let with_prefix_paths: Vec<_> = with_prefix.iter().map(|(path, _)| path).collect(); + assert!(with_prefix_paths.contains(&[1, 2].as_slice())); + assert!(with_prefix_paths.contains(&[1, 2, 3, 4].as_slice())); + assert!(with_prefix_paths.contains(&[1, 2, 5, 6].as_slice())); + assert!(with_prefix_paths.contains(&[1, 2, 7, 8].as_slice())); + + // Verify after slice - should contain paths > [1, 2, ...] range + assert_eq!(after.len(), 3); + let after_paths: Vec<_> = after.iter().map(|(path, _)| path).collect(); + assert!(after_paths.contains(&[1, 3, 0, 0].as_slice())); + assert!(after_paths.contains(&[2, 0, 0, 0].as_slice())); + assert!(after_paths.contains(&[5, 5, 5, 5].as_slice())); + + // Verify that all slices together contain all original entries + assert_eq!(before.len() + with_prefix.len() + after.len(), frozen.len()); + + // Test with empty overlay + let empty_frozen = OverlayStateMut::new().freeze(); + let (empty_before, empty_with, empty_after) = empty_frozen.sub_slice_by_prefix(&prefix); + assert!(empty_before.is_empty()); + assert!(empty_with.is_empty()); + assert!(empty_after.is_empty()); + + // Test with prefix that doesn't match anything + let no_match_prefix = Nibbles::from_nibbles([9, 9, 9]); + let (no_before, no_with, no_after) = frozen.sub_slice_by_prefix(&no_match_prefix); + assert_eq!(no_before.len(), frozen.len()); // All entries should be "before" + assert!(no_with.is_empty()); + assert!(no_after.is_empty()); + + // Test edge case: prefix with all 0xF's (should handle increment() returning None) + let max_prefix = Nibbles::from_nibbles([0xF, 0xF]); + let (max_before, max_with, max_after) = frozen.sub_slice_by_prefix(&max_prefix); + // All our test entries should be before 0xFF + assert_eq!(max_before.len(), frozen.len()); + assert!(max_with.is_empty()); + assert!(max_after.is_empty()); + } + + #[test] + fn test_sub_slice_by_prefix_with_exact_matches() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + // Test case where we have exact prefix matches and no extensions + let prefix = Nibbles::from_nibbles([2, 3]); + + mutable.insert(Nibbles::from_nibbles([1, 9]), Some(OverlayValue::Account(account.clone()))); + mutable.insert(Nibbles::from_nibbles([2, 2]), Some(OverlayValue::Account(account.clone()))); + mutable.insert(prefix.clone(), Some(OverlayValue::Account(account.clone()))); // Exact match + mutable.insert(Nibbles::from_nibbles([2, 4]), Some(OverlayValue::Account(account.clone()))); + mutable.insert(Nibbles::from_nibbles([3, 0]), Some(OverlayValue::Account(account.clone()))); + + let frozen = mutable.freeze(); + let (before, with_prefix, after) = frozen.sub_slice_by_prefix(&prefix); + + assert_eq!(before.len(), 2); // [1,9] and [2,2] + assert_eq!(with_prefix.len(), 1); // [2,3] exactly + assert_eq!(after.len(), 2); // [2,4] and [3,0] + + // Verify the exact match is in with_prefix + let with_prefix_paths: Vec<_> = with_prefix.iter().map(|(path, _)| path).collect(); + assert!(with_prefix_paths.contains(&prefix.as_slice())); + } + + #[test] + fn test_sub_slice_by_prefix_with_prefix_offset() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + // Create paths that simulate account (4 nibbles) + storage (4 nibbles) patterns + let _account_prefix = [1, 0, 0, 0]; + + // Storage paths for the same account + let storage_paths = [ + [1, 0, 0, 0, 2, 0, 0, 0], // account + storage [2,0,0,0] + [1, 0, 0, 0, 5, 0, 0, 0], // account + storage [5,0,0,0] + [1, 0, 0, 0, 5, 1, 0, 0], // account + storage [5,1,0,0] + [1, 0, 0, 0, 8, 0, 0, 0], // account + storage [8,0,0,0] + ]; + + for path in storage_paths.iter() { + mutable + .insert(Nibbles::from_nibbles(path), Some(OverlayValue::Account(account.clone()))); + } + + let frozen = mutable.freeze(); + + // Apply prefix offset to strip the account prefix + let storage_overlay = frozen.with_prefix_offset(4); + + // Test partitioning by storage prefix [5] + let storage_prefix = [5]; + let (before, with_prefix, after) = storage_overlay.sub_slice_by_prefix(&storage_prefix); + + // Before should contain [2,0,0,0] + assert_eq!(before.len(), 1); + let before_paths: Vec<_> = before.iter().map(|(path, _)| path).collect(); + assert!(before_paths.contains(&[2, 0, 0, 0].as_slice())); + + // With prefix should contain [5,0,0,0] and [5,1,0,0] + assert_eq!(with_prefix.len(), 2); + let with_prefix_paths: Vec<_> = with_prefix.iter().map(|(path, _)| path).collect(); + assert!(with_prefix_paths.contains(&[5, 0, 0, 0].as_slice())); + assert!(with_prefix_paths.contains(&[5, 1, 0, 0].as_slice())); + + // After should contain [8,0,0,0] + assert_eq!(after.len(), 1); + let after_paths: Vec<_> = after.iter().map(|(path, _)| path).collect(); + assert!(after_paths.contains(&[8, 0, 0, 0].as_slice())); + } +} diff --git a/src/storage.rs b/src/storage.rs index b86be74..2638784 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -1,5 +1,6 @@ pub mod debug; pub mod engine; +pub mod overlay_root; pub mod proofs; mod test_utils; pub mod value; diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs new file mode 100644 index 0000000..482c04d --- /dev/null +++ b/src/storage/overlay_root.rs @@ -0,0 +1,1757 @@ +use std::rc::Rc; + +use crate::{ + account::Account, + context::TransactionContext, + node::{encode_account_leaf, Node, NodeKind}, + overlay::{OverlayState, OverlayValue}, + page::SlottedPage, + pointer::Pointer, + storage::engine::{Error, StorageEngine}, +}; +use alloy_primitives::{ + map::{B256Map, HashMap}, + B256, U256, +}; +use alloy_rlp::encode_fixed_size; +use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, EMPTY_ROOT_HASH}; +use arrayvec::ArrayVec; + +#[derive(Debug)] +enum TriePosition<'a> { + Node(Nibbles, Rc>, Node), + Pointer(Nibbles, Rc>, Pointer, bool), + None, +} + +struct TraversalStack<'a> { + stack: Vec<(TriePosition<'a>, OverlayState)>, +} + +impl<'a> TraversalStack<'a> { + fn new() -> Self { + Self { stack: vec![] } + } + + fn push_node( + &mut self, + path: Nibbles, + node: Node, + page: Rc>, + overlay: OverlayState, + ) { + self.push(TriePosition::Node(path, page, node), overlay); + } + + fn push_pointer( + &mut self, + path: Nibbles, + pointer: Pointer, + page: Rc>, + can_add_by_hash: bool, + overlay: OverlayState, + ) { + self.push(TriePosition::Pointer(path, page, pointer, can_add_by_hash), overlay); + } + + fn push_overlay(&mut self, overlay: OverlayState) { + self.push(TriePosition::None, overlay); + } + + fn push(&mut self, position: TriePosition<'a>, overlay: OverlayState) { + self.stack.push((position, overlay)); + } + + fn pop(&mut self) -> Option<(TriePosition<'a>, OverlayState)> { + self.stack.pop() + } +} + +#[derive(Debug)] +pub struct OverlayedRoot { + pub root: B256, + pub updated_branch_nodes: HashMap, + pub storage_branch_updates: B256Map>, +} + +impl OverlayedRoot { + pub fn new( + root: B256, + updated_branch_nodes: HashMap, + storage_branch_updates: B256Map>, + ) -> Self { + Self { root, updated_branch_nodes, storage_branch_updates } + } + + pub fn new_hash(root: B256) -> Self { + Self { + root, + updated_branch_nodes: HashMap::default(), + storage_branch_updates: B256Map::default(), + } + } +} +struct RootBuilder { + hash_builder: HashBuilder, + storage_branch_updates: B256Map>, +} + +impl RootBuilder { + fn new() -> Self { + Self { + hash_builder: HashBuilder::default().with_updates(true), + storage_branch_updates: B256Map::default(), + } + } + + fn add_leaf(&mut self, key: Nibbles, value: &[u8]) { + self.hash_builder.add_leaf(key, value); + } + + fn add_branch(&mut self, key: Nibbles, value: B256, stored_in_database: bool) { + self.hash_builder.add_branch(key, value, stored_in_database); + } + + fn add_storage_branch_updates( + &mut self, + account: B256, + updates: HashMap, + ) { + self.storage_branch_updates.insert(account, updates); + } + + fn finalize(self) -> OverlayedRoot { + let (mut hash_builder, updated_branch_nodes) = self.hash_builder.split(); + OverlayedRoot::new(hash_builder.root(), updated_branch_nodes, self.storage_branch_updates) + } +} + +impl StorageEngine { + pub fn compute_state_root_with_overlay( + &self, + context: &TransactionContext, + overlay: OverlayState, + ) -> Result { + if overlay.is_empty() { + return Ok(OverlayedRoot::new_hash(context.root_node_hash)); + } + + let mut root_builder = RootBuilder::new(); + + let root_page = if let Some(root_page_id) = context.root_node_page_id { + let page = self.get_page(context, root_page_id)?; + SlottedPage::try_from(page).unwrap() + } else { + self.add_overlay_to_root_builder(&mut root_builder, &overlay); + return Ok(root_builder.finalize()); + }; + + let root_node: Node = root_page.get_value(0)?; + let mut stack = TraversalStack::new(); + stack.push_node(root_node.prefix().clone(), root_node, Rc::new(root_page), overlay); + + self.compute_root_with_overlay(context, &mut stack, &mut root_builder)?; + + Ok(root_builder.finalize()) + } + + fn compute_root_with_overlay<'a>( + &'a self, + context: &TransactionContext, + stack: &mut TraversalStack<'a>, + root_builder: &mut RootBuilder, + ) -> Result<(), Error> { + // Depth first traversal of the trie, starting at the root node. + // This applies any overlay state to the trie, taking precedence over the trie's own values. + // Whenever a branch or leaf is known to be the final unchanged value, we can add it to the + // hash builder. + while let Some((position, overlay)) = stack.pop() { + match position { + TriePosition::None => { + // No trie position, process whatever is in the overlay + self.add_overlay_to_root_builder(root_builder, &overlay); + } + TriePosition::Pointer(path, page, pointer, can_add_by_hash) => { + if overlay.is_empty() && can_add_by_hash { + if let Some(hash) = pointer.rlp().as_hash() { + // No overlay, just add the pointer by hash + root_builder.add_branch(path, hash, true); + continue; + } + } + // We have an overlay, need to process the child + self.process_overlayed_child( + context, + overlay, + root_builder, + path, + &pointer, + page, + stack, + )?; + } + TriePosition::Node(path, page, node) => { + let (pre_overlay, matching_overlay, post_overlay) = + overlay.sub_slice_by_prefix(&path); + if pre_overlay.contains_prefix_of(&path) { + // The pre_overlay invalidates the current node, so we can simply add the + // full overlay. We need to process it all together, + // as the post_overlay may contain descendants of the pre_overlay. + self.add_overlay_to_root_builder(root_builder, &overlay); + continue; + } + + self.add_overlay_to_root_builder(root_builder, &pre_overlay); + // Defer the post_overlay to be processed after the node is traversed + stack.push_overlay(post_overlay); + + match node.into_kind() { + NodeKind::Branch { children } => { + if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = + matching_overlay.first() + { + if overlay_path == &path { + // the overlay invalidates the current node, so just add this + // and skip the rest of the db traversal + self.add_overlay_to_root_builder( + root_builder, + &matching_overlay, + ); + continue; + } + } + self.process_branch_node_with_overlay( + matching_overlay, + &path, + children, + page, + stack, + )?; + } + NodeKind::AccountLeaf { + nonce_rlp, + balance_rlp, + code_hash, + storage_root, + } => { + self.process_account_leaf_with_overlay( + context, + &matching_overlay, + root_builder, + path, + page, + nonce_rlp, + balance_rlp, + code_hash, + storage_root, + )?; + } + NodeKind::StorageLeaf { value_rlp } => { + if let Some((overlay_path, _)) = matching_overlay.first() { + if overlay_path == &path { + // the overlay invalidates the current node, so just add this + // and skip the rest of the db traversal + self.add_overlay_to_root_builder( + root_builder, + &matching_overlay, + ); + continue; + } + } + // Leaf node, add it to the hash builder + root_builder.add_leaf(path, &value_rlp); + } + } + } + } + } + Ok(()) + } + + fn process_branch_node_with_overlay<'a>( + &'a self, + mut overlay: OverlayState, + path: &Nibbles, + mut children: [Option; 16], + current_page: Rc>, + stack: &mut TraversalStack<'a>, + ) -> Result<(), Error> { + let mut child_data = ArrayVec::<_, 16>::new(); + + let mut minimum_possible_child_count = 0; + for idx in 0..16 { + let child_pointer = children[idx as usize].take(); + if child_pointer.is_none() && overlay.is_empty() { + continue; + } + + let mut child_path = path.clone(); + child_path.push(idx); + let (_, child_overlay, overlay_after_child) = overlay.sub_slice_by_prefix(&child_path); + + if child_pointer.is_some() && child_overlay.is_empty() { + minimum_possible_child_count += 1; + } else if let Some((_, Some(_))) = child_overlay.first() { + // we have a non-tombstone overlay, so there must be at least one descendant + // in this child index + minimum_possible_child_count += 1; + } + + child_data.push((child_path, child_pointer, child_overlay)); + overlay = overlay_after_child; + } + let can_add_by_hash = minimum_possible_child_count > 1; + + for (child_path, child_pointer, child_overlay) in child_data.into_iter().rev() { + match child_pointer { + Some(pointer) => { + stack.push_pointer( + child_path, + pointer, + current_page.clone(), + can_add_by_hash, + child_overlay, + ); + } + None => { + if child_overlay.is_empty() { + // nothing here to add + } else { + // we have a nonconflicting overlay, add all of it to the hash builder + stack.push_overlay(child_overlay); + } + } + } + } + Ok(()) + } + + fn process_account_leaf_with_overlay<'a>( + &'a self, + context: &TransactionContext, + overlay: &OverlayState, + root_builder: &mut RootBuilder, + path: Nibbles, + current_page: Rc>, + mut nonce_rlp: ArrayVec, + mut balance_rlp: ArrayVec, + mut code_hash: B256, + storage_root: Option, + ) -> Result<(), Error> { + let overlayed_account = overlay.lookup(&path); + match overlayed_account { + Some(None) => { + // The account is removed in the overlay + return Ok(()); + } + Some(Some(OverlayValue::Account(overlayed_account))) => { + // The account is updated in the overlay + nonce_rlp = alloy_rlp::encode_fixed_size(&overlayed_account.nonce); + balance_rlp = alloy_rlp::encode_fixed_size(&overlayed_account.balance); + code_hash = overlayed_account.code_hash; + } + _ => { + // The account is not updated in the overlay + } + }; + + let has_storage_overlays = overlay.iter().any(|(path, _)| path.len() > 64); + if !has_storage_overlays { + let storage_root_hash = storage_root + .as_ref() + .map_or(EMPTY_ROOT_HASH, |p| p.rlp().as_hash().unwrap_or(EMPTY_ROOT_HASH)); + + self.add_account_leaf_to_root_builder( + root_builder, + path, + &nonce_rlp, + &balance_rlp, + &code_hash, + &storage_root_hash, + ); + return Ok(()); + } + + let mut storage_root_builder = RootBuilder::new(); + + // We have storage overlays, need to compute a new storage root + let storage_overlay = overlay.with_prefix_offset(64); + + match storage_root { + Some(pointer) => { + let mut storage_stack = TraversalStack::new(); + + // load the root storage node + if let Some(child_cell) = pointer.location().cell_index() { + let root_storage_node: Node = current_page.get_value(child_cell)?; + storage_stack.push_node( + root_storage_node.prefix().clone(), + root_storage_node, + current_page, + storage_overlay, + ); + self.compute_root_with_overlay( + context, + &mut storage_stack, + &mut storage_root_builder, + )? + } else { + let storage_page = + self.get_page(context, pointer.location().page_id().unwrap())?; + let slotted_page = SlottedPage::try_from(storage_page)?; + let root_storage_node: Node = slotted_page.get_value(0)?; + storage_stack.push_node( + root_storage_node.prefix().clone(), + root_storage_node, + Rc::new(slotted_page), + storage_overlay, + ); + self.compute_root_with_overlay( + context, + &mut storage_stack, + &mut storage_root_builder, + )?; + } + } + None => { + // No existing storage root, just add the overlay + self.add_overlay_to_root_builder(&mut storage_root_builder, &storage_overlay); + } + }; + let (mut storage_hash_builder, updated_storage_branch_nodes) = + storage_root_builder.hash_builder.split(); + let new_root = storage_hash_builder.root(); + + root_builder.add_storage_branch_updates( + B256::from_slice(&path.pack()), + updated_storage_branch_nodes, + ); + + self.add_account_leaf_to_root_builder( + root_builder, + path, + &nonce_rlp, + &balance_rlp, + &code_hash, + &new_root, + ); + + Ok(()) + } + + fn add_account_leaf_to_root_builder( + &self, + root_builder: &mut RootBuilder, + path: Nibbles, + nonce_rlp: &ArrayVec, + balance_rlp: &ArrayVec, + code_hash: &B256, + storage_root: &B256, + ) { + let mut buf = [0u8; 110]; // max RLP length for an account: 2 bytes for list length, 9 for nonce, 33 for + // balance, 33 for storage root, 33 for code hash + let mut value_rlp = buf.as_mut(); + let account_rlp_length = + encode_account_leaf(nonce_rlp, balance_rlp, code_hash, storage_root, &mut value_rlp); + root_builder.add_leaf(path, &buf[..account_rlp_length]); + } + + fn process_overlayed_child<'a>( + &'a self, + context: &TransactionContext, + overlay: OverlayState, + root_builder: &mut RootBuilder, + mut child_path: Nibbles, + child: &Pointer, + current_page: Rc>, + stack: &mut TraversalStack<'a>, + ) -> Result<(), Error> { + // First consider the overlay. All values in it must already contain the child_path prefix. + // If the overlay matches the child path, we can add it to the hash builder and skip + // actually reading the child node. + // Account values cannot be directly overlayed, as they may need to be merged with the + // existing storage trie. + if let Some((overlay_path, overlay_value)) = overlay.first() { + if &child_path == overlay_path && + !matches!(overlay_value, Some(OverlayValue::Account(_))) + { + // the child path is directly overlayed, so only use the overlay state + self.add_overlay_to_root_builder(root_builder, &overlay); + return Ok(()); + } + } + + if let Some(child_cell) = child.location().cell_index() { + let child_node: Node = current_page.get_value(child_cell)?; + child_path.extend_from_slice(child_node.prefix()); + stack.push_node(child_path, child_node, current_page, overlay); + } else { + let child_page_id = child.location().page_id().unwrap(); + let child_page = self.get_page(context, child_page_id)?; + let child_slotted_page = SlottedPage::try_from(child_page).unwrap(); + let child_node: Node = child_slotted_page.get_value(0)?; + child_path.extend_from_slice(child_node.prefix()); + stack.push_node(child_path, child_node, Rc::new(child_slotted_page), overlay); + } + Ok(()) + } + + fn process_overlayed_account( + &self, + root_builder: &mut RootBuilder, + path: Nibbles, + account: &Account, + storage_overlay: OverlayState, + ) -> Result<(), Error> { + if storage_overlay.is_empty() { + let encoded = self.encode_account(account); + root_builder.add_leaf(path, &encoded); + return Ok(()); + } + + let mut storage_root_builder = RootBuilder::new(); + self.add_overlay_to_root_builder(&mut storage_root_builder, &storage_overlay); + + let (mut storage_hash_builder, updated_storage_branch_nodes) = + storage_root_builder.hash_builder.split(); + let storage_root = storage_hash_builder.root(); + + root_builder.add_storage_branch_updates( + B256::from_slice(&path.pack()), + updated_storage_branch_nodes, + ); + + let encoded = self.encode_account_with_root(account, storage_root); + root_builder.add_leaf(path, &encoded); + Ok(()) + } + + fn add_overlay_to_root_builder(&self, root_builder: &mut RootBuilder, overlay: &OverlayState) { + let mut last_processed_path: Option<&[u8]> = None; + for (path, value) in overlay.iter() { + if let Some(last_processed_path) = last_processed_path { + if path.starts_with(last_processed_path) { + // skip over all descendants of a processed path + continue; + } + } + + match value { + Some(OverlayValue::Account(account)) => { + let storage_overlay = overlay.sub_slice_for_prefix(path).with_prefix_offset(64); + self.process_overlayed_account( + root_builder, + Nibbles::from_nibbles(path), + account, + storage_overlay, + ) + .unwrap(); + last_processed_path = Some(path); + } + Some(OverlayValue::Storage(storage_value)) => { + let encoded = self.encode_storage(storage_value); + root_builder.add_leaf(Nibbles::from_nibbles(path), &encoded); + } + Some(OverlayValue::Hash(hash)) => { + root_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); + last_processed_path = Some(path); + } + None => { + // Tombstone - skip + last_processed_path = Some(path); + } + } + } + } + + #[inline] + pub fn encode_account(&self, account: &Account) -> ArrayVec { + let trie_account = Account { + nonce: account.nonce, + balance: account.balance, + storage_root: account.storage_root, + code_hash: account.code_hash, + }; + encode_fixed_size(&trie_account) + } + + #[inline] + pub fn encode_account_with_root(&self, account: &Account, root: B256) -> ArrayVec { + let trie_account = Account { + nonce: account.nonce, + balance: account.balance, + storage_root: root, + code_hash: account.code_hash, + }; + encode_fixed_size(&trie_account) + } + + #[inline] + pub fn encode_storage(&self, storage_value: &U256) -> ArrayVec { + encode_fixed_size(storage_value) + } +} + +#[cfg(test)] +mod tests { + use alloy_primitives::{address, Address, U256}; + use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY}; + use rand::Rng; + use tempdir::TempDir; + + use crate::{ + account::Account, + database::Database, + node::TrieValue, + overlay::{OverlayStateMut, OverlayValue}, + path::AddressPath, + }; + + use super::*; + + fn compare_overlay_with_committed_root( + db: &Database, + context: &mut TransactionContext, + overlay: &OverlayState, + ) -> B256 { + let initial_root = context.root_node_hash; + let output = + db.storage_engine.compute_state_root_with_overlay(context, overlay.clone()).unwrap(); + let (overlay_root, account_branch_updates, storage_branch_updates) = + (output.root, output.updated_branch_nodes, output.storage_branch_updates); + assert_ne!(overlay_root, initial_root, "Overlay should not match initial root"); + + let mut overlay_mut_with_branches = OverlayStateMut::new(); + + overlay.data().iter().for_each(|(path, value)| { + overlay_mut_with_branches.insert(path.clone(), value.clone()); + }); + + for (path, branch) in account_branch_updates.iter() { + if let Some(root_hash) = branch.root_hash { + overlay_mut_with_branches.insert(path.clone(), Some(OverlayValue::Hash(root_hash))); + } + let mut hash_idx = 0; + let mut path = path.clone(); + for i in 0..16 { + if branch.hash_mask.is_bit_set(i) { + path.push(i); + overlay_mut_with_branches + .insert(path.clone(), Some(OverlayValue::Hash(branch.hashes[hash_idx]))); + hash_idx += 1; + path.pop(); + } + } + } + + for (account, branches) in storage_branch_updates.iter() { + for (path, branch) in branches.iter() { + if let Some(root_hash) = branch.root_hash { + overlay_mut_with_branches.insert( + Nibbles::unpack(account).join(path), + Some(OverlayValue::Hash(root_hash)), + ); + } + let mut hash_idx = 0; + let mut path = path.clone(); + for i in 0..16 { + if branch.hash_mask.is_bit_set(i) { + path.push(i); + overlay_mut_with_branches.insert( + Nibbles::unpack(account).join(&path), + Some(OverlayValue::Hash(branch.hashes[hash_idx])), + ); + hash_idx += 1; + path.pop(); + } + } + } + } + + let overlay_with_branches = overlay_mut_with_branches.freeze(); + + let output = db + .storage_engine + .compute_state_root_with_overlay(context, overlay_with_branches.clone()) + .unwrap(); + let (overlay_root_with_branches, _, _) = + (output.root, output.updated_branch_nodes, output.storage_branch_updates); + assert_eq!(overlay_root_with_branches, overlay_root); + + let mut changes: Vec<(Nibbles, Option)> = overlay + .data() + .iter() + .map(|(path, value)| (path.clone(), value.clone().map(|v| v.try_into().unwrap()))) + .collect(); + db.storage_engine.set_values(context, &mut changes).unwrap(); + let committed_root = context.root_node_hash; + assert_eq!(overlay_root, committed_root, "Overlay should match committed root"); + + // recompute the root with overlayed state that is already committed. This should match the + // committed root. + let output = db + .storage_engine + .compute_state_root_with_overlay(context, overlay_with_branches) + .unwrap(); + let (overlay_root_after_commit, _, _) = + (output.root, output.updated_branch_nodes, output.storage_branch_updates); + assert_eq!(overlay_root_after_commit, committed_root); + + overlay_root + } + + #[test] + fn test_empty_overlay_root() { + let tmp_dir = TempDir::new("test_overlay_root_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let context = db.storage_engine.read_context(); + let empty_overlay = OverlayStateMut::new().freeze(); + + let output = + db.storage_engine.compute_state_root_with_overlay(&context, empty_overlay).unwrap(); + assert_eq!(output.root, context.root_node_hash); + } + + #[test] + fn test_overlay_root_with_empty_db() { + let tmp_dir = TempDir::new("test_overlay_root_changes_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let context = db.storage_engine.read_context(); + + // Create overlay with some changes + let mut overlay_mut = OverlayStateMut::new(); + let address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045"); + let address_path = AddressPath::for_address(address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + overlay_mut.insert(address_path.into(), Some(OverlayValue::Account(account))); + let overlay = overlay_mut.freeze(); + + // Compute root with overlay + let output = db.storage_engine.compute_state_root_with_overlay(&context, overlay).unwrap(); + + // The root should be different from the empty root (since we have changes) + assert_ne!(output.root, EMPTY_ROOT_HASH); + } + + #[test] + fn test_overlay_root_with_changes() { + let tmp_dir = TempDir::new("test_overlay_root_changes_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // First, add an account using set_values (following the same pattern as other tests) + let address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045"); + let address_path = AddressPath::for_address(address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + db.storage_engine + .set_values( + &mut context, + &mut [(address_path.clone().into(), Some(TrieValue::Account(account)))], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + assert_ne!(initial_root, EMPTY_ROOT_HASH); + + // Now test with actual overlay changes - modify the same account with different values + let mut overlay_mut = OverlayStateMut::new(); + let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + overlay_mut + .insert(address_path.clone().into(), Some(OverlayValue::Account(account2.clone()))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_overlay_with_controlled_paths() { + let tmp_dir = TempDir::new("test_overlay_controlled_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create specific address paths to control trie structure + // These paths are designed to create branch nodes at specific positions + + // Path 1: starts with 0x0... (first nibble = 0) + let path1 = AddressPath::new(Nibbles::from_nibbles([ + 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + ])); + + // Path 2: starts with 0x1... (first nibble = 1) + let path2 = AddressPath::new(Nibbles::from_nibbles([ + 0x1, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + ])); + + let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // Add both accounts to disk - this should create a branch node at the root + db.storage_engine + .set_values( + &mut context, + &mut [ + (path1.clone().into(), Some(TrieValue::Account(account1.clone()))), + (path2.clone().into(), Some(TrieValue::Account(account2.clone()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + assert_ne!(initial_root, EMPTY_ROOT_HASH); + + // Test Case 1: Overlay that affects only path1 (child 0) - path2 subtree should use + // add_branch optimization + let account1_updated = Account::new(10, U256::from(1000), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut + .insert(path1.clone().into(), Some(OverlayValue::Account(account1_updated.clone()))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + + // Test Case 2: Overlay that creates a new account in an empty subtree (None child case), + // affects an existing subtree, and leaves one unaffected Path 3: starts with 0x2... + // (first nibble = 2) - this child doesn't exist on disk + let path3 = AddressPath::new(Nibbles::from_nibbles([ + 0x2, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + ])); + + let account3 = Account::new(3, U256::from(300), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let mut overlay_mut2 = OverlayStateMut::new(); + overlay_mut2.insert(path1.clone().into(), Some(OverlayValue::Account(account3.clone()))); + overlay_mut2.insert(path3.clone().into(), Some(OverlayValue::Account(account3.clone()))); + let overlay2 = overlay_mut2.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay2); + } + + #[test] + fn test_overlay_tombstones() { + let tmp_dir = TempDir::new("test_overlay_tombstones_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + // Step 1: Write account 1 and compute root + let mut context = db.storage_engine.write_context(); + let path1 = AddressPath::new(Nibbles::from_nibbles([ + 0x0, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + ])); + let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // Step 2: Add account 2 + let path2 = AddressPath::new(Nibbles::from_nibbles([ + 0x1, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + ])); + let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // Step 3: Add account 3 + let path3 = AddressPath::new(Nibbles::from_nibbles([ + 0x2, 0x0, 0x0, 0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, + ])); + let account3 = Account::new(3, U256::from(300), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + db.storage_engine + .set_values( + &mut context, + &mut [ + (path1.clone().into(), Some(TrieValue::Account(account1.clone()))), + (path3.clone().into(), Some(TrieValue::Account(account3.clone()))), + ], + ) + .unwrap(); + let root_without_account2 = context.root_node_hash; + + db.storage_engine + .set_values( + &mut context, + &mut [ + (path1.clone().into(), Some(TrieValue::Account(account1.clone()))), + (path2.clone().into(), Some(TrieValue::Account(account2.clone()))), + (path3.clone().into(), Some(TrieValue::Account(account3.clone()))), + ], + ) + .unwrap(); + let root_with_all_accounts = context.root_node_hash; + assert_ne!(root_with_all_accounts, root_without_account2); + + // Step 4: Overlay a tombstone for account 2 and compute root + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(path2.clone().into(), None); // Delete account2 + let overlay = overlay_mut.freeze(); + + let overlay_root = compare_overlay_with_committed_root(&db, &mut context, &overlay); + + assert_eq!( + overlay_root, root_without_account2, + "After deleting account2, committed root should match original single-account root" + ); + } + + #[test] + fn test_overlay_with_storage_changes() { + let tmp_dir = TempDir::new("test_overlay_storage_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create an account with some storage + let account_address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045"); + let account_path = AddressPath::for_address(account_address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // Create storage paths for the account + let storage_key1 = U256::from(42); + let storage_key2 = U256::from(99); + let storage_path1 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key1.into()); + let storage_path2 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key2.into()); + + let storage_value1 = U256::from(123); + let storage_value2 = U256::from(456); + + // Set up initial state with account and storage + db.storage_engine + .set_values( + &mut context, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1))), + (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2))), + ], + ) + .unwrap(); + + // Test Case 1: Overlay that modifies existing storage + let mut overlay_mut = OverlayStateMut::new(); + let new_storage_value1 = U256::from(999); + overlay_mut + .insert(storage_path1.full_path(), Some(OverlayValue::Storage(new_storage_value1))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + + // Test Case 2: Overlay that adds new storage + let mut overlay_mut2 = OverlayStateMut::new(); + let storage_key3 = U256::from(200); + let storage_path3 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key3.into()); + let storage_value3 = U256::from(789); + overlay_mut2.insert(storage_path3.full_path(), Some(OverlayValue::Storage(storage_value3))); + let overlay2 = overlay_mut2.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay2); + + // Test Case 3: Overlay that deletes storage (tombstone) + let mut overlay_mut3 = OverlayStateMut::new(); + overlay_mut3.insert(storage_path2.full_path(), None); // Delete storage slot + let overlay3 = overlay_mut3.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay3); + + // Test Case 4: Combined account and storage changes + let mut overlay_mut4 = OverlayStateMut::new(); + let updated_account = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + overlay_mut4.insert( + account_path.clone().into(), + Some(OverlayValue::Account(updated_account.clone())), + ); + overlay_mut4 + .insert(storage_path1.full_path(), Some(OverlayValue::Storage(new_storage_value1))); + let overlay4 = overlay_mut4.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay4); + + // Test Case 5: Overlay that deletes storage slot via a zero value + let mut overlay_mut5 = OverlayStateMut::new(); + overlay_mut5.insert(storage_path1.full_path(), Some(OverlayValue::Storage(U256::ZERO))); + let overlay5 = overlay_mut5.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay5); + } + + #[test] + fn test_debug_adding_storage_slot_overlay() { + let tmp_dir = TempDir::new("test_add_storage_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create account with 1 storage slot + let account_address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(account_address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + let storage_key1 = U256::from(10); + let storage_path1 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key1.into()); + + // Set up initial state with 1 storage slot + db.storage_engine + .set_values( + &mut context, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), + ], + ) + .unwrap(); + + // Test: Add a NEW storage slot via overlay + let mut overlay_mut = OverlayStateMut::new(); + let storage_key2 = U256::from(20); // New storage key + let storage_path2 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key2.into()); + + overlay_mut.insert(storage_path2.full_path(), Some(OverlayValue::Storage(U256::from(222)))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_overlay_account_with_storage() { + let tmp_dir = TempDir::new("test_overlay_account_with_storage_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create an account with some storage + let account_address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(account_address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + let storage_key = U256::from(10); + let storage_path = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key.into()); + + // Set up initial state with account and storage + db.storage_engine + .set_values( + &mut context, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + (storage_path.full_path(), Some(TrieValue::Storage(U256::from(111)))), + ], + ) + .unwrap(); + + // Test: Overlay that modifies the account value (but not the storage root) + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert( + account_path.clone().into(), + Some(OverlayValue::Account(Account::new( + 2, + U256::from(200), + EMPTY_ROOT_HASH, + KECCAK_EMPTY, + ))), + ); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_debug_minimal_multi_account_overlay() { + let tmp_dir = TempDir::new("test_minimal_multi_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create just 2 accounts with 1 storage slot each + let account1_address = address!("0x0000000000000000000000000000000000000001"); + let account2_address = address!("0x0000000000000000000000000000000000000002"); + + let account1_path = AddressPath::for_address(account1_address); + let account2_path = AddressPath::for_address(account2_address); + + let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // One storage slot for each account + let storage1_key = U256::from(10); + let storage2_key = U256::from(20); + let storage1_path = + crate::path::StoragePath::for_address_and_slot(account1_address, storage1_key.into()); + let storage2_path = + crate::path::StoragePath::for_address_and_slot(account2_address, storage2_key.into()); + + // Set up initial state + db.storage_engine + .set_values( + &mut context, + &mut [ + (account1_path.clone().into(), Some(TrieValue::Account(account1.clone()))), + (account2_path.clone().into(), Some(TrieValue::Account(account2.clone()))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222)))), + ], + ) + .unwrap(); + + // Test: Modify just one storage value per account via overlay + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(storage1_path.full_path(), Some(OverlayValue::Storage(U256::from(999)))); + overlay_mut.insert(storage2_path.full_path(), Some(OverlayValue::Storage(U256::from(888)))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_debug_multiple_storage_overlays_same_account() { + let tmp_dir = TempDir::new("test_debug_multiple_overlays_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create one account with 2 initial storage slots + let account_address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045"); + let account_path = AddressPath::for_address(account_address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + let storage_key1 = U256::from(10); + let storage_key2 = U256::from(20); + let storage_path1 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key1.into()); + let storage_path2 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key2.into()); + + // Set up initial state + db.storage_engine + .set_values( + &mut context, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), + ], + ) + .unwrap(); + + // Test: Apply MULTIPLE storage overlays to the same account + let mut overlay_mut = OverlayStateMut::new(); + + // Modify existing storage slot 1 + overlay_mut + .insert(storage_path1.full_path(), Some(OverlayValue::Storage(U256::from(1111)))); + + // Add new storage slot 3 + let storage_key3 = U256::from(40); + let storage_path3 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key3.into()); + + overlay_mut.insert(storage_path3.full_path(), Some(OverlayValue::Storage(U256::from(444)))); + + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_debug_overlay_vs_committed_single_account() { + let tmp_dir = TempDir::new("test_debug_overlay_committed_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create one account with 2 storage slots + let account_address = address!("0x0000000000000000000000000000000000000001"); + let account_path = AddressPath::for_address(account_address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + let storage_key1 = U256::from(10); + let storage_key2 = U256::from(20); + let storage_path1 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key1.into()); + let storage_path2 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key2.into()); + + // Set up initial state with 2 storage slots + db.storage_engine + .set_values( + &mut context, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), + ], + ) + .unwrap(); + + // Test: Overlay that modifies ONLY ONE storage slot, leaving the other unchanged + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(storage_path1.full_path(), Some(OverlayValue::Storage(U256::from(999)))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_multiple_accounts_with_storage_overlays() { + let tmp_dir = TempDir::new("test_multi_account_storage_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create two accounts with different storage + let account1_address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045"); + let account2_address = address!("0x1234567890abcdef1234567890abcdef12345678"); + + let account1_path = AddressPath::for_address(account1_address); + let account2_path = AddressPath::for_address(account2_address); + + let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // Storage for account1 + let storage1_key1 = U256::from(10); + let storage1_key2 = U256::from(20); + let storage1_path1 = + crate::path::StoragePath::for_address_and_slot(account1_address, storage1_key1.into()); + let storage1_path2 = + crate::path::StoragePath::for_address_and_slot(account1_address, storage1_key2.into()); + + // Storage for account2 + let storage2_key1 = U256::from(30); + let storage2_path1 = + crate::path::StoragePath::for_address_and_slot(account2_address, storage2_key1.into()); + + // Set up initial state + db.storage_engine + .set_values( + &mut context, + &mut [ + (account1_path.clone().into(), Some(TrieValue::Account(account1.clone()))), + (account2_path.clone().into(), Some(TrieValue::Account(account2.clone()))), + (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), + (storage2_path1.full_path(), Some(TrieValue::Storage(U256::from(333)))), + ], + ) + .unwrap(); + + // Test: Overlay changes to both accounts' storage + let mut overlay_mut = OverlayStateMut::new(); + + // Modify account1's storage + overlay_mut + .insert(storage1_path1.full_path(), Some(OverlayValue::Storage(U256::from(1111)))); + + // Add new storage to account1 + let storage1_key3 = U256::from(40); + let storage1_path3 = + crate::path::StoragePath::for_address_and_slot(account1_address, storage1_key3.into()); + overlay_mut + .insert(storage1_path3.full_path(), Some(OverlayValue::Storage(U256::from(444)))); + + // Modify account2's storage + overlay_mut + .insert(storage2_path1.full_path(), Some(OverlayValue::Storage(U256::from(3333)))); + + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_debug_multi_account_storage() { + let tmp_dir = TempDir::new("test_debug_multi_account_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create two accounts with very simple, distinct addresses + let account1_address = address!("0x0000000000000000000000000000000000000001"); + let account2_address = address!("0x0000000000000000000000000000000000000002"); + + let account1_path = AddressPath::for_address(account1_address); + let account2_path = AddressPath::for_address(account2_address); + + let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // One storage slot for each account + let storage1_key = U256::from(10); + let storage2_key = U256::from(20); + let storage1_path = + crate::path::StoragePath::for_address_and_slot(account1_address, storage1_key.into()); + let storage2_path = + crate::path::StoragePath::for_address_and_slot(account2_address, storage2_key.into()); + + // Set up initial state + db.storage_engine + .set_values( + &mut context, + &mut [ + (account1_path.clone().into(), Some(TrieValue::Account(account1.clone()))), + (account2_path.clone().into(), Some(TrieValue::Account(account2.clone()))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222)))), + ], + ) + .unwrap(); + + // Test: Modify just one storage slot for account1 + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(storage1_path.full_path(), Some(OverlayValue::Storage(U256::from(999)))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_debug_both_accounts_storage_change() { + let tmp_dir = TempDir::new("test_debug_both_accounts_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create two accounts with simple addresses + let account1_address = address!("0x0000000000000000000000000000000000000001"); + let account2_address = address!("0x0000000000000000000000000000000000000002"); + + let account1_path = AddressPath::for_address(account1_address); + let account2_path = AddressPath::for_address(account2_address); + + let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // One storage slot for each account + let storage1_key = U256::from(10); + let storage2_key = U256::from(20); + let storage1_path = + crate::path::StoragePath::for_address_and_slot(account1_address, storage1_key.into()); + let storage2_path = + crate::path::StoragePath::for_address_and_slot(account2_address, storage2_key.into()); + + // Set up initial state + db.storage_engine + .set_values( + &mut context, + &mut [ + (account1_path.clone().into(), Some(TrieValue::Account(account1.clone()))), + (account2_path.clone().into(), Some(TrieValue::Account(account2.clone()))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222)))), + ], + ) + .unwrap(); + + // Test: Modify storage for BOTH accounts + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert(storage1_path.full_path(), Some(OverlayValue::Storage(U256::from(999)))); + overlay_mut.insert(storage2_path.full_path(), Some(OverlayValue::Storage(U256::from(888)))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_debug_adding_new_storage_multi_account() { + let tmp_dir = TempDir::new("test_debug_new_storage_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create two accounts (similar to the original failing test) + let account1_address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045"); + let account2_address = address!("0x1234567890abcdef1234567890abcdef12345678"); + + let account1_path = AddressPath::for_address(account1_address); + let account2_path = AddressPath::for_address(account2_address); + + let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let account2 = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // Initial storage + let storage1_key1 = U256::from(10); + let storage1_path1 = + crate::path::StoragePath::for_address_and_slot(account1_address, storage1_key1.into()); + + // Set up initial state with just one storage slot + db.storage_engine + .set_values( + &mut context, + &mut [ + (account1_path.clone().into(), Some(TrieValue::Account(account1.clone()))), + (account2_path.clone().into(), Some(TrieValue::Account(account2.clone()))), + (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), + ], + ) + .unwrap(); + + // Test: Add NEW storage to account1 + let mut overlay_mut = OverlayStateMut::new(); + let storage1_key2 = U256::from(40); // New storage key + let storage1_path2 = + crate::path::StoragePath::for_address_and_slot(account1_address, storage1_key2.into()); + + overlay_mut + .insert(storage1_path2.full_path(), Some(OverlayValue::Storage(U256::from(444)))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_storage_overlay_with_empty_account() { + let tmp_dir = TempDir::new("test_empty_account_storage_db").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create an account with no initial storage + let account_address = address!("0xd8da6bf26964af9d7eed9e03e53415d37aa96045"); + let account_path = AddressPath::for_address(account_address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // Set up initial state with just the account (no storage) + db.storage_engine + .set_values( + &mut context, + &mut [(account_path.clone().into(), Some(TrieValue::Account(account.clone())))], + ) + .unwrap(); + + // Test: Add storage to account that had no storage before + let mut overlay_mut = OverlayStateMut::new(); + let storage_key = U256::from(42); + let storage_path = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key.into()); + let storage_value = U256::from(123); + overlay_mut.insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value))); + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_1000_accounts_with_10_overlay() { + for _ in 0..10 { + let tmp_dir = TempDir::new("test_1000_accounts_with_10_overlay").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + let mut rng = rand::rng(); + + let mut changes: Vec<(Nibbles, Option)> = Vec::with_capacity(1000); + + for i in 0..1000 { + let account_address = Address::random(); + let account = + Account::new(i, U256::from(rng.random::()), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let account_path = AddressPath::for_address(account_address); + + changes.push((account_path.into(), Some(TrieValue::Account(account)))); + } + + changes.sort_by(|a, b| a.0.cmp(&b.0)); + + db.storage_engine.set_values(&mut context, &mut changes).unwrap(); + + // Create overlay with modifications to every 100th account + let mut overlay_mut = OverlayStateMut::new(); + + // Take every 100th account from the changes + for (i, (path, value)) in changes.iter().step_by(100).enumerate() { + if let Some(TrieValue::Account(account)) = value { + if i % 2 == 0 { + // For half of the sampled accounts, create new modified account + let mut new_account = account.clone(); + new_account.balance = U256::from(rng.random::()); // Random new balance + overlay_mut.insert(path.clone(), Some(OverlayValue::Account(new_account))); + } else { + // For other half, mark for deletion + overlay_mut.insert(path.clone(), None); + } + } + } + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + } + + #[test] + fn test_overlay_new_account_with_storage() { + let tmp_dir = TempDir::new("test_overlay_new_account_with_storage").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + let account_address = Address::random(); + let account_path = AddressPath::for_address(account_address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + db.storage_engine + .set_values( + &mut context, + &mut [(account_path.clone().into(), Some(TrieValue::Account(account.clone())))], + ) + .unwrap(); + + let mut overlay_mut = OverlayStateMut::new(); + let new_address = Address::random(); + let new_account_path = AddressPath::for_address(new_address); + let new_account = Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + overlay_mut.insert( + new_account_path.clone().into(), + Some(OverlayValue::Account(new_account.clone())), + ); + + let storage_key1 = B256::right_padding_from(&[1, 1, 2, 3]); + let storage_path1 = crate::path::StoragePath::for_address_path_and_slot_hash( + new_account_path.clone(), + Nibbles::unpack(storage_key1), + ); + let storage_value1 = U256::from(123); + overlay_mut.insert(storage_path1.full_path(), Some(OverlayValue::Storage(storage_value1))); + + let storage_key2 = B256::right_padding_from(&[1, 1, 2, 0]); + let storage_path2 = crate::path::StoragePath::for_address_path_and_slot_hash( + new_account_path.clone(), + Nibbles::unpack(storage_key2), + ); + let storage_value2 = U256::from(234); + overlay_mut.insert(storage_path2.full_path(), Some(OverlayValue::Storage(storage_value2))); + + let storage_key3 = B256::right_padding_from(&[2, 2, 0, 0]); + let storage_path3 = crate::path::StoragePath::for_address_path_and_slot_hash( + new_account_path.clone(), + Nibbles::unpack(storage_key3), + ); + let storage_value3 = U256::from(345); + overlay_mut.insert(storage_path3.full_path(), Some(OverlayValue::Storage(storage_value3))); + + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_overlay_update_account_with_storage() { + let tmp_dir = TempDir::new("test_overlay_update_account_with_storage").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + let account_address = Address::random(); + let account_path = AddressPath::for_address(account_address); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + let storage_key1 = U256::from(42); + let storage_key2 = U256::from(43); + let storage_path1 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key1.into()); + let storage_path2 = + crate::path::StoragePath::for_address_and_slot(account_address, storage_key2.into()); + + db.storage_engine + .set_values( + &mut context, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), + ], + ) + .unwrap(); + + let mut overlay_mut = OverlayStateMut::new(); + let new_account = Account::new(1, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY); + overlay_mut.insert(account_path.clone().into(), Some(OverlayValue::Account(new_account))); + overlay_mut.insert(storage_path1.full_path(), Some(OverlayValue::Storage(U256::from(333)))); + + let overlay = overlay_mut.freeze(); + + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_branch_node_prefix_ordering_bug() { + let tmp_dir = TempDir::new("test_branch_prefix_ordering").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + // Create the specific account path that causes the issue + // Account path: 0x1dfdadc9cfa121d06649309ad0233f7c14e54b7df756a84e7340eaf8b9912261 + let account_nibbles = Nibbles::from_nibbles([ + 0x1, 0xd, 0xf, 0xd, 0xa, 0xd, 0xc, 0x9, 0xc, 0xf, 0xa, 0x1, 0x2, 0x1, 0xd, 0x0, 0x6, + 0x6, 0x4, 0x9, 0x3, 0x0, 0x9, 0xa, 0xd, 0x0, 0x2, 0x3, 0x3, 0xf, 0x7, 0xc, 0x1, 0x4, + 0xe, 0x5, 0x4, 0xb, 0x7, 0xd, 0xf, 0x7, 0x5, 0x6, 0xa, 0x8, 0x4, 0xe, 0x7, 0x3, 0x4, + 0x0, 0xe, 0xa, 0xf, 0x8, 0xb, 0x9, 0x9, 0x1, 0x2, 0x2, 0x6, 0x1, + ]); + let account_path = AddressPath::new(account_nibbles); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + // Create storage paths that will create a branch node with prefix 0x340 + // These paths are designed to create a branch at storage path 0x340 with children at: + // - 0x340123aa...aa + // - 0x340123bb...bb + // - 0x3411...11 + + // First storage path: 0x340123aa...aa (remaining 60 nibbles are 'a') + let mut storage1_nibbles = vec![0x3, 0x4, 0x0, 0x0, 0x0, 0x0]; // 6 nibbles + storage1_nibbles.extend(vec![0xa; 58]); // Fill remaining 58 nibbles with 'a' to make 64 total + let storage1_path = crate::path::StoragePath::for_address_path_and_slot_hash( + account_path.clone(), + Nibbles::from_nibbles(storage1_nibbles), + ); + + // Second storage path: 0x340123bb...bb (remaining 60 nibbles are 'b') + let mut storage2_nibbles = vec![0x3, 0x4, 0x0, 0x0, 0x0, 0x0]; // 6 nibbles + storage2_nibbles.extend(vec![0xb; 58]); // Fill remaining 58 nibbles with 'b' to make 64 total + let storage2_path = crate::path::StoragePath::for_address_path_and_slot_hash( + account_path.clone(), + Nibbles::from_nibbles(storage2_nibbles), + ); + + // Third storage path: 0x3411...11 (remaining 62 nibbles are '1') + let mut storage3_nibbles = vec![0x3, 0x4, 0x1]; // 3 nibbles + storage3_nibbles.extend(vec![0x1; 61]); // Fill remaining 61 nibbles with '1' to make 64 total + let storage3_path = crate::path::StoragePath::for_address_path_and_slot_hash( + account_path.clone(), + Nibbles::from_nibbles(storage3_nibbles), + ); + + // Set up initial state with the account and storage that creates the branch structure + db.storage_engine + .set_values( + &mut context, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222)))), + (storage3_path.full_path(), Some(TrieValue::Storage(U256::from(333)))), + ], + ) + .unwrap(); + + // Now create an overlay that adds a storage value that will cause the ordering issue + // This path should be: account_path + + // 0x34044c6a65488ba0051b7669dae97b8b1fe0cdbb72351b0ca7b4dc42f50dd9b8 + let overlay_storage_nibbles = Nibbles::from_nibbles([ + 0x3, 0x4, 0x0, 0x4, 0x4, 0xc, 0x6, 0xa, 0x6, 0x5, 0x4, 0x8, 0x8, 0xb, 0xa, 0x0, 0x0, + 0x5, 0x1, 0xb, 0x7, 0x6, 0x6, 0x9, 0xd, 0xa, 0xe, 0x9, 0x7, 0xb, 0x8, 0xb, 0x1, 0xf, + 0xe, 0x0, 0xc, 0xd, 0xb, 0xb, 0x7, 0x2, 0x3, 0x5, 0x1, 0xb, 0x0, 0xc, 0xa, 0x7, 0xb, + 0x4, 0xd, 0xc, 0x4, 0x2, 0xf, 0x5, 0x0, 0xd, 0xd, 0x9, 0xb, 0x8, + ]); + let overlay_storage_path = crate::path::StoragePath::for_address_path_and_slot_hash( + account_path.clone(), + overlay_storage_nibbles, + ); + + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut + .insert(overlay_storage_path.full_path(), Some(OverlayValue::Storage(U256::from(999)))); + let overlay = overlay_mut.freeze(); + + // This triggered a panic due to lexicographic ordering violation + // The branch node at path ending in 0x340 will be added after its descendant + // at path ending in 0x34044c6a65488ba0051b7669dae97b8b1fe0cdbb72351b0ca7b4dc42f50dd9b8 + compare_overlay_with_committed_root(&db, &mut context, &overlay); + } + + #[test] + fn test_overlay_root_with_branch_node_prefix_ordering_bug() { + let tmp_dir = + TempDir::new("test_overlay_root_with_branch_node_prefix_ordering_bug").unwrap(); + let file_path = tmp_dir.path().join("test.db"); + let db = Database::create_new(file_path).unwrap(); + + let mut context = db.storage_engine.write_context(); + + let account_path = AddressPath::new(Nibbles::from_nibbles([ + 0x4, 0x5, 0x7, 0x0, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, + 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, + 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, + 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, 0x1, + ])); + let account = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); + + db.storage_engine + .set_values( + &mut context, + &mut [(account_path.clone().into(), Some(TrieValue::Account(account.clone())))], + ) + .unwrap(); + + let account_path2 = AddressPath::new(Nibbles::from_nibbles([ + 0x4, 0x5, 0x7, 0x0, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, + 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, + 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, + 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, 0x2, + ])); + + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut.insert( + Nibbles::from_nibbles([0x4, 0x5, 0x7, 0x0]), + Some(OverlayValue::Hash(B256::random())), + ); + overlay_mut.insert( + account_path2.clone().into(), + Some(OverlayValue::Account(Account::new( + 2, + U256::from(200), + EMPTY_ROOT_HASH, + KECCAK_EMPTY, + ))), + ); + let overlay = overlay_mut.freeze(); + + let initial_root = context.root_node_hash; + // This triggered a panic due to the addition of a leaf node after adding an ancestor branch + // node. + let output = db.storage_engine.compute_state_root_with_overlay(&context, overlay).unwrap(); + let (overlay_root, _, _) = + (output.root, output.updated_branch_nodes, output.storage_branch_updates); + assert_ne!(overlay_root, initial_root, "Overlay should not match initial root"); + } +} diff --git a/src/transaction.rs b/src/transaction.rs index 5005bb8..f7600a4 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -6,15 +6,16 @@ use crate::{ context::TransactionContext, database::Database, node::TrieValue, + overlay::OverlayState, path::{AddressPath, StoragePath}, - storage::proofs::AccountProof, + storage::{overlay_root::OverlayedRoot, proofs::AccountProof}, }; -use alloy_primitives::{StorageValue, B256}; +use alloy_primitives::{map::HashMap, StorageValue, B256}; use alloy_trie::Nibbles; pub use error::TransactionError; pub use manager::TransactionManager; use sealed::sealed; -use std::{collections::HashMap, fmt::Debug, ops::Deref, sync::Arc}; +use std::{fmt::Debug, ops::Deref, sync::Arc}; #[sealed] pub trait TransactionKind: Debug {} @@ -55,7 +56,7 @@ impl, K: TransactionKind> Transaction { committed: false, context, database, - pending_changes: HashMap::new(), + pending_changes: HashMap::default(), _marker: std::marker::PhantomData, } } @@ -84,6 +85,16 @@ impl, K: TransactionKind> Transaction { self.context.root_node_hash } + pub fn compute_root_with_overlay( + &self, + overlay_state: OverlayState, + ) -> Result { + self.database + .storage_engine + .compute_state_root_with_overlay(&self.context, overlay_state) + .map_err(|_| TransactionError) + } + pub fn get_account_with_proof( &self, address_path: AddressPath,