From 5ac409e5dda0e18862f78682c6fc8ad383739838 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Thu, 24 Jul 2025 18:58:46 -0700 Subject: [PATCH 01/17] WIP: Support computing state root with overlayed state --- src/lib.rs | 1 + src/overlay.rs | 734 ++++++++++++++++ src/storage.rs | 1 + src/storage/overlay_root.rs | 1658 +++++++++++++++++++++++++++++++++++ src/transaction.rs | 3 + src/transaction/error.rs | 12 +- src/transaction/manager.rs | 2 +- 7 files changed, 2408 insertions(+), 3 deletions(-) create mode 100644 src/overlay.rs create mode 100644 src/storage/overlay_root.rs diff --git a/src/lib.rs b/src/lib.rs index 022db229..98e9c5c1 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/overlay.rs b/src/overlay.rs new file mode 100644 index 00000000..2a12698f --- /dev/null +++ b/src/overlay.rs @@ -0,0 +1,734 @@ +use crate::{ + node::{Node, TrieValue}, + pointer::Pointer, +}; +use alloy_trie::Nibbles; + +/// 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. + 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 { changes: Vec::new().into_boxed_slice(), 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(); + } + } + } + + OverlayState { changes: deduped.into_boxed_slice(), prefix_offset: 0 } + } +} + +/// Immutable overlay state with sorted changes for efficient querying and sub-slicing. +/// This is created by freezing an OverlayStateMut and allows for efficient prefix operations. +#[derive(Debug, Clone)] +pub struct OverlayState { + changes: Box<[(Nibbles, Option)]>, + prefix_offset: usize, +} + +impl OverlayState { + /// Creates an empty overlay state. + pub fn empty() -> Self { + Self { changes: Vec::new().into_boxed_slice(), prefix_offset: 0 } + } + + /// 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() + } + + /// 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: &Nibbles) -> Option<&Option> { + match self.changes.binary_search_by(|(p, _)| self.compare_with_offset(p, path)) { + Ok(index) => Some(&self.changes[index].1), + Err(_) => None, + } + } + + /// Finds all entries that have the given prefix. + /// Returns a sub-slice of the overlay containing only matching entries. + /// The prefix is compared against offset-adjusted paths. + pub fn find_prefix_range(&self, prefix: &Nibbles) -> OverlayState { + if self.changes.is_empty() { + return OverlayState::empty(); + } + + let mut start_idx = None; + let mut end_idx = self.changes.len(); + + // Find the range of entries that start with the prefix after applying offset + for (i, (path, _)) in self.changes.iter().enumerate() { + if path.len() >= self.prefix_offset { + let adjusted_path = path.slice(self.prefix_offset..); + if adjusted_path.len() >= prefix.len() && adjusted_path[..prefix.len()] == *prefix { + if start_idx.is_none() { + start_idx = Some(i); + } + } else if start_idx.is_some() { + // We've found the end of the matching range + end_idx = i; + break; + } + } + } + + match start_idx { + Some(start) => OverlayState { + changes: self.changes[start..end_idx].to_vec().into_boxed_slice(), + prefix_offset: self.prefix_offset, + }, + None => OverlayState::empty(), + } + } + + /// 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.changes.len() > 1 { + let first_path = &self.changes[0].0; + let last_path = &self.changes[self.changes.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 { changes: self.changes.clone(), prefix_offset: total_offset } + } + + /// 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: &Nibbles, + target_path: &Nibbles, + ) -> std::cmp::Ordering { + if stored_path.len() < self.prefix_offset { + std::cmp::Ordering::Less + } else { + let adjusted_path = stored_path.slice(self.prefix_offset..); + adjusted_path.cmp(target_path) + } + } + + /// Creates a 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. + pub fn sub_slice(&self, start: usize, end: usize) -> OverlayState { + let end = end.min(self.changes.len()); + let start = start.min(end); + + OverlayState { + changes: self.changes[start..end].to_vec().into_boxed_slice(), + prefix_offset: self.prefix_offset, + } + } + + /// Creates a sub-slice containing only changes that come before the given prefix. + /// The prefix comparison takes prefix_offset into account. + pub fn sub_slice_before_prefix(&self, prefix: &Nibbles) -> OverlayState { + let index = self + .changes + .binary_search_by(|(p, _)| self.compare_with_offset(p, prefix)) + .unwrap_or_else(|i| i); // Insert position is exactly what we want for "before" + self.sub_slice(0, index) + } + + /// Creates a sub-slice containing only changes that come strictly after the given prefix. + /// The prefix comparison takes prefix_offset into account. + pub fn sub_slice_after_prefix(&self, prefix: &Nibbles) -> OverlayState { + // Find the next key after prefix by incrementing prefix + let next_prefix = match prefix.increment() { + Some(next) => next, + None => { + // Prefix is all 0xF's, nothing can come after it + return OverlayState::empty(); + } + }; + + let index = self + .changes + .binary_search_by(|(p, _)| self.compare_with_offset(p, &next_prefix)) + .unwrap_or_else(|i| i); + self.sub_slice(index, self.changes.len()) + } + + /// Creates a 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: &Nibbles) -> OverlayState { + self.find_prefix_range(prefix) + } + + /// Returns an iterator over all changes in the overlay. + /// The paths are adjusted by the prefix_offset. + pub fn iter(&self) -> impl Iterator)> { + self.changes.iter().filter_map(move |(path, value)| { + if path.len() >= self.prefix_offset { + let adjusted_path = path.slice(self.prefix_offset..); + Some((adjusted_path, value)) + } else { + None + } + }) + } + + /// Checks if the overlay contains any changes that would affect the given path. + /// This includes exact matches and any changes to descendants of the path. + /// The path comparison takes prefix_offset into account. + pub fn affects_path(&self, path: &Nibbles) -> bool { + if self.is_empty() { + return false; + } + + // Check for exact match or any path that starts with the given path after applying offset + self.iter().any(|(overlay_path, _)| overlay_path.starts_with(path)) + } +} + +/// Pointer that can reference either on-disk nodes or in-memory overlay nodes. +/// This allows the trie traversal to seamlessly handle both persisted and overlayed state. +#[derive(Debug, Clone)] +pub enum OverlayPointer { + /// Reference to a node stored on disk + OnDisk(Pointer), + /// Reference to a node stored in memory as part of the overlay + InMemory(Box), +} + +impl OverlayPointer { + /// Creates a new on-disk overlay pointer. + pub fn on_disk(pointer: Pointer) -> Self { + Self::OnDisk(pointer) + } + + /// Creates a new in-memory overlay pointer. + pub fn in_memory(node: Node) -> Self { + Self::InMemory(Box::new(node)) + } + + /// Returns true if this pointer references an on-disk node. + pub fn is_on_disk(&self) -> bool { + matches!(self, Self::OnDisk(_)) + } + + /// Returns true if this pointer references an in-memory node. + pub fn is_in_memory(&self) -> bool { + matches!(self, Self::InMemory(_)) + } + + /// Returns the underlying disk pointer if this is an on-disk reference. + pub fn as_disk_pointer(&self) -> Option<&Pointer> { + match self { + Self::OnDisk(pointer) => Some(pointer), + Self::InMemory(_) => None, + } + } + + /// Returns the underlying in-memory node if this is an in-memory reference. + pub fn as_memory_node(&self) -> Option<&Node> { + match self { + Self::OnDisk(_) => None, + Self::InMemory(node) => Some(node), + } + } + + /// Converts this overlay pointer to an owned Node. + /// For on-disk pointers, this would require loading the node from storage. + /// For in-memory pointers, this clones the existing node. + pub fn into_node(self) -> Result { + match self { + Self::OnDisk(_) => Err(OverlayError::DiskNodeNotLoaded), + Self::InMemory(node) => Ok(*node), + } + } +} + +impl From for OverlayPointer { + fn from(pointer: Pointer) -> Self { + Self::OnDisk(pointer) + } +} + +impl From for OverlayPointer { + fn from(node: Node) -> Self { + Self::InMemory(Box::new(node)) + } +} + +/// Errors that can occur when working with overlay pointers. +#[derive(Debug)] +pub enum OverlayError { + /// Attempted to access a disk node without loading it first + DiskNodeNotLoaded, +} + +impl std::fmt::Display for OverlayError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + Self::DiskNodeNotLoaded => write!(f, "disk node not loaded"), + } + } +} + +impl std::error::Error for OverlayError {} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{account::Account, node::TrieValue}; + 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(TrieValue::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(TrieValue::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(TrieValue::Account(account1))); + mutable.insert(path.clone(), Some(TrieValue::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(TrieValue::Account(account.clone()))); + mutable + .insert(Nibbles::from_nibbles([1, 2, 5, 6]), Some(TrieValue::Account(account.clone()))); + mutable + .insert(Nibbles::from_nibbles([1, 3, 7, 8]), Some(TrieValue::Account(account.clone()))); + mutable + .insert(Nibbles::from_nibbles([2, 3, 4, 5]), Some(TrieValue::Account(account.clone()))); + + let frozen = mutable.freeze(); + + // Find entries with prefix [1, 2] + let prefix = Nibbles::from_nibbles([1, 2]); + let subset = frozen.find_prefix_range(&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.find_prefix_range(&prefix); + + assert_eq!(subset.len(), 3); // Should match first three entries + } + + #[test] + fn test_affects_path() { + let mut mutable = OverlayStateMut::new(); + let account = test_account(); + + mutable + .insert(Nibbles::from_nibbles([1, 2, 3, 4]), Some(TrieValue::Account(account.clone()))); + mutable + .insert(Nibbles::from_nibbles([1, 2, 5, 6]), Some(TrieValue::Account(account.clone()))); + + let frozen = mutable.freeze(); + + // Test exact match + assert!(frozen.affects_path(&Nibbles::from_nibbles([1, 2, 3, 4]))); + + // Test parent path + assert!(frozen.affects_path(&Nibbles::from_nibbles([1, 2]))); + + // Test unrelated path + assert!(!frozen.affects_path(&Nibbles::from_nibbles([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.affects_path(&path)); + + let subset = frozen.find_prefix_range(&path); + assert!(subset.is_empty()); + } + + #[test] + fn test_overlay_pointer_creation() { + use crate::{location::Location, node::Node, pointer::Pointer}; + use alloy_trie::{nodes::RlpNode, Nibbles, EMPTY_ROOT_HASH}; + + // Test on-disk pointer + let location = Location::for_cell(42); + let rlp = RlpNode::word_rlp(&EMPTY_ROOT_HASH); + let disk_pointer = Pointer::new(location, rlp); + let overlay_pointer = OverlayPointer::on_disk(disk_pointer.clone()); + + assert!(overlay_pointer.is_on_disk()); + assert!(!overlay_pointer.is_in_memory()); + assert_eq!(overlay_pointer.as_disk_pointer(), Some(&disk_pointer)); + assert!(overlay_pointer.as_memory_node().is_none()); + + // Test in-memory pointer + let path = Nibbles::from_nibbles([1, 2, 3, 4]); + let account = test_account(); + let node = Node::new_leaf(path, &TrieValue::Account(account)).unwrap(); + let overlay_pointer = OverlayPointer::in_memory(node.clone()); + + assert!(!overlay_pointer.is_on_disk()); + assert!(overlay_pointer.is_in_memory()); + assert!(overlay_pointer.as_disk_pointer().is_none()); + assert_eq!(overlay_pointer.as_memory_node(), Some(&node)); + } + + #[test] + fn test_overlay_pointer_conversion() { + use crate::{location::Location, pointer::Pointer}; + use alloy_trie::{nodes::RlpNode, Nibbles, EMPTY_ROOT_HASH}; + + // Test From + let location = Location::for_cell(42); + let rlp = RlpNode::word_rlp(&EMPTY_ROOT_HASH); + let disk_pointer = Pointer::new(location, rlp); + let overlay_pointer: OverlayPointer = disk_pointer.into(); + assert!(overlay_pointer.is_on_disk()); + + // Test From + let path = Nibbles::from_nibbles([1, 2, 3, 4]); + let account = test_account(); + let node = Node::new_leaf(path, &TrieValue::Account(account)).unwrap(); + let overlay_pointer: OverlayPointer = node.into(); + assert!(overlay_pointer.is_in_memory()); + } + + #[test] + fn test_overlay_pointer_into_node() { + use alloy_trie::Nibbles; + + // Test in-memory pointer + let path = Nibbles::from_nibbles([1, 2, 3, 4]); + let account = test_account(); + let original_node = Node::new_leaf(path, &TrieValue::Account(account)).unwrap(); + let overlay_pointer = OverlayPointer::in_memory(original_node.clone()); + + let extracted_node = overlay_pointer.into_node().unwrap(); + assert_eq!(extracted_node.prefix(), original_node.prefix()); + assert_eq!(extracted_node.value().unwrap(), original_node.value().unwrap()); + + // Test on-disk pointer (should fail) + use crate::{location::Location, pointer::Pointer}; + use alloy_trie::{nodes::RlpNode, EMPTY_ROOT_HASH}; + + let location = Location::for_cell(42); + let rlp = RlpNode::word_rlp(&EMPTY_ROOT_HASH); + let disk_pointer = Pointer::new(location, rlp); + let overlay_pointer = OverlayPointer::on_disk(disk_pointer); + + assert!(overlay_pointer.into_node().is_err()); + } + + #[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(TrieValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles(storage_path2.clone()), + Some(TrieValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles(storage_path3.clone()), + Some(TrieValue::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]); + } + } + + #[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 = vec![1, 0, 0, 0]; // 4-nibble account prefix + + let paths = [ + vec![1, 0, 0, 0, 2, 0], // account + storage [2, 0] + vec![1, 0, 0, 0, 5, 0], // account + storage [5, 0] + vec![1, 0, 0, 0, 8, 0], // account + storage [8, 0] + vec![1, 0, 0, 0, 9, 0], /* account + storage [9, 0] - changed from different account + * to same account */ + ]; + + for path in &paths { + mutable.insert( + Nibbles::from_nibbles(path.clone()), + Some(TrieValue::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 = Nibbles::from_nibbles([5, 0]); // Storage key [5, 0] + let before = storage_overlay.sub_slice_before_prefix(&split_point); + let after = storage_overlay.sub_slice_after_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(&Nibbles::from_nibbles([2, 0]))); + + // 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(&Nibbles::from_nibbles([5, 0]))); // Strictly after, so [5, 0] not included + assert!(after_paths.contains(&Nibbles::from_nibbles([8, 0]))); + assert!(after_paths.contains(&Nibbles::from_nibbles([9, 0]))); + } + + #[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 = vec![1, 0, 0, 0]; // Account 1: 4 nibbles + let _account2_prefix = vec![2, 0, 0, 0]; // Account 2: 4 nibbles + + // Storage for account 1 + mutable.insert( + Nibbles::from_nibbles([1, 0, 0, 0, 5, 5]), + Some(TrieValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles([1, 0, 0, 0, 5, 6]), + Some(TrieValue::Account(account.clone())), + ); + + // Storage for account 2 + mutable.insert( + Nibbles::from_nibbles([2, 0, 0, 0, 7, 7]), + Some(TrieValue::Account(account.clone())), + ); + + let frozen = mutable.freeze(); + + // Test finding storage for account 1 using find_prefix_range on original overlay + let account1_storage = + frozen.find_prefix_range(&Nibbles::from_nibbles(account1_prefix.clone())); + 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: Vec<_> = storage_overlay + .find_prefix_range(&Nibbles::from_nibbles([5])) + .iter() + .map(|(path, _)| path) + .collect(); + + assert_eq!(storage_with_prefix5.len(), 2); + assert!(storage_with_prefix5.contains(&Nibbles::from_nibbles([5, 5]))); + assert!(storage_with_prefix5.contains(&Nibbles::from_nibbles([5, 6]))); + } +} diff --git a/src/storage.rs b/src/storage.rs index b86be74f..2638784e 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 00000000..779f7de8 --- /dev/null +++ b/src/storage/overlay_root.rs @@ -0,0 +1,1658 @@ +use crate::{ + context::TransactionContext, + node::{Node, NodeKind, TrieValue}, + overlay::OverlayState, + page::{PageId, SlottedPage}, + storage::engine::{Error, StorageEngine}, +}; +use alloy_primitives::B256; +use alloy_rlp::encode; +use alloy_trie::{HashBuilder, Nibbles, TrieAccount}; + +impl StorageEngine { + /// Computes the root hash with overlay changes without persisting them. + /// This uses alloy-trie's HashBuilder for efficient merkle root computation. + pub fn compute_root_with_overlay( + &self, + context: &TransactionContext, + overlay: &OverlayState, + ) -> Result { + if overlay.is_empty() { + // No overlay changes, return current root + return Ok(context.root_node_hash); + } + + let mut hash_builder = HashBuilder::default(); + + // Use proper trie traversal with overlay integration + self.traverse_with_overlay(context, overlay, &mut hash_builder)?; + + let root = hash_builder.root(); + Ok(root) + } + + /// Helper method to traverse the existing trie and integrate with overlay + fn traverse_with_overlay( + &self, + context: &TransactionContext, + overlay: &OverlayState, + hash_builder: &mut HashBuilder, + ) -> Result<(), Error> { + // First, collect all existing values from disk + if let Some(root_page_id) = context.root_node_page_id { + self.traverse_page_with_overlay( + context, + overlay, + hash_builder, + root_page_id, + &Nibbles::new(), + Some(context.root_node_hash), + )?; + } else { + // No root page, just add all overlay changes to the hash builder + for (path, value) in overlay.iter() { + if let Some(trie_value) = value { + let rlp_encoded = self.encode_trie_value(trie_value)?; + hash_builder.add_leaf(path, &rlp_encoded); + } + } + } + + Ok(()) + } + + /// Traverse a specific page with overlay integration + fn traverse_page_with_overlay( + &self, + context: &TransactionContext, + overlay: &OverlayState, + hash_builder: &mut HashBuilder, + page_id: PageId, + path_prefix: &Nibbles, + node_hash: Option, + ) -> Result<(), Error> { + let page = self.get_page(context, page_id)?; + let slotted_page = SlottedPage::try_from(page)?; + + // Start at the root cell (cell_index = 0) and recurse properly + self.traverse_node_with_overlay( + context, + overlay, + hash_builder, + &slotted_page, + 0, // Start from root cell + path_prefix, + node_hash, + ) + } + + /// Traverse a specific node with overlay integration + fn traverse_node_with_overlay( + &self, + context: &TransactionContext, + overlay: &OverlayState, + hash_builder: &mut HashBuilder, + slotted_page: &SlottedPage, + cell_index: u8, + path_prefix: &Nibbles, + node_hash: Option, + ) -> Result<(), Error> { + let node: Node = slotted_page.get_value(cell_index)?; + let full_path = { + let mut full = path_prefix.clone(); + full.extend_from_slice_unchecked(node.prefix()); + full + }; + + let pre_overlay = overlay.sub_slice_before_prefix(&full_path); + let post_overlay = overlay.sub_slice_after_prefix(&full_path); + + // Add all pre-overlay changes + for (path, value) in pre_overlay.iter() { + if let Some(trie_value) = value { + let rlp_encoded = self.encode_trie_value(trie_value)?; + hash_builder.add_leaf(path, &rlp_encoded); + } + } + + // Check if this node is affected by overlay + if overlay.affects_path(&full_path) { + // This node or its descendants are affected by overlay + self.handle_affected_node_with_overlay( + context, + overlay, + hash_builder, + slotted_page, + &node, + &full_path, + )?; + } else { + // Node not affected by overlay, use disk value as-is + self.handle_unaffected_node( + context, + hash_builder, + slotted_page, + &node, + &full_path, + node_hash, + )?; + } + + for (path, value) in post_overlay.iter() { + if let Some(trie_value) = value { + let rlp_encoded = self.encode_trie_value(trie_value)?; + hash_builder.add_leaf(path, &rlp_encoded); + } + } + + Ok(()) + } + + /// Handle a node that is affected by overlay changes + fn handle_affected_node_with_overlay( + &self, + context: &TransactionContext, + overlay: &OverlayState, + hash_builder: &mut HashBuilder, + slotted_page: &SlottedPage, + node: &Node, + full_path: &Nibbles, + ) -> Result<(), Error> { + use crate::node::NodeKind; + + match node.kind() { + NodeKind::AccountLeaf { .. } if full_path.len() == 64 => { + // Account leaf at address level - handle potential storage overlays + self.handle_account_leaf_with_overlay( + context, + overlay, + hash_builder, + slotted_page, + node, + full_path, + )?; + } + NodeKind::StorageLeaf { .. } => { + // Storage leaf - handle with simple leaf logic + // This can be either 128-nibble paths (account-level traversal) or 64-nibble paths + // (storage-level traversal) + if let Some(overlay_value) = overlay.lookup(full_path) { + // Use overlay value + if let Some(trie_value) = overlay_value { + let rlp_encoded = self.encode_trie_value(trie_value)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } + // If overlay_value is None, it's a deletion - skip + } else { + // No exact overlay match, use disk value + let node_value = node.value()?; + let rlp_encoded = self.encode_trie_value(&node_value)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } + } + NodeKind::Branch { .. } => { + // Branch node - recurse into each child with overlay sub-slicing + self.traverse_branch_node_with_overlay( + context, + overlay, + hash_builder, + slotted_page, + node, + full_path, + )?; + } + _ => { + // Other leaf nodes - handle with simple leaf logic + if let Some(overlay_value) = overlay.lookup(full_path) { + if let Some(trie_value) = overlay_value { + let rlp_encoded = self.encode_trie_value(trie_value)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } + } else { + let node_value = node.value()?; + let rlp_encoded = self.encode_trie_value(&node_value)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } + } + } + Ok(()) + } + + /// Handle a node that is not affected by overlay changes + fn handle_unaffected_node( + &self, + _context: &TransactionContext, + hash_builder: &mut HashBuilder, + _slotted_page: &SlottedPage, + node: &Node, + full_path: &Nibbles, + node_hash: Option, + ) -> Result<(), Error> { + match node.kind() { + NodeKind::AccountLeaf { .. } | NodeKind::StorageLeaf { .. } => { + // Account leaf - use disk value directly + let node_value = node.value()?; + let rlp_encoded = self.encode_trie_value(&node_value)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } + NodeKind::Branch { .. } => { + // Branch node - since it's unaffected by overlay, we can use its existing hash + // This is much more efficient than recursing into all children + if let Some(hash) = node_hash { + // Use the precomputed hash directly for this branch subtree + // stored_in_database = true since this is an existing node from disk + hash_builder.add_branch(full_path.clone(), hash, true); + } else { + // Fallback: if we don't have the hash, we need to compute it + // This shouldn't normally happen but provides a safety net + let node_value = node.value()?; + let rlp_encoded = self.encode_trie_value(&node_value)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } + } + } + Ok(()) + } + + /// Handle traversal of branch nodes with overlay integration + fn traverse_branch_node_with_overlay( + &self, + context: &TransactionContext, + overlay: &OverlayState, + hash_builder: &mut HashBuilder, + slotted_page: &SlottedPage, + node: &Node, + full_path: &Nibbles, + ) -> Result<(), Error> { + // For each child in the branch node (0-15), check if there are relevant overlay changes + // and recursively traverse child nodes + for child_index in 0..16 { + // Construct the path for this child + let mut child_path = full_path.clone(); + child_path.push(child_index); + + // Create sub-slice of overlay for this child's subtree + let child_overlay = overlay.sub_slice_for_prefix(&child_path); + + if let Ok(Some(child_pointer)) = node.child(child_index) { + // Child exists on disk - traverse it with overlay integration + let child_hash = child_pointer.rlp().as_hash(); + let child_location = child_pointer.location(); + + if let Some(child_cell_index) = child_location.cell_index() { + // Child is in the same page + self.traverse_node_with_overlay( + context, + &child_overlay, + hash_builder, + slotted_page, + child_cell_index, + &child_path, + child_hash, + )?; + } else if let Some(child_page_id) = child_location.page_id() { + // Child is in a different page - add safety check + if self.page_exists(context, child_page_id) { + self.traverse_page_with_overlay( + context, + &child_overlay, + hash_builder, + child_page_id, + &child_path, + child_hash, + )?; + } + } + } else { + // No child exists on disk, but check if overlay has leaves for this subtree + if !child_overlay.is_empty() { + // There are overlay changes that create new nodes in this subtree + // Add all overlay leaves for this child subtree + for (path, value) in child_overlay.iter() { + if let Some(trie_value) = value { + let rlp_encoded = self.encode_trie_value(trie_value)?; + hash_builder.add_leaf(path, &rlp_encoded); + } + // Tombstones (None values) are skipped + } + } + } + } + Ok(()) + } + + /// Handle an account leaf with potential storage overlays using iterative HashBuilder approach + fn handle_account_leaf_with_overlay( + &self, + context: &TransactionContext, + overlay: &OverlayState, + hash_builder: &mut HashBuilder, + slotted_page: &SlottedPage, + node: &Node, + full_path: &Nibbles, + ) -> Result<(), Error> { + // Get the original account from the node + let original_account = match node.value()? { + TrieValue::Account(account) => account, + _ => { + // This shouldn't happen for AccountLeaf nodes, but handle gracefully + let rlp_encoded = self.encode_trie_value(&node.value()?)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + return Ok(()); + } + }; + + // Check if this account is directly overlaid (account-level changes) + if let Some(overlay_value) = overlay.lookup(full_path) { + // Direct account overlay - use overlay value, but still check for storage overlays + if let Some(trie_value) = overlay_value { + // Even with account overlay, we might need to compute storage root for storage + // overlays + if let TrieValue::Account(overlay_account) = trie_value { + // Check if there are storage overlays for this account + let storage_overlays = overlay.find_prefix_range(full_path); + let has_storage_overlays = storage_overlays + .iter() + .any(|(path, _)| path.len() == 128 && path.len() > full_path.len()); + + if has_storage_overlays { + // Compute new storage root with overlays + let new_storage_root = self.compute_storage_root_iteratively( + context, + node, // Use the original node for storage traversal + full_path, + &storage_overlays, + slotted_page, + )?; + + // Create modified account with new storage root + let mut modified_account = overlay_account.clone(); + modified_account.storage_root = new_storage_root; + let rlp_encoded = + self.encode_trie_value(&TrieValue::Account(modified_account))?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } else { + // No storage overlays, use account overlay as-is + let rlp_encoded = self.encode_trie_value(trie_value)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } + } else { + // Not an account value, just use as-is + let rlp_encoded = self.encode_trie_value(trie_value)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } + } + // If overlay_value is None, it's a deletion - skip + return Ok(()); + } + + // Check if there are any storage overlays for this account + // Storage overlays have 128-nibble paths that start with this account's 64-nibble path + let storage_overlays = overlay.find_prefix_range(full_path); + let has_storage_overlays = storage_overlays + .iter() + .any(|(path, _)| path.len() == 128 && path.len() > full_path.len()); + + if !has_storage_overlays { + // No storage overlays for this account, use original account + let rlp_encoded = self.encode_trie_value(&TrieValue::Account(original_account))?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + return Ok(()); + } + + // We have storage overlays, need to compute a new storage root using iterative approach + let new_storage_root = self.compute_storage_root_iteratively( + context, + node, + full_path, + &storage_overlays, + slotted_page, + )?; + + // Create a modified account with the new storage root + let mut modified_account = original_account; + modified_account.storage_root = new_storage_root; + + // Add the modified account to the main HashBuilder + let rlp_encoded = self.encode_trie_value(&TrieValue::Account(modified_account))?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + + Ok(()) + } + + /// Compute the storage root for an account using iterative HashBuilder approach with prefix + /// offset + fn compute_storage_root_iteratively( + &self, + context: &TransactionContext, + account_node: &Node, + _account_path: &Nibbles, + storage_overlays: &OverlayState, + slotted_page: &SlottedPage, + ) -> Result { + // Create a storage-specific overlay with 64-nibble prefix offset + // This converts 128-nibble storage paths to 64-nibble storage-relative paths + let storage_overlay = storage_overlays.with_prefix_offset(64); + + let mut storage_hash_builder = HashBuilder::default(); + + // Get the original account's storage root to determine if we need to traverse existing + // storage + let original_storage_root = match account_node.value()? { + TrieValue::Account(account) => account.storage_root, + _ => return Err(Error::NodeError(crate::node::NodeError::NoValue)), /* This shouldn't happen for account nodes */ + }; + + // If the account has existing storage, traverse the existing storage trie + if original_storage_root != alloy_trie::EMPTY_ROOT_HASH { + // Try to find the storage root page from the account node's direct child + if let Ok(Some(storage_root_pointer)) = account_node.direct_child() { + let storage_root_location = storage_root_pointer.location(); + let storage_root_hash = storage_root_pointer.rlp().as_hash(); + + if let Some(storage_page_id) = storage_root_location.page_id() { + // Traverse the existing storage trie with storage overlays + self.traverse_page_with_overlay( + context, + &storage_overlay, + &mut storage_hash_builder, + storage_page_id, + &Nibbles::new(), // Start with empty path for storage root + storage_root_hash, + )?; + } else if let Some(storage_cell_index) = storage_root_location.cell_index() { + // Storage root is in the same page as the account + self.traverse_node_with_overlay( + context, + &storage_overlay, + &mut storage_hash_builder, + slotted_page, + storage_cell_index, + &Nibbles::new(), // Start with empty path for storage root + storage_root_hash, + )?; + } + } + } else { + // No existing storage, just add overlay changes + for (path, value) in storage_overlay.iter() { + if let Some(trie_value) = value { + let rlp_encoded = self.encode_trie_value(trie_value)?; + storage_hash_builder.add_leaf(path, &rlp_encoded); + } + } + } + + let computed_root = storage_hash_builder.root(); + Ok(computed_root) + } + + /// Check if a page exists in the storage + fn page_exists(&self, context: &TransactionContext, page_id: PageId) -> bool { + self.get_page(context, page_id).is_ok() + } + + /// Helper to encode TrieValue as RLP + fn encode_trie_value(&self, trie_value: &TrieValue) -> Result, Error> { + let rlp_encoded = match trie_value { + TrieValue::Account(account) => { + let trie_account = TrieAccount { + nonce: account.nonce, + balance: account.balance, + storage_root: account.storage_root, + code_hash: account.code_hash, + }; + encode(trie_account) + } + TrieValue::Storage(storage_value) => encode(storage_value), + }; + Ok(rlp_encoded) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::{ + account::Account, database::Database, node::TrieValue, overlay::OverlayStateMut, + path::AddressPath, + }; + use alloy_primitives::{address, U256}; + use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY}; + use tempdir::TempDir; + + #[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 root = db.storage_engine.compute_root_with_overlay(&context, &empty_overlay).unwrap(); + assert_eq!(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(TrieValue::Account(account))); + let overlay = overlay_mut.freeze(); + + // Compute root with overlay + let root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + + // The root should be different from the empty root (since we have changes) + assert_ne!(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); + + // Test that we can compute root with empty overlay first (should match initial_root) + let empty_overlay = OverlayStateMut::new().freeze(); + let root_with_empty_overlay = + db.storage_engine.compute_root_with_overlay(&context, &empty_overlay).unwrap(); + assert_eq!(root_with_empty_overlay, initial_root); + + // 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(TrieValue::Account(account2.clone()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify: commit the overlay changes and compare roots + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context, + &mut [(address_path.into(), Some(TrieValue::Account(account2)))], + ) + .unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!(overlay_root, committed_root, "Overlay root should match committed root"); + } + + #[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(TrieValue::Account(account1_updated.clone()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the change + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context, + &mut [ + (path1.clone().into(), Some(TrieValue::Account(account1_updated.clone()))), + (path2.clone().into(), Some(TrieValue::Account(account2.clone()))), + ], + ) + .unwrap(); + assert_eq!( + overlay_root, verification_context.root_node_hash, + "Case 1: Overlay root should match committed root" + ); + + // Test Case 2: Overlay that creates a new account in an empty subtree (None child case) + // 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(path3.clone().into(), Some(TrieValue::Account(account3.clone()))); + let overlay2 = overlay_mut2.freeze(); + + let overlay_root2 = + db.storage_engine.compute_root_with_overlay(&context, &overlay2).unwrap(); + assert_ne!(overlay_root2, initial_root); + assert_ne!(overlay_root2, overlay_root); + + // Verify by committing the change + let mut verification_context2 = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context2, + &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(); + assert_eq!( + overlay_root2, verification_context2.root_node_hash, + "Case 2: Overlay root should match committed root" + ); + + // Test Case 3: Mixed overlay - affects one subtree, creates new subtree, leaves one + // unaffected + let mut overlay_mut3 = OverlayStateMut::new(); + overlay_mut3 + .insert(path1.clone().into(), Some(TrieValue::Account(account1_updated.clone()))); // Modify existing + overlay_mut3.insert(path3.clone().into(), Some(TrieValue::Account(account3.clone()))); // Create new + // path2 is left unaffected - should use add_branch optimization + let overlay3 = overlay_mut3.freeze(); + + let overlay_root3 = + db.storage_engine.compute_root_with_overlay(&context, &overlay3).unwrap(); + assert_ne!(overlay_root3, initial_root); + assert_ne!(overlay_root3, overlay_root); + assert_ne!(overlay_root3, overlay_root2); + + // Verify by committing the changes + let mut verification_context3 = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context3, + &mut [ + (path1.clone().into(), Some(TrieValue::Account(account1_updated))), + (path2.clone().into(), Some(TrieValue::Account(account2))), + (path3.clone().into(), Some(TrieValue::Account(account3))), + ], + ) + .unwrap(); + assert_eq!( + overlay_root3, verification_context3.root_node_hash, + "Case 3: Overlay root should match committed root" + ); + } + + #[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 context1 = 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); + + db.storage_engine + .set_values( + &mut context1, + &mut [(path1.clone().into(), Some(TrieValue::Account(account1.clone())))], + ) + .unwrap(); + let root_with_account1 = context1.root_node_hash; + + // Step 2: Add account 2 + let mut context2 = db.storage_engine.write_context(); + 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); + + db.storage_engine + .set_values( + &mut context2, + &mut [ + (path1.clone().into(), Some(TrieValue::Account(account1.clone()))), + (path2.clone().into(), Some(TrieValue::Account(account2.clone()))), + ], + ) + .unwrap(); + let root_with_both_accounts = context2.root_node_hash; + assert_ne!(root_with_both_accounts, root_with_account1); + + // Step 3: 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_with_deletion = + db.storage_engine.compute_root_with_overlay(&context2, &overlay).unwrap(); + assert_ne!(overlay_root_with_deletion, root_with_both_accounts); + + // Step 4: Verify by actually erasing account 2 and computing root + let mut context3 = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut context3, + &mut [ + (path1.clone().into(), Some(TrieValue::Account(account1))), + // path2 is omitted - effectively deleted + ], + ) + .unwrap(); + let root_after_deletion = context3.root_node_hash; + + // The overlay root with tombstone should match the root after actual deletion + assert_eq!( + overlay_root_with_deletion, root_after_deletion, + "Tombstone overlay root should match actual deletion root" + ); + + // Both should equal the original root with just account1 + assert_eq!( + overlay_root_with_deletion, root_with_account1, + "After deleting account2, root should match original single-account root" + ); + assert_eq!( + root_after_deletion, root_with_account1, + "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.into()))), + (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2.into()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + assert_ne!(initial_root, EMPTY_ROOT_HASH); + + // 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(TrieValue::Storage(new_storage_value1.into()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the storage change + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + ( + storage_path1.full_path(), + Some(TrieValue::Storage(new_storage_value1.into())), + ), + (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2.into()))), + ], + ) + .unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!( + overlay_root, committed_root, + "Storage overlay root should match committed root" + ); + + // 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(TrieValue::Storage(storage_value3.into()))); + let overlay2 = overlay_mut2.freeze(); + + let overlay_root2 = + db.storage_engine.compute_root_with_overlay(&context, &overlay2).unwrap(); + assert_ne!(overlay_root2, initial_root); + assert_ne!(overlay_root2, overlay_root); + + // Verify by committing the new storage + let mut verification_context2 = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context2, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1.into()))), + (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2.into()))), + (storage_path3.full_path(), Some(TrieValue::Storage(storage_value3.into()))), + ], + ) + .unwrap(); + let committed_root2 = verification_context2.root_node_hash; + + assert_eq!( + overlay_root2, committed_root2, + "New storage overlay root should match committed root" + ); + + // 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(); + + let overlay_root3 = + db.storage_engine.compute_root_with_overlay(&context, &overlay3).unwrap(); + assert_ne!(overlay_root3, initial_root); + + // Verify by committing the deletion + let mut verification_context3 = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context3, + &mut [ + (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), + (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1.into()))), + // storage_path2 is omitted - effectively deleted + ], + ) + .unwrap(); + let committed_root3 = verification_context3.root_node_hash; + + assert_eq!( + overlay_root3, committed_root3, + "Storage deletion overlay root should match committed root" + ); + + // 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(TrieValue::Account(updated_account.clone()))); + overlay_mut4 + .insert(storage_path1.full_path(), Some(TrieValue::Storage(new_storage_value1.into()))); + let overlay4 = overlay_mut4.freeze(); + + let overlay_root4 = + db.storage_engine.compute_root_with_overlay(&context, &overlay4).unwrap(); + assert_ne!(overlay_root4, initial_root); + + // Note: For combined account+storage changes, the account overlay takes precedence + // and storage overlays should still be applied to compute the new storage root + // This is a complex case that exercises our account overlay + storage overlay logic + } + + #[test] + fn test_overlay_performance_cases() { + let tmp_dir = TempDir::new("test_overlay_performance_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 a larger trie with multiple levels to test add_branch optimization + let accounts_data = [ + ([0x0, 0x0], Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY)), + ([0x0, 0x1], Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY)), + ([0x1, 0x0], Account::new(3, U256::from(300), EMPTY_ROOT_HASH, KECCAK_EMPTY)), + ([0x1, 0x1], Account::new(4, U256::from(400), EMPTY_ROOT_HASH, KECCAK_EMPTY)), + ([0x2, 0x0], Account::new(5, U256::from(500), EMPTY_ROOT_HASH, KECCAK_EMPTY)), + ]; + + let mut changes = Vec::new(); + for (prefix, account) in accounts_data.iter() { + let mut nibbles = vec![prefix[0], prefix[1]]; + // Pad to 64 nibbles for address path + nibbles.extend(vec![0x0; 62]); + let path = AddressPath::new(Nibbles::from_nibbles(nibbles)); + changes.push((path.into(), Some(TrieValue::Account(account.clone())))); + } + + db.storage_engine.set_values(&mut context, &mut changes).unwrap(); + let initial_root = context.root_node_hash; + + // Test: Modify only one leaf in a large subtree + // This should trigger add_branch optimization for unaffected subtrees + let mut nibbles = vec![0x0, 0x0]; + nibbles.extend(vec![0x0; 62]); + let modify_path = AddressPath::new(Nibbles::from_nibbles(nibbles)); + + let updated_account = Account::new(10, U256::from(1000), EMPTY_ROOT_HASH, KECCAK_EMPTY); + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut + .insert(modify_path.clone().into(), Some(TrieValue::Account(updated_account.clone()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify the computation is correct by comparing with direct modification + let mut verification_context = db.storage_engine.write_context(); + let mut verification_changes = Vec::new(); + for (i, (prefix, account)) in accounts_data.iter().enumerate() { + let mut nibbles = vec![prefix[0], prefix[1]]; + nibbles.extend(vec![0x0; 62]); + let path = AddressPath::new(Nibbles::from_nibbles(nibbles)); + + if i == 0 { + // Use updated account for first entry + verification_changes + .push((path.into(), Some(TrieValue::Account(updated_account.clone())))); + } else { + verification_changes.push((path.into(), Some(TrieValue::Account(account.clone())))); + } + } + + db.storage_engine.set_values(&mut verification_context, &mut verification_changes).unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!( + overlay_root, committed_root, + "Performance test: Overlay root should match committed root" + ); + } + + #[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).into()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + + // 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(TrieValue::Storage(U256::from(222).into()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the addition + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine.set_values(&mut verification_context, &mut [ + (account_path.into(), Some(TrieValue::Account(account))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111).into()))), // Original + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), // New + ]).unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!(overlay_root, committed_root, "Adding storage slot via overlay should match"); + } + + #[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).into()))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + + // Test: Modify just one storage value per account via overlay + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut + .insert(storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))); + overlay_mut + .insert(storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888).into()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the changes + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context, + &mut [ + (account1_path.into(), Some(TrieValue::Account(account1))), + (account2_path.into(), Some(TrieValue::Account(account2))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888).into()))), + ], + ) + .unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!( + overlay_root, committed_root, + "Minimal multi-account storage overlay should match" + ); + } + + #[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).into()))), + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + + // 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(TrieValue::Storage(U256::from(1111).into()))); + + // 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(TrieValue::Storage(U256::from(444).into()))); + + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing all changes + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine.set_values(&mut verification_context, &mut [ + (account_path.into(), Some(TrieValue::Account(account))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(1111).into()))), // Modified + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), // Unchanged + (storage_path3.full_path(), Some(TrieValue::Storage(U256::from(444).into()))), // New + ]).unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!( + overlay_root, committed_root, + "Multiple storage overlays same account should match" + ); + } + + #[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).into()))), + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + + // 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(TrieValue::Storage(U256::from(999).into()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing: modify slot1, keep slot2 unchanged + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine.set_values(&mut verification_context, &mut [ + (account_path.into(), Some(TrieValue::Account(account))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(999).into()))), // Modified + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), // Unchanged + ]).unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!( + overlay_root, committed_root, + "Single account partial storage overlay should match" + ); + } + + #[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).into()))), + (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + (storage2_path1.full_path(), Some(TrieValue::Storage(U256::from(333).into()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + + // 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(TrieValue::Storage(U256::from(1111).into()))); + + // 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(TrieValue::Storage(U256::from(444).into()))); + + // Modify account2's storage + overlay_mut + .insert(storage2_path1.full_path(), Some(TrieValue::Storage(U256::from(3333).into()))); + + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing all changes + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context, + &mut [ + (account1_path.into(), Some(TrieValue::Account(account1))), + (account2_path.into(), Some(TrieValue::Account(account2))), + (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(1111).into()))), + (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + (storage1_path3.full_path(), Some(TrieValue::Storage(U256::from(444).into()))), + (storage2_path1.full_path(), Some(TrieValue::Storage(U256::from(3333).into()))), + ], + ) + .unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!( + overlay_root, committed_root, + "Multi-account storage overlay root should match committed root" + ); + } + + #[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).into()))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + + // Test: Modify just one storage slot for account1 + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut + .insert(storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the change + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine.set_values(&mut verification_context, &mut [ + (account1_path.into(), Some(TrieValue::Account(account1))), + (account2_path.into(), Some(TrieValue::Account(account2))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))), // Modified + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), // Unchanged + ]).unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!( + overlay_root, committed_root, + "Debug: Simple multi-account storage should match" + ); + } + + #[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).into()))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + + // Test: Modify storage for BOTH accounts + let mut overlay_mut = OverlayStateMut::new(); + overlay_mut + .insert(storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))); + overlay_mut + .insert(storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888).into()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing both changes + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine.set_values(&mut verification_context, &mut [ + (account1_path.into(), Some(TrieValue::Account(account1))), + (account2_path.into(), Some(TrieValue::Account(account2))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))), // Modified + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888).into()))), // Modified + ]).unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!( + overlay_root, committed_root, + "Debug: Both accounts storage changes should match" + ); + } + + #[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).into()))), + ], + ) + .unwrap(); + + let initial_root = context.root_node_hash; + + // 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(TrieValue::Storage(U256::from(444).into()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the addition + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine.set_values(&mut verification_context, &mut [ + (account1_path.into(), Some(TrieValue::Account(account1))), + (account2_path.into(), Some(TrieValue::Account(account2))), + (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(111).into()))), // Original + (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(444).into()))), // New + ]).unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!(overlay_root, committed_root, "Debug: Adding new storage should match"); + } + + #[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(); + + let initial_root = context.root_node_hash; + + // 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(TrieValue::Storage(storage_value.into()))); + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the storage addition + let mut verification_context = db.storage_engine.write_context(); + db.storage_engine + .set_values( + &mut verification_context, + &mut [ + (account_path.into(), Some(TrieValue::Account(account))), + (storage_path.full_path(), Some(TrieValue::Storage(storage_value.into()))), + ], + ) + .unwrap(); + let committed_root = verification_context.root_node_hash; + + assert_eq!( + overlay_root, committed_root, + "Empty account storage overlay root should match committed root" + ); + } +} diff --git a/src/transaction.rs b/src/transaction.rs index 5005bb88..778bf9bc 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -6,6 +6,7 @@ use crate::{ context::TransactionContext, database::Database, node::TrieValue, + overlay::OverlayStateMut, path::{AddressPath, StoragePath}, storage::proofs::AccountProof, }; @@ -46,6 +47,7 @@ pub struct Transaction { context: TransactionContext, database: DB, pending_changes: HashMap>, + overlay_state: Option, _marker: std::marker::PhantomData, } @@ -56,6 +58,7 @@ impl, K: TransactionKind> Transaction { context, database, pending_changes: HashMap::new(), + overlay_state: None, _marker: std::marker::PhantomData, } } diff --git a/src/transaction/error.rs b/src/transaction/error.rs index eb3ea809..b4262da4 100644 --- a/src/transaction/error.rs +++ b/src/transaction/error.rs @@ -1,11 +1,19 @@ use std::{error::Error, fmt}; #[derive(Debug)] -pub struct TransactionError; +pub enum TransactionError { + /// Generic transaction error (for backward compatibility) + Generic, + /// Overlay functionality is not enabled for this transaction + OverlayNotEnabled, +} impl fmt::Display for TransactionError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - write!(f, "transaction error") + match self { + Self::Generic => write!(f, "transaction error"), + Self::OverlayNotEnabled => write!(f, "overlay functionality is not enabled for this transaction"), + } } } diff --git a/src/transaction/manager.rs b/src/transaction/manager.rs index 3cb8efd5..465ee9b0 100644 --- a/src/transaction/manager.rs +++ b/src/transaction/manager.rs @@ -25,7 +25,7 @@ impl TransactionManager { pub fn begin_rw(&mut self, snapshot_id: SnapshotId) -> Result { // only allow one writable transaction at a time if self.has_writer { - return Err(TransactionError); + return Err(TransactionError::Generic); } self.has_writer = true; self.open_txs.push(snapshot_id - 1); From fef1020668fcb7c0b4e88025a19eabb5b03a4eef Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Fri, 25 Jul 2025 12:38:00 -0700 Subject: [PATCH 02/17] Zero copy subslicing of overlay state --- src/overlay.rs | 203 ++++++++++++++++++++++++++++++++------- src/transaction/error.rs | 4 +- 2 files changed, 172 insertions(+), 35 deletions(-) diff --git a/src/overlay.rs b/src/overlay.rs index 2a12698f..74cd8c28 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -3,6 +3,7 @@ use crate::{ pointer::Pointer, }; use alloy_trie::Nibbles; +use std::sync::Arc; /// Mutable overlay state that accumulates changes during transaction building. /// Changes are stored unsorted for fast insertion, then sorted when frozen. @@ -44,7 +45,7 @@ impl OverlayStateMut { /// 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 { changes: Vec::new().into_boxed_slice(), prefix_offset: 0 }; + return OverlayState { data: Arc::new([]), start_idx: 0, end_idx: 0, prefix_offset: 0 }; } // Sort by path @@ -66,57 +67,89 @@ impl OverlayStateMut { } } - OverlayState { changes: deduped.into_boxed_slice(), prefix_offset: 0 } + 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 sub-slicing. +/// 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 { - changes: Box<[(Nibbles, Option)]>, + data: Arc<[(Nibbles, Option)]>, + start_idx: usize, + end_idx: usize, prefix_offset: usize, } +// SAFETY: OverlayState is thread-safe because: +// - Arc<[...]> provides thread-safe reference counting +// - All fields are immutable after construction +// - No interior mutability is used +// - All operations are read-only or create new instances +unsafe impl Send for OverlayState {} +unsafe impl Sync for OverlayState {} + impl OverlayState { /// Creates an empty overlay state. pub fn empty() -> Self { - Self { changes: Vec::new().into_boxed_slice(), prefix_offset: 0 } + Self { data: Arc::new([]), start_idx: 0, end_idx: 0, prefix_offset: 0 } + } + + pub fn data(&self) -> &[(Nibbles, Option)] { + &self.data + } + + /// Returns the effective slice of data for this overlay, respecting bounds. + 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.changes.len() + self.end_idx - self.start_idx } /// Returns true if the overlay is empty. pub fn is_empty(&self) -> bool { - self.changes.is_empty() + 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: &Nibbles) -> Option<&Option> { - match self.changes.binary_search_by(|(p, _)| self.compare_with_offset(p, path)) { - Ok(index) => Some(&self.changes[index].1), + 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, } } /// Finds all entries that have the given prefix. - /// Returns a sub-slice of the overlay containing only matching entries. + /// Returns a zero-copy sub-slice of the overlay containing only matching entries. /// The prefix is compared against offset-adjusted paths. pub fn find_prefix_range(&self, prefix: &Nibbles) -> OverlayState { - if self.changes.is_empty() { + if self.is_empty() { return OverlayState::empty(); } + let slice = self.effective_slice(); let mut start_idx = None; - let mut end_idx = self.changes.len(); + let mut end_idx = slice.len(); // Find the range of entries that start with the prefix after applying offset - for (i, (path, _)) in self.changes.iter().enumerate() { + for (i, (path, _)) in slice.iter().enumerate() { if path.len() >= self.prefix_offset { let adjusted_path = path.slice(self.prefix_offset..); if adjusted_path.len() >= prefix.len() && adjusted_path[..prefix.len()] == *prefix { @@ -133,7 +166,9 @@ impl OverlayState { match start_idx { Some(start) => OverlayState { - changes: self.changes[start..end_idx].to_vec().into_boxed_slice(), + data: Arc::clone(&self.data), + start_idx: self.start_idx + start, + end_idx: self.start_idx + end_idx, prefix_offset: self.prefix_offset, }, None => OverlayState::empty(), @@ -155,9 +190,10 @@ impl OverlayState { // Validate that all paths share the same prefix up to total_offset // and that offset-adjusted paths maintain sort order - if self.changes.len() > 1 { - let first_path = &self.changes[0].0; - let last_path = &self.changes[self.changes.len() - 1].0; + 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 { @@ -177,7 +213,12 @@ impl OverlayState { } } - OverlayState { changes: self.changes.clone(), prefix_offset: total_offset } + OverlayState { + data: Arc::clone(&self.data), + start_idx: self.start_idx, + end_idx: self.end_idx, + prefix_offset: total_offset, + } } /// Helper method to compare a stored path with a target path, applying prefix_offset. @@ -195,31 +236,40 @@ impl OverlayState { } } - /// Creates a sub-slice of the overlay from the given range. + /// 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. + /// 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. + /// Perfect for thread pool scenarios where each thread processes a disjoint range. pub fn sub_slice(&self, start: usize, end: usize) -> OverlayState { - let end = end.min(self.changes.len()); + let current_len = self.len(); + let end = end.min(current_len); let start = start.min(end); OverlayState { - changes: self.changes[start..end].to_vec().into_boxed_slice(), + data: Arc::clone(&self.data), + start_idx: self.start_idx + start, + end_idx: self.start_idx + end, prefix_offset: self.prefix_offset, } } - /// Creates a sub-slice containing only changes that come before the given prefix. + /// Creates a zero-copy sub-slice containing only changes that come before the given prefix. /// The prefix comparison takes prefix_offset into account. pub fn sub_slice_before_prefix(&self, prefix: &Nibbles) -> OverlayState { - let index = self - .changes + let slice = self.effective_slice(); + let index = slice .binary_search_by(|(p, _)| self.compare_with_offset(p, prefix)) .unwrap_or_else(|i| i); // Insert position is exactly what we want for "before" self.sub_slice(0, index) } - /// Creates a sub-slice containing only changes that come strictly after the given prefix. - /// The prefix comparison takes prefix_offset into account. + /// Creates a zero-copy sub-slice containing only changes that come strictly after the given + /// prefix. The prefix comparison takes prefix_offset into account. pub fn sub_slice_after_prefix(&self, prefix: &Nibbles) -> OverlayState { // Find the next key after prefix by incrementing prefix let next_prefix = match prefix.increment() { @@ -230,25 +280,25 @@ impl OverlayState { } }; - let index = self - .changes + let slice = self.effective_slice(); + let index = slice .binary_search_by(|(p, _)| self.compare_with_offset(p, &next_prefix)) .unwrap_or_else(|i| i); - self.sub_slice(index, self.changes.len()) + self.sub_slice(index, slice.len()) } - /// Creates a sub-slice containing only changes that affect the subtree + /// 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: &Nibbles) -> OverlayState { self.find_prefix_range(prefix) } - /// Returns an iterator over all changes in the overlay. + /// 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.changes.iter().filter_map(move |(path, value)| { - if path.len() >= self.prefix_offset { + self.effective_slice().iter().filter_map(move |(path, value)| { + if path.len() > self.prefix_offset { let adjusted_path = path.slice(self.prefix_offset..); Some((adjusted_path, value)) } else { @@ -731,4 +781,89 @@ mod tests { assert!(storage_with_prefix5.contains(&Nibbles::from_nibbles([5, 5]))); assert!(storage_with_prefix5.contains(&Nibbles::from_nibbles([5, 6]))); } + + #[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(TrieValue::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::{sync::Arc as StdArc, 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(TrieValue::Account(account.clone())), + ); + } + + let frozen = mutable.freeze(); + let shared_overlay = StdArc::new(frozen); + + // Spawn multiple threads that work with different sub-slices + let handles: Vec<_> = (0..4) + .map(|i| { + let overlay = StdArc::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]); + } } diff --git a/src/transaction/error.rs b/src/transaction/error.rs index b4262da4..de0bf916 100644 --- a/src/transaction/error.rs +++ b/src/transaction/error.rs @@ -12,7 +12,9 @@ impl fmt::Display for TransactionError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { match self { Self::Generic => write!(f, "transaction error"), - Self::OverlayNotEnabled => write!(f, "overlay functionality is not enabled for this transaction"), + Self::OverlayNotEnabled => { + write!(f, "overlay functionality is not enabled for this transaction") + } } } } From f19588dd52fc9eb7fd8fe16debc1fe9572018372 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Fri, 25 Jul 2025 12:39:17 -0700 Subject: [PATCH 03/17] Support overlayed new accounts with storage --- src/storage/overlay_root.rs | 422 +++++++++++++++++++++++------------- 1 file changed, 277 insertions(+), 145 deletions(-) diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 779f7de8..13307144 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -22,12 +22,16 @@ impl StorageEngine { return Ok(context.root_node_hash); } - let mut hash_builder = HashBuilder::default(); + let mut hash_builder = HashBuilder::default().with_updates(true); // Use proper trie traversal with overlay integration self.traverse_with_overlay(context, overlay, &mut hash_builder)?; - let root = hash_builder.root(); + let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); + // println!("Updated branch nodes: {:?}", updated_branch_nodes); + + let root = hash_builder.root(); // This will clear the hash builder + // println!("Root: {:?}", root); Ok(root) } @@ -50,12 +54,7 @@ impl StorageEngine { )?; } else { // No root page, just add all overlay changes to the hash builder - for (path, value) in overlay.iter() { - if let Some(trie_value) = value { - let rlp_encoded = self.encode_trie_value(trie_value)?; - hash_builder.add_leaf(path, &rlp_encoded); - } - } + self.process_overlay_state(context, overlay, hash_builder)?; } Ok(()) @@ -108,12 +107,7 @@ impl StorageEngine { let post_overlay = overlay.sub_slice_after_prefix(&full_path); // Add all pre-overlay changes - for (path, value) in pre_overlay.iter() { - if let Some(trie_value) = value { - let rlp_encoded = self.encode_trie_value(trie_value)?; - hash_builder.add_leaf(path, &rlp_encoded); - } - } + self.process_overlay_state(context, &pre_overlay, hash_builder)?; // Check if this node is affected by overlay if overlay.affects_path(&full_path) { @@ -131,20 +125,43 @@ impl StorageEngine { self.handle_unaffected_node( context, hash_builder, - slotted_page, &node, + path_prefix, &full_path, node_hash, )?; } - for (path, value) in post_overlay.iter() { - if let Some(trie_value) = value { - let rlp_encoded = self.encode_trie_value(trie_value)?; - hash_builder.add_leaf(path, &rlp_encoded); + // Add all post-overlay changes + self.process_overlay_state(context, &post_overlay, hash_builder)?; + + Ok(()) + } + + fn process_overlay_state( + &self, + context: &TransactionContext, + overlay: &OverlayState, + hash_builder: &mut HashBuilder, + ) -> Result<(), Error> { + for (path, value) in overlay.iter() { + // only process leaves in the current trie + if path.len() == 64 { + match value { + Some(TrieValue::Account(_)) => { + self.handle_overlayed_account(context, overlay, hash_builder, &path)?; + } + Some(TrieValue::Storage(storage_value)) => { + let rlp_encoded = + self.encode_trie_value(&TrieValue::Storage(*storage_value))?; + hash_builder.add_leaf(path, &rlp_encoded); + } + None => { + // Tombstone - skip + } + } } } - Ok(()) } @@ -158,17 +175,15 @@ impl StorageEngine { node: &Node, full_path: &Nibbles, ) -> Result<(), Error> { - use crate::node::NodeKind; - match node.kind() { - NodeKind::AccountLeaf { .. } if full_path.len() == 64 => { + NodeKind::AccountLeaf { .. } => { // Account leaf at address level - handle potential storage overlays self.handle_account_leaf_with_overlay( context, overlay, hash_builder, slotted_page, - node, + Some(node), full_path, )?; } @@ -201,19 +216,6 @@ impl StorageEngine { full_path, )?; } - _ => { - // Other leaf nodes - handle with simple leaf logic - if let Some(overlay_value) = overlay.lookup(full_path) { - if let Some(trie_value) = overlay_value { - let rlp_encoded = self.encode_trie_value(trie_value)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - } - } else { - let node_value = node.value()?; - let rlp_encoded = self.encode_trie_value(&node_value)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - } - } } Ok(()) } @@ -223,8 +225,8 @@ impl StorageEngine { &self, _context: &TransactionContext, hash_builder: &mut HashBuilder, - _slotted_page: &SlottedPage, node: &Node, + path_prefix: &Nibbles, full_path: &Nibbles, node_hash: Option, ) -> Result<(), Error> { @@ -233,22 +235,15 @@ impl StorageEngine { // Account leaf - use disk value directly let node_value = node.value()?; let rlp_encoded = self.encode_trie_value(&node_value)?; + // println!("Adding unaffected leaf: {:?}", full_path); hash_builder.add_leaf(full_path.clone(), &rlp_encoded); } NodeKind::Branch { .. } => { // Branch node - since it's unaffected by overlay, we can use its existing hash // This is much more efficient than recursing into all children - if let Some(hash) = node_hash { - // Use the precomputed hash directly for this branch subtree - // stored_in_database = true since this is an existing node from disk - hash_builder.add_branch(full_path.clone(), hash, true); - } else { - // Fallback: if we don't have the hash, we need to compute it - // This shouldn't normally happen but provides a safety net - let node_value = node.value()?; - let rlp_encoded = self.encode_trie_value(&node_value)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - } + // println!("Adding branch node: {:?}", path_prefix); + let hash = node_hash.unwrap(); + hash_builder.add_branch(path_prefix.clone(), hash, true); } } Ok(()) @@ -291,17 +286,15 @@ impl StorageEngine { child_hash, )?; } else if let Some(child_page_id) = child_location.page_id() { - // Child is in a different page - add safety check - if self.page_exists(context, child_page_id) { - self.traverse_page_with_overlay( - context, - &child_overlay, - hash_builder, - child_page_id, - &child_path, - child_hash, - )?; - } + // Child is in a different page + self.traverse_page_with_overlay( + context, + &child_overlay, + hash_builder, + child_page_id, + &child_path, + child_hash, + )?; } } else { // No child exists on disk, but check if overlay has leaves for this subtree @@ -328,95 +321,128 @@ impl StorageEngine { overlay: &OverlayState, hash_builder: &mut HashBuilder, slotted_page: &SlottedPage, - node: &Node, + node: Option<&Node>, full_path: &Nibbles, ) -> Result<(), Error> { // Get the original account from the node - let original_account = match node.value()? { + let original_account = node.map(|n| match n.value().unwrap() { TrieValue::Account(account) => account, _ => { - // This shouldn't happen for AccountLeaf nodes, but handle gracefully - let rlp_encoded = self.encode_trie_value(&node.value()?)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - return Ok(()); + panic!("Node is not an account leaf"); } - }; + }); // Check if this account is directly overlaid (account-level changes) - if let Some(overlay_value) = overlay.lookup(full_path) { + let account = if let Some(overlay_value) = overlay.lookup(full_path) { // Direct account overlay - use overlay value, but still check for storage overlays if let Some(trie_value) = overlay_value { // Even with account overlay, we might need to compute storage root for storage // overlays if let TrieValue::Account(overlay_account) = trie_value { - // Check if there are storage overlays for this account - let storage_overlays = overlay.find_prefix_range(full_path); - let has_storage_overlays = storage_overlays - .iter() - .any(|(path, _)| path.len() == 128 && path.len() > full_path.len()); - - if has_storage_overlays { - // Compute new storage root with overlays - let new_storage_root = self.compute_storage_root_iteratively( - context, - node, // Use the original node for storage traversal - full_path, - &storage_overlays, - slotted_page, - )?; - - // Create modified account with new storage root - let mut modified_account = overlay_account.clone(); - modified_account.storage_root = new_storage_root; - let rlp_encoded = - self.encode_trie_value(&TrieValue::Account(modified_account))?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - } else { - // No storage overlays, use account overlay as-is - let rlp_encoded = self.encode_trie_value(trie_value)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - } + Some(overlay_account.clone()) } else { - // Not an account value, just use as-is - let rlp_encoded = self.encode_trie_value(trie_value)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + panic!("Overlay value is not an account"); } + } else { + // If overlay_value is None, it's a deletion - skip + println!("Skipping deletion: {full_path:?}"); + None } - // If overlay_value is None, it's a deletion - skip - return Ok(()); - } + } else { + original_account + }; - // Check if there are any storage overlays for this account - // Storage overlays have 128-nibble paths that start with this account's 64-nibble path - let storage_overlays = overlay.find_prefix_range(full_path); - let has_storage_overlays = storage_overlays - .iter() - .any(|(path, _)| path.len() == 128 && path.len() > full_path.len()); + if let Some(account) = account { + // Check if there are any storage overlays for this account + // Storage overlays have 128-nibble paths that start with this account's 64-nibble path + let storage_overlays = overlay.find_prefix_range(full_path); + let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() == 128); - if !has_storage_overlays { - // No storage overlays for this account, use original account - let rlp_encoded = self.encode_trie_value(&TrieValue::Account(original_account))?; + if !has_storage_overlays { + // No storage overlays for this account, use original account + let rlp_encoded = self.encode_trie_value(&TrieValue::Account(account))?; + println!("Adding account leaf: {full_path:?}"); + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + return Ok(()); + } + + // We have storage overlays, need to compute a new storage root using iterative approach + let new_storage_root = self.compute_storage_root_iteratively( + context, + node, + &storage_overlays, + slotted_page, + )?; + + // Create a modified account with the new storage root + let mut modified_account = account; + modified_account.storage_root = new_storage_root; + + // Add the modified account to the main HashBuilder + let rlp_encoded = self.encode_trie_value(&TrieValue::Account(modified_account))?; hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - return Ok(()); } + Ok(()) + } - // We have storage overlays, need to compute a new storage root using iterative approach - let new_storage_root = self.compute_storage_root_iteratively( - context, - node, - full_path, - &storage_overlays, - slotted_page, - )?; + fn handle_overlayed_account( + &self, + context: &TransactionContext, + overlay: &OverlayState, + hash_builder: &mut HashBuilder, + full_path: &Nibbles, + ) -> Result<(), Error> { + // Check if the overlay is an account leaf or None + let account = if let Some(overlay_value) = overlay.lookup(full_path) { + // Direct account overlay - use overlay value, but still check for storage overlays + if let Some(trie_value) = overlay_value { + // Even with account overlay, we might need to compute storage root for storage + // overlays + if let TrieValue::Account(overlay_account) = trie_value { + Some(overlay_account.clone()) + } else { + panic!("Overlay value is not an account"); + } + } else { + // If overlay_value is None, it's a deletion - skip + println!("Skipping deletion: {full_path:?}"); + None + } + } else { + None + }; + + if let Some(account) = account { + // Check if there are any storage overlays for this account + // Storage overlays have 128-nibble paths that start with this account's 64-nibble path + let storage_overlays = overlay.find_prefix_range(full_path); + let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() == 128); + + if !has_storage_overlays { + // No storage overlays for this account, use original account + let rlp_encoded = self.encode_trie_value(&TrieValue::Account(account))?; + println!("Adding account leaf: {full_path:?}"); + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + return Ok(()); + } + + // We have storage overlays, need to compute a new storage root + let storage_overlay = storage_overlays.with_prefix_offset(64); + + let mut storage_hash_builder = HashBuilder::default().with_updates(true); - // Create a modified account with the new storage root - let mut modified_account = original_account; - modified_account.storage_root = new_storage_root; + self.process_overlay_state(context, &storage_overlay, &mut storage_hash_builder)?; - // Add the modified account to the main HashBuilder - let rlp_encoded = self.encode_trie_value(&TrieValue::Account(modified_account))?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + let new_storage_root = storage_hash_builder.root(); + // Create a modified account with the new storage root + let mut modified_account = account; + modified_account.storage_root = new_storage_root; + + // Add the modified account to the main HashBuilder + let rlp_encoded = self.encode_trie_value(&TrieValue::Account(modified_account))?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + } Ok(()) } @@ -425,8 +451,7 @@ impl StorageEngine { fn compute_storage_root_iteratively( &self, context: &TransactionContext, - account_node: &Node, - _account_path: &Nibbles, + account_node: Option<&Node>, storage_overlays: &OverlayState, slotted_page: &SlottedPage, ) -> Result { @@ -434,19 +459,22 @@ impl StorageEngine { // This converts 128-nibble storage paths to 64-nibble storage-relative paths let storage_overlay = storage_overlays.with_prefix_offset(64); - let mut storage_hash_builder = HashBuilder::default(); + let mut storage_hash_builder = HashBuilder::default().with_updates(true); // Get the original account's storage root to determine if we need to traverse existing // storage - let original_storage_root = match account_node.value()? { - TrieValue::Account(account) => account.storage_root, - _ => return Err(Error::NodeError(crate::node::NodeError::NoValue)), /* This shouldn't happen for account nodes */ - }; + let original_storage_root = account_node + .map(|n| match n.value().unwrap() { + TrieValue::Account(account) => account.storage_root, + _ => panic!("Node is not an account leaf"), /* This shouldn't happen for account + * nodes */ + }) + .unwrap_or(alloy_trie::EMPTY_ROOT_HASH); // If the account has existing storage, traverse the existing storage trie if original_storage_root != alloy_trie::EMPTY_ROOT_HASH { // Try to find the storage root page from the account node's direct child - if let Ok(Some(storage_root_pointer)) = account_node.direct_child() { + if let Ok(Some(storage_root_pointer)) = account_node.unwrap().direct_child() { let storage_root_location = storage_root_pointer.location(); let storage_root_hash = storage_root_pointer.rlp().as_hash(); @@ -475,23 +503,17 @@ impl StorageEngine { } } else { // No existing storage, just add overlay changes - for (path, value) in storage_overlay.iter() { - if let Some(trie_value) = value { - let rlp_encoded = self.encode_trie_value(trie_value)?; - storage_hash_builder.add_leaf(path, &rlp_encoded); - } - } + self.process_overlay_state(context, &storage_overlay, &mut storage_hash_builder)?; } + let (mut storage_hash_builder, updated_branch_nodes) = storage_hash_builder.split(); + println!("Updated storage branch nodes: {updated_branch_nodes:?}"); + let computed_root = storage_hash_builder.root(); + println!("Computed storage root: {computed_root:?}"); Ok(computed_root) } - /// Check if a page exists in the storage - fn page_exists(&self, context: &TransactionContext, page_id: PageId) -> bool { - self.get_page(context, page_id).is_ok() - } - /// Helper to encode TrieValue as RLP fn encode_trie_value(&self, trie_value: &TrieValue) -> Result, Error> { let rlp_encoded = match trie_value { @@ -517,8 +539,9 @@ mod tests { account::Account, database::Database, node::TrieValue, overlay::OverlayStateMut, path::AddressPath, }; - use alloy_primitives::{address, U256}; + use alloy_primitives::{address, Address, U256}; use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY}; + use rand::Rng; use tempdir::TempDir; #[test] @@ -1655,4 +1678,113 @@ mod tests { "Empty account storage overlay root should match committed root" ); } + + #[test] + fn test_1000_accounts_with_10_overlay() { + 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(); + + let initial_root = context.root_node_hash; + + // 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(TrieValue::Account(new_account))); + } else { + // For other half, mark for deletion + overlay_mut.insert(path.clone(), None); + } + } + } + + let overlay = overlay_mut.freeze(); + + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the storage addition + db.storage_engine.set_values(&mut context, overlay.data().to_vec().as_mut()).unwrap(); + let committed_root = context.root_node_hash; + + db.storage_engine + .print_page(&context, std::io::stdout(), context.root_node_page_id) + .unwrap(); + + assert_eq!(overlay_root, committed_root, "1000 accounts with 10 overlay should match"); + } + + #[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 initial_root = context.root_node_hash; + + 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.into(), Some(TrieValue::Account(new_account.clone()))); + + let storage_key = U256::from(42); + let storage_path = + crate::path::StoragePath::for_address_and_slot(new_address, storage_key.into()); + let storage_value = U256::from(123); + overlay_mut + .insert(storage_path.full_path(), Some(TrieValue::Storage(storage_value.into()))); + + let overlay = overlay_mut.freeze(); + let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the storage addition + db.storage_engine.set_values(&mut context, overlay.data().to_vec().as_mut()).unwrap(); + let committed_root = context.root_node_hash; + + db.storage_engine + .print_page(&context, std::io::stdout(), context.root_node_page_id) + .unwrap(); + + assert_eq!(overlay_root, committed_root, "Overlay new account with storage should match"); + } } From 1bc786777de60f4a69ae44428f15dbdb93f637a4 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Mon, 28 Jul 2025 16:22:04 -0700 Subject: [PATCH 04/17] Correctly compute root with overlayed branches --- src/overlay.rs | 310 ++++++++++++++++--- src/storage/overlay_root.rs | 593 ++++++++++++++++++++---------------- src/transaction.rs | 3 - 3 files changed, 608 insertions(+), 298 deletions(-) diff --git a/src/overlay.rs b/src/overlay.rs index 74cd8c28..b6dc0910 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -1,15 +1,36 @@ use crate::{ + account::Account, node::{Node, TrieValue}, pointer::Pointer, }; +use alloy_primitives::{StorageValue, B256}; use alloy_trie::Nibbles; use std::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)>, + changes: Vec<(Nibbles, Option)>, } impl OverlayStateMut { @@ -25,7 +46,7 @@ impl OverlayStateMut { /// 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) { + 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. self.changes.push((path, value)); @@ -67,7 +88,7 @@ impl OverlayStateMut { } } - let data: Arc<[(Nibbles, Option)]> = Arc::from(deduped.into_boxed_slice()); + 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 } } @@ -86,7 +107,7 @@ impl OverlayStateMut { /// The Arc<[...]> provides thread-safe reference counting for the underlying data. #[derive(Debug, Clone)] pub struct OverlayState { - data: Arc<[(Nibbles, Option)]>, + data: Arc<[(Nibbles, Option)]>, start_idx: usize, end_idx: usize, prefix_offset: usize, @@ -106,12 +127,21 @@ impl OverlayState { Self { data: Arc::new([]), start_idx: 0, end_idx: 0, prefix_offset: 0 } } - pub fn data(&self) -> &[(Nibbles, Option)] { + pub fn data(&self) -> &[(Nibbles, Option)] { &self.data } + pub fn get(&self, index: usize) -> Option<(Nibbles, &Option)> { + let slice = self.effective_slice(); + if index < slice.len() { + Some((slice[index].0.slice(self.prefix_offset..), &slice[index].1)) + } else { + None + } + } + /// Returns the effective slice of data for this overlay, respecting bounds. - fn effective_slice(&self) -> &[(Nibbles, Option)] { + pub fn effective_slice(&self) -> &[(Nibbles, Option)] { &self.data[self.start_idx..self.end_idx] } @@ -128,7 +158,7 @@ impl OverlayState { /// 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: &Nibbles) -> Option<&Option> { + pub fn lookup(&self, path: &Nibbles) -> 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), @@ -258,6 +288,35 @@ impl OverlayState { } } + /// 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: &Nibbles, + ) -> (OverlayState, OverlayState, OverlayState) { + let slice = self.effective_slice(); + + let start_index = slice + .binary_search_by(|(p, _)| self.compare_with_offset(p, prefix)) + .unwrap_or_else(|i| i); + let end_index = match prefix.increment() { + Some(next) => slice + .binary_search_by(|(p, _)| self.compare_with_offset(p, &next)) + .unwrap_or_else(|i| i), + None => { + // Prefix is all 0xF's, nothing can come after it + slice.len() + } + }; + + 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 come before the given prefix. /// The prefix comparison takes prefix_offset into account. pub fn sub_slice_before_prefix(&self, prefix: &Nibbles) -> OverlayState { @@ -296,7 +355,7 @@ impl OverlayState { /// 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)> { + 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.slice(self.prefix_offset..); @@ -311,10 +370,6 @@ impl OverlayState { /// This includes exact matches and any changes to descendants of the path. /// The path comparison takes prefix_offset into account. pub fn affects_path(&self, path: &Nibbles) -> bool { - if self.is_empty() { - return false; - } - // Check for exact match or any path that starts with the given path after applying offset self.iter().any(|(overlay_path, _)| overlay_path.starts_with(path)) } @@ -428,7 +483,7 @@ mod tests { let path2 = Nibbles::from_nibbles([5, 6, 7, 8]); let account = test_account(); - overlay.insert(path1.clone(), Some(TrieValue::Account(account.clone()))); + overlay.insert(path1.clone(), Some(OverlayValue::Account(account.clone()))); overlay.insert(path2.clone(), None); // Tombstone assert!(!overlay.is_empty()); @@ -442,7 +497,7 @@ mod tests { let path2 = Nibbles::from_nibbles([5, 6, 7, 8]); let account = test_account(); - mutable.insert(path1.clone(), Some(TrieValue::Account(account.clone()))); + mutable.insert(path1.clone(), Some(OverlayValue::Account(account.clone()))); mutable.insert(path2.clone(), None); let frozen = mutable.freeze(); @@ -471,8 +526,8 @@ mod tests { 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(TrieValue::Account(account1))); - mutable.insert(path.clone(), Some(TrieValue::Account(account2.clone()))); + 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(); @@ -489,14 +544,22 @@ mod tests { let account = test_account(); // Add some paths with common prefixes - mutable - .insert(Nibbles::from_nibbles([1, 2, 3, 4]), Some(TrieValue::Account(account.clone()))); - mutable - .insert(Nibbles::from_nibbles([1, 2, 5, 6]), Some(TrieValue::Account(account.clone()))); - mutable - .insert(Nibbles::from_nibbles([1, 3, 7, 8]), Some(TrieValue::Account(account.clone()))); - mutable - .insert(Nibbles::from_nibbles([2, 3, 4, 5]), Some(TrieValue::Account(account.clone()))); + 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(); @@ -518,10 +581,14 @@ mod tests { let mut mutable = OverlayStateMut::new(); let account = test_account(); - mutable - .insert(Nibbles::from_nibbles([1, 2, 3, 4]), Some(TrieValue::Account(account.clone()))); - mutable - .insert(Nibbles::from_nibbles([1, 2, 5, 6]), Some(TrieValue::Account(account.clone()))); + 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(); @@ -649,15 +716,15 @@ mod tests { mutable.insert( Nibbles::from_nibbles(storage_path1.clone()), - Some(TrieValue::Account(account.clone())), + Some(OverlayValue::Account(account.clone())), ); mutable.insert( Nibbles::from_nibbles(storage_path2.clone()), - Some(TrieValue::Account(account.clone())), + Some(OverlayValue::Account(account.clone())), ); mutable.insert( Nibbles::from_nibbles(storage_path3.clone()), - Some(TrieValue::Account(account.clone())), + Some(OverlayValue::Account(account.clone())), ); let frozen = mutable.freeze(); @@ -712,7 +779,7 @@ mod tests { for path in &paths { mutable.insert( Nibbles::from_nibbles(path.clone()), - Some(TrieValue::Account(account.clone())), + Some(OverlayValue::Account(account.clone())), ); } @@ -749,17 +816,17 @@ mod tests { // Storage for account 1 mutable.insert( Nibbles::from_nibbles([1, 0, 0, 0, 5, 5]), - Some(TrieValue::Account(account.clone())), + Some(OverlayValue::Account(account.clone())), ); mutable.insert( Nibbles::from_nibbles([1, 0, 0, 0, 5, 6]), - Some(TrieValue::Account(account.clone())), + Some(OverlayValue::Account(account.clone())), ); // Storage for account 2 mutable.insert( Nibbles::from_nibbles([2, 0, 0, 0, 7, 7]), - Some(TrieValue::Account(account.clone())), + Some(OverlayValue::Account(account.clone())), ); let frozen = mutable.freeze(); @@ -791,7 +858,7 @@ mod tests { for i in 0..10 { mutable.insert( Nibbles::from_nibbles([i, 0, 0, 0]), - Some(TrieValue::Account(account.clone())), + Some(OverlayValue::Account(account.clone())), ); } @@ -834,7 +901,7 @@ mod tests { for i in 0..100 { mutable.insert( Nibbles::from_nibbles([i % 16, (i / 16) % 16, 0, 0]), - Some(TrieValue::Account(account.clone())), + Some(OverlayValue::Account(account.clone())), ); } @@ -866,4 +933,173 @@ mod tests { 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 = vec![ + // 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 = Nibbles::from_nibbles(vec![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(&Nibbles::from_nibbles(vec![0, 5, 6, 7]))); + assert!(before_paths.contains(&Nibbles::from_nibbles(vec![1, 0, 0, 0]))); + assert!(before_paths.contains(&Nibbles::from_nibbles(vec![1, 1, 9, 9]))); + + // 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(&Nibbles::from_nibbles(vec![1, 2]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles(vec![1, 2, 3, 4]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles(vec![1, 2, 5, 6]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles(vec![1, 2, 7, 8]))); + + // 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(&Nibbles::from_nibbles(vec![1, 3, 0, 0]))); + assert!(after_paths.contains(&Nibbles::from_nibbles(vec![2, 0, 0, 0]))); + assert!(after_paths.contains(&Nibbles::from_nibbles(vec![5, 5, 5, 5]))); + + // 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(vec![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(vec![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(vec![2, 3]); + + mutable.insert( + Nibbles::from_nibbles(vec![1, 9]), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles(vec![2, 2]), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert(prefix.clone(), Some(OverlayValue::Account(account.clone()))); // Exact match + mutable.insert( + Nibbles::from_nibbles(vec![2, 4]), + Some(OverlayValue::Account(account.clone())), + ); + mutable.insert( + Nibbles::from_nibbles(vec![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)); + } + + #[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 = vec![1, 0, 0, 0]; + + // Storage paths for the same account + let storage_paths = vec![ + vec![1, 0, 0, 0, 2, 0, 0, 0], // account + storage [2,0,0,0] + vec![1, 0, 0, 0, 5, 0, 0, 0], // account + storage [5,0,0,0] + vec![1, 0, 0, 0, 5, 1, 0, 0], // account + storage [5,1,0,0] + vec![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.clone()), + 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 = Nibbles::from_nibbles(vec![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(&Nibbles::from_nibbles(vec![2, 0, 0, 0]))); + + // 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(&Nibbles::from_nibbles(vec![5, 0, 0, 0]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles(vec![5, 1, 0, 0]))); + + // 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(&Nibbles::from_nibbles(vec![8, 0, 0, 0]))); + } } diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 13307144..413818ab 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -1,13 +1,14 @@ use crate::{ + account::Account, context::TransactionContext, node::{Node, NodeKind, TrieValue}, - overlay::OverlayState, + overlay::{OverlayState, OverlayValue}, page::{PageId, SlottedPage}, storage::engine::{Error, StorageEngine}, }; -use alloy_primitives::B256; +use alloy_primitives::{map::HashMap, B256, U256}; use alloy_rlp::encode; -use alloy_trie::{HashBuilder, Nibbles, TrieAccount}; +use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, TrieAccount}; impl StorageEngine { /// Computes the root hash with overlay changes without persisting them. @@ -16,10 +17,10 @@ impl StorageEngine { &self, context: &TransactionContext, overlay: &OverlayState, - ) -> Result { + ) -> Result<(B256, HashMap), Error> { if overlay.is_empty() { // No overlay changes, return current root - return Ok(context.root_node_hash); + return Ok((context.root_node_hash, HashMap::default())); } let mut hash_builder = HashBuilder::default().with_updates(true); @@ -32,7 +33,7 @@ impl StorageEngine { let root = hash_builder.root(); // This will clear the hash builder // println!("Root: {:?}", root); - Ok(root) + Ok((root, updated_branch_nodes)) } /// Helper method to traverse the existing trie and integrate with overlay @@ -103,38 +104,77 @@ impl StorageEngine { full }; - let pre_overlay = overlay.sub_slice_before_prefix(&full_path); - let post_overlay = overlay.sub_slice_after_prefix(&full_path); + let mut current_path = path_prefix.clone(); + while current_path.len() <= full_path.len() { + // println!("Current path: {:?}", current_path); + // println!("Full path: {:?}", full_path); + let (pre_overlay, with_prefix, _) = overlay.sub_slice_by_prefix(¤t_path); + // println!("Pre-overlay: {:?}", pre_overlay.effective_slice()); + // println!("With-prefix: {:?}", with_prefix.effective_slice()); + // println!("Post-overlay: {:?}", post_overlay.effective_slice()); + + // Add all pre-overlay changes + self.process_overlay_state(context, &pre_overlay, hash_builder)?; + + // Check if this node is affected by overlay + if with_prefix.is_empty() { + // Node not affected by overlay, use disk value as-is + match node.kind() { + NodeKind::AccountLeaf { .. } | NodeKind::StorageLeaf { .. } => { + // Leaf node - use disk value directly + let node_value = node.value()?; + let rlp_encoded = self.encode_trie_value(&node_value)?; + // println!("Adding unaffected leaf: {:?}", full_path); + hash_builder.add_leaf(full_path, &rlp_encoded); + } + NodeKind::Branch { .. } => { + // Branch node - since it's unaffected by overlay, we can use its existing + // hash This is much more efficient than recursing + // into all children println!("Adding branch node: + // {:?}", path_prefix); + let hash = node_hash.unwrap(); + // println!("Adding unaffected branch: {:?}", path_prefix); + hash_builder.add_branch(path_prefix.clone(), hash, true); + } + } + break; + } - // Add all pre-overlay changes - self.process_overlay_state(context, &pre_overlay, hash_builder)?; + let next_overlay_length; + if let Some((first_overlay_path, _)) = overlay.get(0) { + // If the first overlay path is the same as the current path, we've found an overlay + if first_overlay_path == current_path { + // println!("Direct overlay: {:?}", first_overlay_path); + // Direct overlay - the overlay path exactly matches this node's path + self.process_overlay_state(context, &with_prefix, hash_builder)?; + break; + } else { + next_overlay_length = first_overlay_path.len(); + } + } else { + panic!("Overlay path does not match node path {current_path:?}"); + } - // Check if this node is affected by overlay - if overlay.affects_path(&full_path) { - // This node or its descendants are affected by overlay - self.handle_affected_node_with_overlay( - context, - overlay, - hash_builder, - slotted_page, - &node, - &full_path, - )?; - } else { - // Node not affected by overlay, use disk value as-is - self.handle_unaffected_node( - context, - hash_builder, - &node, - path_prefix, - &full_path, - node_hash, - )?; - } + // If we've reached the end of the full path, we've found the node + if current_path.len() == full_path.len() { + // println!("Found node: {:?}", full_path); + self.handle_affected_node_with_overlay( + context, + overlay, + hash_builder, + slotted_page, + &node, + &full_path, + )?; + break; + } - // Add all post-overlay changes + // Fast forward the path length to the next length at which the overlay changes + let path_len = std::cmp::min(full_path.len(), next_overlay_length + 1); + current_path = full_path.slice(..path_len); + } + let post_overlay = overlay.sub_slice_after_prefix(¤t_path); self.process_overlay_state(context, &post_overlay, hash_builder)?; - Ok(()) } @@ -144,21 +184,33 @@ impl StorageEngine { overlay: &OverlayState, hash_builder: &mut HashBuilder, ) -> Result<(), Error> { + let mut last_processed_path: Option = None; for (path, value) in overlay.iter() { - // only process leaves in the current trie - if path.len() == 64 { - match value { - Some(TrieValue::Account(_)) => { - self.handle_overlayed_account(context, overlay, hash_builder, &path)?; - } - Some(TrieValue::Storage(storage_value)) => { - let rlp_encoded = - self.encode_trie_value(&TrieValue::Storage(*storage_value))?; - hash_builder.add_leaf(path, &rlp_encoded); - } - None => { - // Tombstone - skip - } + if let Some(ref last_processed_path) = last_processed_path { + if path.has_prefix(last_processed_path) { + // skip over all descendants of a processed path + continue; + } + } + match value { + Some(OverlayValue::Account(account)) => { + self.handle_overlayed_account(context, overlay, hash_builder, &path, account)?; + last_processed_path = Some(path); + } + Some(OverlayValue::Storage(storage_value)) => { + let rlp_encoded = self.encode_storage(storage_value)?; + // println!("Adding overlayed storage leaf: {:?}", path); + hash_builder.add_leaf(path.clone(), &rlp_encoded); + last_processed_path = Some(path); + } + Some(OverlayValue::Hash(hash)) => { + // println!("Adding overlayed branch: {:?}", path); + hash_builder.add_branch(path.clone(), *hash, false); + last_processed_path = Some(path); + } + None => { + // Tombstone - skip + last_processed_path = Some(path); } } } @@ -176,76 +228,32 @@ impl StorageEngine { full_path: &Nibbles, ) -> Result<(), Error> { match node.kind() { - NodeKind::AccountLeaf { .. } => { - // Account leaf at address level - handle potential storage overlays - self.handle_account_leaf_with_overlay( + NodeKind::Branch { .. } => { + self.traverse_branch_node_with_overlay( context, overlay, hash_builder, slotted_page, - Some(node), + node, full_path, )?; } - NodeKind::StorageLeaf { .. } => { - // Storage leaf - handle with simple leaf logic - // This can be either 128-nibble paths (account-level traversal) or 64-nibble paths - // (storage-level traversal) - if let Some(overlay_value) = overlay.lookup(full_path) { - // Use overlay value - if let Some(trie_value) = overlay_value { - let rlp_encoded = self.encode_trie_value(trie_value)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - } - // If overlay_value is None, it's a deletion - skip - } else { - // No exact overlay match, use disk value - let node_value = node.value()?; - let rlp_encoded = self.encode_trie_value(&node_value)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - } - } - NodeKind::Branch { .. } => { - // Branch node - recurse into each child with overlay sub-slicing - self.traverse_branch_node_with_overlay( + NodeKind::AccountLeaf { .. } => { + // Account leaf with descendant storage overlays + self.traverse_account_leaf_with_overlayed_storage( context, overlay, hash_builder, slotted_page, - node, + Some(node), full_path, )?; } - } - Ok(()) - } - - /// Handle a node that is not affected by overlay changes - fn handle_unaffected_node( - &self, - _context: &TransactionContext, - hash_builder: &mut HashBuilder, - node: &Node, - path_prefix: &Nibbles, - full_path: &Nibbles, - node_hash: Option, - ) -> Result<(), Error> { - match node.kind() { - NodeKind::AccountLeaf { .. } | NodeKind::StorageLeaf { .. } => { - // Account leaf - use disk value directly - let node_value = node.value()?; - let rlp_encoded = self.encode_trie_value(&node_value)?; - // println!("Adding unaffected leaf: {:?}", full_path); - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - } - NodeKind::Branch { .. } => { - // Branch node - since it's unaffected by overlay, we can use its existing hash - // This is much more efficient than recursing into all children - // println!("Adding branch node: {:?}", path_prefix); - let hash = node_hash.unwrap(); - hash_builder.add_branch(path_prefix.clone(), hash, true); + NodeKind::StorageLeaf { .. } => { + panic!("Storage leaf with descendant overlay: {full_path:?}"); } } + Ok(()) } @@ -269,6 +277,15 @@ impl StorageEngine { // Create sub-slice of overlay for this child's subtree let child_overlay = overlay.sub_slice_for_prefix(&child_path); + if let Some((path, _)) = child_overlay.get(0) { + if path == child_path { + // Direct overlay - the overlay path exactly matches this node's path + // No need to traverse the child + self.process_overlay_state(context, &child_overlay, hash_builder)?; + continue; + } + } + if let Ok(Some(child_pointer)) = node.child(child_index) { // Child exists on disk - traverse it with overlay integration let child_hash = child_pointer.rlp().as_hash(); @@ -297,25 +314,15 @@ impl StorageEngine { )?; } } else { - // No child exists on disk, but check if overlay has leaves for this subtree - if !child_overlay.is_empty() { - // There are overlay changes that create new nodes in this subtree - // Add all overlay leaves for this child subtree - for (path, value) in child_overlay.iter() { - if let Some(trie_value) = value { - let rlp_encoded = self.encode_trie_value(trie_value)?; - hash_builder.add_leaf(path, &rlp_encoded); - } - // Tombstones (None values) are skipped - } - } + // No child exists on disk, process any matching overlay changes + self.process_overlay_state(context, &child_overlay, hash_builder)?; } } Ok(()) } /// Handle an account leaf with potential storage overlays using iterative HashBuilder approach - fn handle_account_leaf_with_overlay( + fn traverse_account_leaf_with_overlayed_storage( &self, context: &TransactionContext, overlay: &OverlayState, @@ -338,18 +345,18 @@ impl StorageEngine { if let Some(trie_value) = overlay_value { // Even with account overlay, we might need to compute storage root for storage // overlays - if let TrieValue::Account(overlay_account) = trie_value { - Some(overlay_account.clone()) + if let OverlayValue::Account(overlay_account) = trie_value { + Some(overlay_account) } else { panic!("Overlay value is not an account"); } } else { // If overlay_value is None, it's a deletion - skip - println!("Skipping deletion: {full_path:?}"); + // println!("Skipping deletion: {full_path:?}"); None } } else { - original_account + original_account.as_ref() }; if let Some(account) = account { @@ -360,8 +367,8 @@ impl StorageEngine { if !has_storage_overlays { // No storage overlays for this account, use original account - let rlp_encoded = self.encode_trie_value(&TrieValue::Account(account))?; - println!("Adding account leaf: {full_path:?}"); + let rlp_encoded = self.encode_account(account)?; + // println!("Adding account leaf: {full_path:?}"); hash_builder.add_leaf(full_path.clone(), &rlp_encoded); return Ok(()); } @@ -374,12 +381,8 @@ impl StorageEngine { slotted_page, )?; - // Create a modified account with the new storage root - let mut modified_account = account; - modified_account.storage_root = new_storage_root; - // Add the modified account to the main HashBuilder - let rlp_encoded = self.encode_trie_value(&TrieValue::Account(modified_account))?; + let rlp_encoded = self.encode_account_with_root(account, new_storage_root)?; hash_builder.add_leaf(full_path.clone(), &rlp_encoded); } Ok(()) @@ -391,58 +394,33 @@ impl StorageEngine { overlay: &OverlayState, hash_builder: &mut HashBuilder, full_path: &Nibbles, + account: &Account, ) -> Result<(), Error> { - // Check if the overlay is an account leaf or None - let account = if let Some(overlay_value) = overlay.lookup(full_path) { - // Direct account overlay - use overlay value, but still check for storage overlays - if let Some(trie_value) = overlay_value { - // Even with account overlay, we might need to compute storage root for storage - // overlays - if let TrieValue::Account(overlay_account) = trie_value { - Some(overlay_account.clone()) - } else { - panic!("Overlay value is not an account"); - } - } else { - // If overlay_value is None, it's a deletion - skip - println!("Skipping deletion: {full_path:?}"); - None - } - } else { - None - }; - - if let Some(account) = account { - // Check if there are any storage overlays for this account - // Storage overlays have 128-nibble paths that start with this account's 64-nibble path - let storage_overlays = overlay.find_prefix_range(full_path); - let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() == 128); - - if !has_storage_overlays { - // No storage overlays for this account, use original account - let rlp_encoded = self.encode_trie_value(&TrieValue::Account(account))?; - println!("Adding account leaf: {full_path:?}"); - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - return Ok(()); - } + // Check if there are any storage overlays for this account + // Storage overlays have 128-nibble paths that start with this account's 64-nibble path + let storage_overlays = overlay.find_prefix_range(full_path); + let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() == 128); + + if !has_storage_overlays { + // No storage overlays for this account, use original account + let rlp_encoded = self.encode_account(account)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + return Ok(()); + } - // We have storage overlays, need to compute a new storage root - let storage_overlay = storage_overlays.with_prefix_offset(64); + // We have storage overlays, need to compute a new storage root + let storage_overlay = storage_overlays.with_prefix_offset(64); - let mut storage_hash_builder = HashBuilder::default().with_updates(true); + let mut storage_hash_builder = HashBuilder::default().with_updates(true); - self.process_overlay_state(context, &storage_overlay, &mut storage_hash_builder)?; + self.process_overlay_state(context, &storage_overlay, &mut storage_hash_builder)?; - let new_storage_root = storage_hash_builder.root(); + let new_storage_root = storage_hash_builder.root(); - // Create a modified account with the new storage root - let mut modified_account = account; - modified_account.storage_root = new_storage_root; + // Add the modified account to the main HashBuilder + let rlp_encoded = self.encode_account_with_root(account, new_storage_root)?; + hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - // Add the modified account to the main HashBuilder - let rlp_encoded = self.encode_trie_value(&TrieValue::Account(modified_account))?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); - } Ok(()) } @@ -457,7 +435,11 @@ impl StorageEngine { ) -> Result { // Create a storage-specific overlay with 64-nibble prefix offset // This converts 128-nibble storage paths to 64-nibble storage-relative paths - let storage_overlay = storage_overlays.with_prefix_offset(64); + let mut storage_overlay = storage_overlays.with_prefix_offset(64); + if let Some((_, Some(OverlayValue::Account(_)))) = storage_overlay.get(0) { + // println!("Skipping account overlay: {:?}", path); + storage_overlay = storage_overlay.sub_slice(1, storage_overlay.len()); + } let mut storage_hash_builder = HashBuilder::default().with_updates(true); @@ -506,30 +488,50 @@ impl StorageEngine { self.process_overlay_state(context, &storage_overlay, &mut storage_hash_builder)?; } - let (mut storage_hash_builder, updated_branch_nodes) = storage_hash_builder.split(); - println!("Updated storage branch nodes: {updated_branch_nodes:?}"); + let (mut storage_hash_builder, _updated_branch_nodes) = storage_hash_builder.split(); + // println!("Updated storage branch nodes: {updated_branch_nodes:?}"); let computed_root = storage_hash_builder.root(); - println!("Computed storage root: {computed_root:?}"); + // println!("Computed storage root: {computed_root:?}"); Ok(computed_root) } /// Helper to encode TrieValue as RLP + #[inline] fn encode_trie_value(&self, trie_value: &TrieValue) -> Result, Error> { let rlp_encoded = match trie_value { - TrieValue::Account(account) => { - let trie_account = TrieAccount { - nonce: account.nonce, - balance: account.balance, - storage_root: account.storage_root, - code_hash: account.code_hash, - }; - encode(trie_account) - } - TrieValue::Storage(storage_value) => encode(storage_value), + TrieValue::Account(account) => self.encode_account(account)?, + TrieValue::Storage(storage_value) => self.encode_storage(storage_value)?, }; Ok(rlp_encoded) } + + #[inline] + fn encode_account(&self, account: &Account) -> Result, Error> { + let trie_account = TrieAccount { + nonce: account.nonce, + balance: account.balance, + storage_root: account.storage_root, + code_hash: account.code_hash, + }; + Ok(encode(trie_account)) + } + + #[inline] + fn encode_account_with_root(&self, account: &Account, root: B256) -> Result, Error> { + let trie_account = TrieAccount { + nonce: account.nonce, + balance: account.balance, + storage_root: root, + code_hash: account.code_hash, + }; + Ok(encode(trie_account)) + } + + #[inline] + fn encode_storage(&self, storage_value: &U256) -> Result, Error> { + Ok(encode(storage_value)) + } } #[cfg(test)] @@ -553,7 +555,8 @@ mod tests { let context = db.storage_engine.read_context(); let empty_overlay = OverlayStateMut::new().freeze(); - let root = db.storage_engine.compute_root_with_overlay(&context, &empty_overlay).unwrap(); + let (root, _) = + db.storage_engine.compute_root_with_overlay(&context, &empty_overlay).unwrap(); assert_eq!(root, context.root_node_hash); } @@ -571,11 +574,11 @@ mod tests { 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(TrieValue::Account(account))); + overlay_mut.insert(address_path.into(), Some(OverlayValue::Account(account))); let overlay = overlay_mut.freeze(); // Compute root with overlay - let root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (root, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); // The root should be different from the empty root (since we have changes) assert_ne!(root, EMPTY_ROOT_HASH); @@ -606,7 +609,7 @@ mod tests { // Test that we can compute root with empty overlay first (should match initial_root) let empty_overlay = OverlayStateMut::new().freeze(); - let root_with_empty_overlay = + let (root_with_empty_overlay, _) = db.storage_engine.compute_root_with_overlay(&context, &empty_overlay).unwrap(); assert_eq!(root_with_empty_overlay, initial_root); @@ -614,10 +617,12 @@ mod tests { 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(TrieValue::Account(account2.clone()))); + overlay_mut + .insert(address_path.clone().into(), Some(OverlayValue::Account(account2.clone()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify: commit the overlay changes and compare roots @@ -682,10 +687,11 @@ mod tests { 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(TrieValue::Account(account1_updated.clone()))); + .insert(path1.clone().into(), Some(OverlayValue::Account(account1_updated.clone()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing the change @@ -715,10 +721,10 @@ mod tests { let account3 = Account::new(3, U256::from(300), EMPTY_ROOT_HASH, KECCAK_EMPTY); let mut overlay_mut2 = OverlayStateMut::new(); - overlay_mut2.insert(path3.clone().into(), Some(TrieValue::Account(account3.clone()))); + overlay_mut2.insert(path3.clone().into(), Some(OverlayValue::Account(account3.clone()))); let overlay2 = overlay_mut2.freeze(); - let overlay_root2 = + let (overlay_root2, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay2).unwrap(); assert_ne!(overlay_root2, initial_root); assert_ne!(overlay_root2, overlay_root); @@ -744,12 +750,12 @@ mod tests { // unaffected let mut overlay_mut3 = OverlayStateMut::new(); overlay_mut3 - .insert(path1.clone().into(), Some(TrieValue::Account(account1_updated.clone()))); // Modify existing - overlay_mut3.insert(path3.clone().into(), Some(TrieValue::Account(account3.clone()))); // Create new - // path2 is left unaffected - should use add_branch optimization + .insert(path1.clone().into(), Some(OverlayValue::Account(account1_updated.clone()))); // Modify existing + overlay_mut3.insert(path3.clone().into(), Some(OverlayValue::Account(account3.clone()))); // Create new + // path2 is left unaffected - should use add_branch optimization let overlay3 = overlay_mut3.freeze(); - let overlay_root3 = + let (overlay_root3, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay3).unwrap(); assert_ne!(overlay_root3, initial_root); assert_ne!(overlay_root3, overlay_root); @@ -824,7 +830,7 @@ mod tests { overlay_mut.insert(path2.clone().into(), None); // Delete account2 let overlay = overlay_mut.freeze(); - let overlay_root_with_deletion = + let (overlay_root_with_deletion, _) = db.storage_engine.compute_root_with_overlay(&context2, &overlay).unwrap(); assert_ne!(overlay_root_with_deletion, root_with_both_accounts); @@ -900,11 +906,14 @@ mod tests { // 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(TrieValue::Storage(new_storage_value1.into()))); + overlay_mut.insert( + storage_path1.full_path(), + Some(OverlayValue::Storage(new_storage_value1.into())), + ); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing the storage change @@ -936,10 +945,10 @@ mod tests { 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(TrieValue::Storage(storage_value3.into()))); + .insert(storage_path3.full_path(), Some(OverlayValue::Storage(storage_value3.into()))); let overlay2 = overlay_mut2.freeze(); - let overlay_root2 = + let (overlay_root2, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay2).unwrap(); assert_ne!(overlay_root2, initial_root); assert_ne!(overlay_root2, overlay_root); @@ -969,7 +978,7 @@ mod tests { overlay_mut3.insert(storage_path2.full_path(), None); // Delete storage slot let overlay3 = overlay_mut3.freeze(); - let overlay_root3 = + let (overlay_root3, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay3).unwrap(); assert_ne!(overlay_root3, initial_root); @@ -995,13 +1004,17 @@ mod tests { // 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(TrieValue::Account(updated_account.clone()))); - overlay_mut4 - .insert(storage_path1.full_path(), Some(TrieValue::Storage(new_storage_value1.into()))); + 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.into())), + ); let overlay4 = overlay_mut4.freeze(); - let overlay_root4 = + let (overlay_root4, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay4).unwrap(); assert_ne!(overlay_root4, initial_root); @@ -1047,11 +1060,14 @@ mod tests { let updated_account = Account::new(10, U256::from(1000), EMPTY_ROOT_HASH, KECCAK_EMPTY); let mut overlay_mut = OverlayStateMut::new(); - overlay_mut - .insert(modify_path.clone().into(), Some(TrieValue::Account(updated_account.clone()))); + overlay_mut.insert( + modify_path.clone().into(), + Some(OverlayValue::Account(updated_account.clone())), + ); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify the computation is correct by comparing with direct modification @@ -1097,6 +1113,8 @@ mod tests { let storage_path1 = crate::path::StoragePath::for_address_and_slot(account_address, storage_key1.into()); + // println!("Storage path 1: {:?}", storage_path1.full_path()); + // Set up initial state with 1 storage slot db.storage_engine .set_values( @@ -1115,11 +1133,15 @@ mod tests { let storage_key2 = U256::from(20); // New storage key let storage_path2 = crate::path::StoragePath::for_address_and_slot(account_address, storage_key2.into()); + + // println!("Storage path 2: {:?}", storage_path2.full_path()); + overlay_mut - .insert(storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))); + .insert(storage_path2.full_path(), Some(OverlayValue::Storage(U256::from(222).into()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing the addition @@ -1178,12 +1200,13 @@ mod tests { // Test: Modify just one storage value per account via overlay let mut overlay_mut = OverlayStateMut::new(); overlay_mut - .insert(storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))); + .insert(storage1_path.full_path(), Some(OverlayValue::Storage(U256::from(999).into()))); overlay_mut - .insert(storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888).into()))); + .insert(storage2_path.full_path(), Some(OverlayValue::Storage(U256::from(888).into()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing the changes @@ -1245,30 +1268,32 @@ mod tests { let mut overlay_mut = OverlayStateMut::new(); // Modify existing storage slot 1 - overlay_mut - .insert(storage_path1.full_path(), Some(TrieValue::Storage(U256::from(1111).into()))); + overlay_mut.insert( + storage_path1.full_path(), + Some(OverlayValue::Storage(U256::from(1111).into())), + ); // 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(TrieValue::Storage(U256::from(444).into()))); + .insert(storage_path3.full_path(), Some(OverlayValue::Storage(U256::from(444).into()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing all changes - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine.set_values(&mut verification_context, &mut [ + db.storage_engine.set_values(&mut context, &mut [ (account_path.into(), Some(TrieValue::Account(account))), (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(1111).into()))), // Modified (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), // Unchanged (storage_path3.full_path(), Some(TrieValue::Storage(U256::from(444).into()))), // New ]).unwrap(); - let committed_root = verification_context.root_node_hash; + let committed_root = context.root_node_hash; assert_eq!( overlay_root, committed_root, @@ -1313,10 +1338,11 @@ mod tests { // 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(TrieValue::Storage(U256::from(999).into()))); + .insert(storage_path1.full_path(), Some(OverlayValue::Storage(U256::from(999).into()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing: modify slot1, keep slot2 unchanged @@ -1385,23 +1411,30 @@ mod tests { let mut overlay_mut = OverlayStateMut::new(); // Modify account1's storage - overlay_mut - .insert(storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(1111).into()))); + overlay_mut.insert( + storage1_path1.full_path(), + Some(OverlayValue::Storage(U256::from(1111).into())), + ); // 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(TrieValue::Storage(U256::from(444).into()))); + overlay_mut.insert( + storage1_path3.full_path(), + Some(OverlayValue::Storage(U256::from(444).into())), + ); // Modify account2's storage - overlay_mut - .insert(storage2_path1.full_path(), Some(TrieValue::Storage(U256::from(3333).into()))); + overlay_mut.insert( + storage2_path1.full_path(), + Some(OverlayValue::Storage(U256::from(3333).into())), + ); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing all changes @@ -1471,10 +1504,11 @@ mod tests { // Test: Modify just one storage slot for account1 let mut overlay_mut = OverlayStateMut::new(); overlay_mut - .insert(storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))); + .insert(storage1_path.full_path(), Some(OverlayValue::Storage(U256::from(999).into()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing the change @@ -1537,12 +1571,13 @@ mod tests { // Test: Modify storage for BOTH accounts let mut overlay_mut = OverlayStateMut::new(); overlay_mut - .insert(storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))); + .insert(storage1_path.full_path(), Some(OverlayValue::Storage(U256::from(999).into()))); overlay_mut - .insert(storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888).into()))); + .insert(storage2_path.full_path(), Some(OverlayValue::Storage(U256::from(888).into()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing both changes @@ -1604,11 +1639,14 @@ mod tests { let storage1_path2 = crate::path::StoragePath::for_address_and_slot(account1_address, storage1_key2.into()); - overlay_mut - .insert(storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(444).into()))); + overlay_mut.insert( + storage1_path2.full_path(), + Some(OverlayValue::Storage(U256::from(444).into())), + ); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing the addition @@ -1654,10 +1692,11 @@ mod tests { 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(TrieValue::Storage(storage_value.into()))); + .insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value.into()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing the storage addition @@ -1715,7 +1754,7 @@ mod tests { // 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(TrieValue::Account(new_account))); + overlay_mut.insert(path.clone(), Some(OverlayValue::Account(new_account))); } else { // For other half, mark for deletion overlay_mut.insert(path.clone(), None); @@ -1723,18 +1762,49 @@ mod tests { } } + let mut overlay_mut_with_branches = overlay_mut.clone(); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, overlay_branches) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); + // println!("Overlay branches: {:?}", overlay_branches); + + for (path, branch) in overlay_branches.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 as u8); + overlay_mut_with_branches + .insert(path.clone(), Some(OverlayValue::Hash(branch.hashes[hash_idx]))); + hash_idx += 1; + path.pop(); + } + } + } + let overlay_with_branches = overlay_mut_with_branches.freeze(); + // println!("Overlay with branches: {:?}", overlay_with_branches); + let (overlay_root_with_branches, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay_with_branches).unwrap(); + assert_eq!(overlay_root_with_branches, overlay_root); + // Verify by committing the storage addition - db.storage_engine.set_values(&mut context, overlay.data().to_vec().as_mut()).unwrap(); + 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(&mut context, &mut changes).unwrap(); let committed_root = context.root_node_hash; - db.storage_engine - .print_page(&context, std::io::stdout(), context.root_node_page_id) - .unwrap(); + // db.storage_engine + // .print_page(&context, std::io::stdout(), context.root_node_page_id) + // .unwrap(); assert_eq!(overlay_root, committed_root, "1000 accounts with 10 overlay should match"); } @@ -1764,26 +1834,33 @@ mod tests { 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.into(), Some(TrieValue::Account(new_account.clone()))); + overlay_mut + .insert(new_account_path.into(), Some(OverlayValue::Account(new_account.clone()))); let storage_key = U256::from(42); let storage_path = crate::path::StoragePath::for_address_and_slot(new_address, storage_key.into()); let storage_value = U256::from(123); overlay_mut - .insert(storage_path.full_path(), Some(TrieValue::Storage(storage_value.into()))); + .insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value.into()))); let overlay = overlay_mut.freeze(); - let overlay_root = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); // Verify by committing the storage addition - db.storage_engine.set_values(&mut context, overlay.data().to_vec().as_mut()).unwrap(); + 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(&mut context, &mut changes).unwrap(); let committed_root = context.root_node_hash; - db.storage_engine - .print_page(&context, std::io::stdout(), context.root_node_page_id) - .unwrap(); + // db.storage_engine + // .print_page(&context, std::io::stdout(), context.root_node_page_id) + // .unwrap(); assert_eq!(overlay_root, committed_root, "Overlay new account with storage should match"); } diff --git a/src/transaction.rs b/src/transaction.rs index 778bf9bc..5005bb88 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -6,7 +6,6 @@ use crate::{ context::TransactionContext, database::Database, node::TrieValue, - overlay::OverlayStateMut, path::{AddressPath, StoragePath}, storage::proofs::AccountProof, }; @@ -47,7 +46,6 @@ pub struct Transaction { context: TransactionContext, database: DB, pending_changes: HashMap>, - overlay_state: Option, _marker: std::marker::PhantomData, } @@ -58,7 +56,6 @@ impl, K: TransactionKind> Transaction { context, database, pending_changes: HashMap::new(), - overlay_state: None, _marker: std::marker::PhantomData, } } From ae02d0165abd759d70e4aabc3e05c324b5cb38ab Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Mon, 28 Jul 2025 16:43:02 -0700 Subject: [PATCH 05/17] Lint --- src/overlay.rs | 102 ++++++++++------------ src/storage/overlay_root.rs | 166 +++++++++++++++--------------------- 2 files changed, 112 insertions(+), 156 deletions(-) diff --git a/src/overlay.rs b/src/overlay.rs index b6dc0910..4f8dd442 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -766,14 +766,13 @@ mod tests { let account = test_account(); // Create overlay with account-prefixed storage paths for a SINGLE account - let _account_prefix = vec![1, 0, 0, 0]; // 4-nibble account prefix + let _account_prefix = [1, 0, 0, 0]; // 4-nibble account prefix let paths = [ - vec![1, 0, 0, 0, 2, 0], // account + storage [2, 0] - vec![1, 0, 0, 0, 5, 0], // account + storage [5, 0] - vec![1, 0, 0, 0, 8, 0], // account + storage [8, 0] - vec![1, 0, 0, 0, 9, 0], /* account + storage [9, 0] - changed from different account - * to same account */ + [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 { @@ -810,8 +809,8 @@ mod tests { let account = test_account(); // Create storage overlays for two different accounts - let account1_prefix = vec![1, 0, 0, 0]; // Account 1: 4 nibbles - let _account2_prefix = vec![2, 0, 0, 0]; // Account 2: 4 nibbles + 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( @@ -832,8 +831,7 @@ mod tests { let frozen = mutable.freeze(); // Test finding storage for account 1 using find_prefix_range on original overlay - let account1_storage = - frozen.find_prefix_range(&Nibbles::from_nibbles(account1_prefix.clone())); + let account1_storage = frozen.find_prefix_range(&Nibbles::from_nibbles(account1_prefix)); assert_eq!(account1_storage.len(), 2); // Now test with prefix_offset - should convert to storage-relative paths @@ -892,7 +890,7 @@ mod tests { #[test] fn test_thread_safety_with_arc() { - use std::{sync::Arc as StdArc, thread}; + use std::thread; let mut mutable = OverlayStateMut::new(); let account = test_account(); @@ -906,12 +904,12 @@ mod tests { } let frozen = mutable.freeze(); - let shared_overlay = StdArc::new(frozen); + let shared_overlay = Arc::new(frozen); // Spawn multiple threads that work with different sub-slices let handles: Vec<_> = (0..4) .map(|i| { - let overlay = StdArc::clone(&shared_overlay); + let overlay = Arc::clone(&shared_overlay); thread::spawn(move || { let start = i * 25; let end = (i + 1) * 25; @@ -941,7 +939,7 @@ mod tests { // Create a diverse set of paths to test prefix partitioning // Use vectors instead of arrays to avoid size mismatch issues - let test_paths = vec![ + 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()))), @@ -963,7 +961,7 @@ mod tests { } let frozen = mutable.freeze(); - let prefix = Nibbles::from_nibbles(vec![1, 2]); + let prefix = Nibbles::from_nibbles([1, 2]); // Test sub_slice_by_prefix let (before, with_prefix, after) = frozen.sub_slice_by_prefix(&prefix); @@ -971,24 +969,24 @@ mod tests { // 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(&Nibbles::from_nibbles(vec![0, 5, 6, 7]))); - assert!(before_paths.contains(&Nibbles::from_nibbles(vec![1, 0, 0, 0]))); - assert!(before_paths.contains(&Nibbles::from_nibbles(vec![1, 1, 9, 9]))); + assert!(before_paths.contains(&Nibbles::from_nibbles([0, 5, 6, 7]))); + assert!(before_paths.contains(&Nibbles::from_nibbles([1, 0, 0, 0]))); + assert!(before_paths.contains(&Nibbles::from_nibbles([1, 1, 9, 9]))); // 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(&Nibbles::from_nibbles(vec![1, 2]))); - assert!(with_prefix_paths.contains(&Nibbles::from_nibbles(vec![1, 2, 3, 4]))); - assert!(with_prefix_paths.contains(&Nibbles::from_nibbles(vec![1, 2, 5, 6]))); - assert!(with_prefix_paths.contains(&Nibbles::from_nibbles(vec![1, 2, 7, 8]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([1, 2]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([1, 2, 3, 4]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([1, 2, 5, 6]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([1, 2, 7, 8]))); // 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(&Nibbles::from_nibbles(vec![1, 3, 0, 0]))); - assert!(after_paths.contains(&Nibbles::from_nibbles(vec![2, 0, 0, 0]))); - assert!(after_paths.contains(&Nibbles::from_nibbles(vec![5, 5, 5, 5]))); + assert!(after_paths.contains(&Nibbles::from_nibbles([1, 3, 0, 0]))); + assert!(after_paths.contains(&Nibbles::from_nibbles([2, 0, 0, 0]))); + assert!(after_paths.contains(&Nibbles::from_nibbles([5, 5, 5, 5]))); // Verify that all slices together contain all original entries assert_eq!(before.len() + with_prefix.len() + after.len(), frozen.len()); @@ -1001,14 +999,14 @@ mod tests { assert!(empty_after.is_empty()); // Test with prefix that doesn't match anything - let no_match_prefix = Nibbles::from_nibbles(vec![9, 9, 9]); + 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(vec![0xF, 0xF]); + 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()); @@ -1022,25 +1020,13 @@ mod tests { let account = test_account(); // Test case where we have exact prefix matches and no extensions - let prefix = Nibbles::from_nibbles(vec![2, 3]); + let prefix = Nibbles::from_nibbles([2, 3]); - mutable.insert( - Nibbles::from_nibbles(vec![1, 9]), - Some(OverlayValue::Account(account.clone())), - ); - mutable.insert( - Nibbles::from_nibbles(vec![2, 2]), - Some(OverlayValue::Account(account.clone())), - ); + 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(vec![2, 4]), - Some(OverlayValue::Account(account.clone())), - ); - mutable.insert( - Nibbles::from_nibbles(vec![3, 0]), - Some(OverlayValue::Account(account.clone())), - ); + 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); @@ -1060,21 +1046,19 @@ mod tests { let account = test_account(); // Create paths that simulate account (4 nibbles) + storage (4 nibbles) patterns - let _account_prefix = vec![1, 0, 0, 0]; + let _account_prefix = [1, 0, 0, 0]; // Storage paths for the same account - let storage_paths = vec![ - vec![1, 0, 0, 0, 2, 0, 0, 0], // account + storage [2,0,0,0] - vec![1, 0, 0, 0, 5, 0, 0, 0], // account + storage [5,0,0,0] - vec![1, 0, 0, 0, 5, 1, 0, 0], // account + storage [5,1,0,0] - vec![1, 0, 0, 0, 8, 0, 0, 0], // account + storage [8,0,0,0] + 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.clone()), - Some(OverlayValue::Account(account.clone())), - ); + mutable + .insert(Nibbles::from_nibbles(path), Some(OverlayValue::Account(account.clone()))); } let frozen = mutable.freeze(); @@ -1083,23 +1067,23 @@ mod tests { let storage_overlay = frozen.with_prefix_offset(4); // Test partitioning by storage prefix [5] - let storage_prefix = Nibbles::from_nibbles(vec![5]); + let storage_prefix = Nibbles::from_nibbles([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(&Nibbles::from_nibbles(vec![2, 0, 0, 0]))); + assert!(before_paths.contains(&Nibbles::from_nibbles([2, 0, 0, 0]))); // 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(&Nibbles::from_nibbles(vec![5, 0, 0, 0]))); - assert!(with_prefix_paths.contains(&Nibbles::from_nibbles(vec![5, 1, 0, 0]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([5, 0, 0, 0]))); + assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([5, 1, 0, 0]))); // 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(&Nibbles::from_nibbles(vec![8, 0, 0, 0]))); + assert!(after_paths.contains(&Nibbles::from_nibbles([8, 0, 0, 0]))); } } diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 413818ab..83824bec 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -894,8 +894,8 @@ mod tests { &mut context, &mut [ (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), - (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1.into()))), - (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2.into()))), + (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1))), + (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2))), ], ) .unwrap(); @@ -906,10 +906,8 @@ mod tests { // 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.into())), - ); + overlay_mut + .insert(storage_path1.full_path(), Some(OverlayValue::Storage(new_storage_value1))); let overlay = overlay_mut.freeze(); let (overlay_root, _) = @@ -923,11 +921,8 @@ mod tests { &mut verification_context, &mut [ (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), - ( - storage_path1.full_path(), - Some(TrieValue::Storage(new_storage_value1.into())), - ), - (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2.into()))), + (storage_path1.full_path(), Some(TrieValue::Storage(new_storage_value1))), + (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2))), ], ) .unwrap(); @@ -944,8 +939,7 @@ mod tests { 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.into()))); + overlay_mut2.insert(storage_path3.full_path(), Some(OverlayValue::Storage(storage_value3))); let overlay2 = overlay_mut2.freeze(); let (overlay_root2, _) = @@ -960,9 +954,9 @@ mod tests { &mut verification_context2, &mut [ (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), - (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1.into()))), - (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2.into()))), - (storage_path3.full_path(), Some(TrieValue::Storage(storage_value3.into()))), + (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1))), + (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2))), + (storage_path3.full_path(), Some(TrieValue::Storage(storage_value3))), ], ) .unwrap(); @@ -989,7 +983,7 @@ mod tests { &mut verification_context3, &mut [ (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), - (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1.into()))), + (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1))), // storage_path2 is omitted - effectively deleted ], ) @@ -1008,10 +1002,8 @@ mod tests { account_path.clone().into(), Some(OverlayValue::Account(updated_account.clone())), ); - overlay_mut4.insert( - storage_path1.full_path(), - Some(OverlayValue::Storage(new_storage_value1.into())), - ); + overlay_mut4 + .insert(storage_path1.full_path(), Some(OverlayValue::Storage(new_storage_value1))); let overlay4 = overlay_mut4.freeze(); let (overlay_root4, _) = @@ -1121,7 +1113,7 @@ mod tests { &mut context, &mut [ (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), - (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111).into()))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), ], ) .unwrap(); @@ -1136,8 +1128,7 @@ mod tests { // println!("Storage path 2: {:?}", storage_path2.full_path()); - overlay_mut - .insert(storage_path2.full_path(), Some(OverlayValue::Storage(U256::from(222).into()))); + overlay_mut.insert(storage_path2.full_path(), Some(OverlayValue::Storage(U256::from(222)))); let overlay = overlay_mut.freeze(); let (overlay_root, _) = @@ -1148,8 +1139,8 @@ mod tests { let mut verification_context = db.storage_engine.write_context(); db.storage_engine.set_values(&mut verification_context, &mut [ (account_path.into(), Some(TrieValue::Account(account))), - (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111).into()))), // Original - (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), // New + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), // Original + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), // New ]).unwrap(); let committed_root = verification_context.root_node_hash; @@ -1189,8 +1180,8 @@ mod tests { &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).into()))), - (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222)))), ], ) .unwrap(); @@ -1199,10 +1190,8 @@ mod tests { // 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).into()))); - overlay_mut - .insert(storage2_path.full_path(), Some(OverlayValue::Storage(U256::from(888).into()))); + 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(); let (overlay_root, _) = @@ -1217,8 +1206,8 @@ mod tests { &mut [ (account1_path.into(), Some(TrieValue::Account(account1))), (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))), - (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888).into()))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999)))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888)))), ], ) .unwrap(); @@ -1256,8 +1245,8 @@ mod tests { &mut context, &mut [ (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), - (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111).into()))), - (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), ], ) .unwrap(); @@ -1268,17 +1257,14 @@ mod tests { let mut overlay_mut = OverlayStateMut::new(); // Modify existing storage slot 1 - overlay_mut.insert( - storage_path1.full_path(), - Some(OverlayValue::Storage(U256::from(1111).into())), - ); + 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).into()))); + overlay_mut.insert(storage_path3.full_path(), Some(OverlayValue::Storage(U256::from(444)))); let overlay = overlay_mut.freeze(); @@ -1289,9 +1275,9 @@ mod tests { // Verify by committing all changes db.storage_engine.set_values(&mut context, &mut [ (account_path.into(), Some(TrieValue::Account(account))), - (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(1111).into()))), // Modified - (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), // Unchanged - (storage_path3.full_path(), Some(TrieValue::Storage(U256::from(444).into()))), // New + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(1111)))), // Modified + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), // Unchanged + (storage_path3.full_path(), Some(TrieValue::Storage(U256::from(444)))), // New ]).unwrap(); let committed_root = context.root_node_hash; @@ -1327,8 +1313,8 @@ mod tests { &mut context, &mut [ (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), - (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111).into()))), - (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), ], ) .unwrap(); @@ -1337,8 +1323,7 @@ mod tests { // 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).into()))); + overlay_mut.insert(storage_path1.full_path(), Some(OverlayValue::Storage(U256::from(999)))); let overlay = overlay_mut.freeze(); let (overlay_root, _) = @@ -1349,8 +1334,8 @@ mod tests { let mut verification_context = db.storage_engine.write_context(); db.storage_engine.set_values(&mut verification_context, &mut [ (account_path.into(), Some(TrieValue::Account(account))), - (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(999).into()))), // Modified - (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), // Unchanged + (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(999)))), // Modified + (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), // Unchanged ]).unwrap(); let committed_root = verification_context.root_node_hash; @@ -1398,9 +1383,9 @@ mod tests { &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).into()))), - (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), - (storage2_path1.full_path(), Some(TrieValue::Storage(U256::from(333).into()))), + (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(); @@ -1411,25 +1396,19 @@ mod tests { let mut overlay_mut = OverlayStateMut::new(); // Modify account1's storage - overlay_mut.insert( - storage1_path1.full_path(), - Some(OverlayValue::Storage(U256::from(1111).into())), - ); + 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).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).into())), - ); + overlay_mut + .insert(storage2_path1.full_path(), Some(OverlayValue::Storage(U256::from(3333)))); let overlay = overlay_mut.freeze(); @@ -1445,10 +1424,10 @@ mod tests { &mut [ (account1_path.into(), Some(TrieValue::Account(account1))), (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(1111).into()))), - (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), - (storage1_path3.full_path(), Some(TrieValue::Storage(U256::from(444).into()))), - (storage2_path1.full_path(), Some(TrieValue::Storage(U256::from(3333).into()))), + (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(1111)))), + (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), + (storage1_path3.full_path(), Some(TrieValue::Storage(U256::from(444)))), + (storage2_path1.full_path(), Some(TrieValue::Storage(U256::from(3333)))), ], ) .unwrap(); @@ -1493,8 +1472,8 @@ mod tests { &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).into()))), - (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222)))), ], ) .unwrap(); @@ -1503,8 +1482,7 @@ mod tests { // 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).into()))); + overlay_mut.insert(storage1_path.full_path(), Some(OverlayValue::Storage(U256::from(999)))); let overlay = overlay_mut.freeze(); let (overlay_root, _) = @@ -1516,8 +1494,8 @@ mod tests { db.storage_engine.set_values(&mut verification_context, &mut [ (account1_path.into(), Some(TrieValue::Account(account1))), (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))), // Modified - (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), // Unchanged + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999)))), // Modified + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222)))), // Unchanged ]).unwrap(); let committed_root = verification_context.root_node_hash; @@ -1560,8 +1538,8 @@ mod tests { &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).into()))), - (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222).into()))), + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(111)))), + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222)))), ], ) .unwrap(); @@ -1570,10 +1548,8 @@ mod tests { // 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).into()))); - overlay_mut - .insert(storage2_path.full_path(), Some(OverlayValue::Storage(U256::from(888).into()))); + 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(); let (overlay_root, _) = @@ -1585,8 +1561,8 @@ mod tests { db.storage_engine.set_values(&mut verification_context, &mut [ (account1_path.into(), Some(TrieValue::Account(account1))), (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999).into()))), // Modified - (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888).into()))), // Modified + (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999)))), // Modified + (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888)))), // Modified ]).unwrap(); let committed_root = verification_context.root_node_hash; @@ -1626,7 +1602,7 @@ mod tests { &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).into()))), + (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), ], ) .unwrap(); @@ -1639,10 +1615,8 @@ mod tests { 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).into())), - ); + overlay_mut + .insert(storage1_path2.full_path(), Some(OverlayValue::Storage(U256::from(444)))); let overlay = overlay_mut.freeze(); let (overlay_root, _) = @@ -1654,8 +1628,8 @@ mod tests { db.storage_engine.set_values(&mut verification_context, &mut [ (account1_path.into(), Some(TrieValue::Account(account1))), (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(111).into()))), // Original - (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(444).into()))), // New + (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), // Original + (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(444)))), // New ]).unwrap(); let committed_root = verification_context.root_node_hash; @@ -1691,8 +1665,7 @@ mod tests { 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.into()))); + overlay_mut.insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value))); let overlay = overlay_mut.freeze(); let (overlay_root, _) = @@ -1706,7 +1679,7 @@ mod tests { &mut verification_context, &mut [ (account_path.into(), Some(TrieValue::Account(account))), - (storage_path.full_path(), Some(TrieValue::Storage(storage_value.into()))), + (storage_path.full_path(), Some(TrieValue::Storage(storage_value))), ], ) .unwrap(); @@ -1779,7 +1752,7 @@ mod tests { let mut path = path.clone(); for i in 0..16 { if branch.hash_mask.is_bit_set(i) { - path.push(i as u8); + path.push(i); overlay_mut_with_branches .insert(path.clone(), Some(OverlayValue::Hash(branch.hashes[hash_idx]))); hash_idx += 1; @@ -1841,8 +1814,7 @@ mod tests { let storage_path = crate::path::StoragePath::for_address_and_slot(new_address, storage_key.into()); let storage_value = U256::from(123); - overlay_mut - .insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value.into()))); + overlay_mut.insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value))); let overlay = overlay_mut.freeze(); let (overlay_root, _) = From 2631577a7f2040326ebf1860af6253f61c28993b Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Thu, 31 Jul 2025 11:21:49 -0700 Subject: [PATCH 06/17] WIP: stats/benchmarking --- Cargo.toml | 1 + benches/crud_benchmarks.rs | 137 ++++++++++++++- src/overlay.rs | 2 +- src/page/manager/mmap.rs | 2 +- src/storage/overlay_root.rs | 324 +++++++++++++++++++++++++++++------- src/transaction.rs | 20 ++- 6 files changed, 418 insertions(+), 68 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index ed0f05ca..056c2024 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,6 +27,7 @@ rand = "0.9.1" walkdir = "2" serde_json = "1.0.140" tempfile = "3.19.1" +test-log = { version = "0.2.18", features = ["trace", "color"] } [lib] bench = false diff --git a/benches/crud_benchmarks.rs b/benches/crud_benchmarks.rs index 450601d5..35e38b89 100644 --- a/benches/crud_benchmarks.rs +++ b/benches/crud_benchmarks.rs @@ -20,10 +20,11 @@ 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, hint::black_box, io, path::Path, sync::Arc, thread, time::Duration}; use tempdir::TempDir; use triedb::{ account::Account, + overlay::{OverlayStateMut, OverlayValue}, path::{AddressPath, StoragePath}, Database, }; @@ -82,6 +83,56 @@ 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); + Arc::new(Database::open(db_path.clone()).unwrap()) + }, + |db| { + let thread_count = 4; + // Divide addresses into 4 chunks + let addresses = addresses.clone(); + let chunk_size = addresses.len() / thread_count; + let mut handles = Vec::new(); + + // Spawn 4 threads + 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 handle = thread::spawn(move || { + // 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.clone()).unwrap(); + assert!(a.is_some()); + } + + // Commit the transaction + tx.commit().unwrap(); + }); + + handles.push(handle); + } + + // Wait for all threads to complete + for handle in handles { + handle.join().unwrap(); + } + }, + ); + }); group.finish(); } @@ -538,6 +589,89 @@ 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 file_name = base_dir.main_file_name.clone(); + + let mut rng = StdRng::seed_from_u64(SEED_CONTRACT); + // let total_storage_per_address = DEFAULT_SETUP_DB_STORAGE_PER_CONTRACT; + let total_addresses = BATCH_SIZE; + let addresses: Vec = + (0..total_addresses).map(|_| generate_random_address(&mut rng)).collect(); + // let storage_paths_values = generate_storage_paths_values(&addresses, total_storage_per_address); + + 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(); + + // Build overlay state from storage paths and values + // let mut storage_overlay_mut = OverlayStateMut::new(); + + // for (storage_path, storage_value) in &storage_paths_values { + // // Convert storage path to nibbles for overlay + // let nibbles = storage_path.full_path(); + // storage_overlay_mut.insert(nibbles, Some(OverlayValue::Storage(*storage_value))); + // } + + // // Freeze the mutable overlay to get an immutable one + // let storage_overlay = storage_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 dir = TempDir::new("triedb_bench_state_root_with_account_overlay").unwrap(); + copy_files(&base_dir, dir.path()).unwrap(); + 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).unwrap(); + + tx.commit().unwrap(); + }, + ); + }); + + // group.bench_function(BenchmarkId::new("state_root_with_storage_overlay", BATCH_SIZE), |b| { + // b.iter_with_setup( + // || { + // let dir = TempDir::new("triedb_bench_state_root_with_storage_overlay").unwrap(); + // copy_files(&base_dir, dir.path()).unwrap(); + // let db_path = dir.path().join(&file_name); + // Database::open(db_path).unwrap() + // }, + // |db| { + // black_box({ + // for _ in 0..100 { + // let tx = db.begin_ro().unwrap(); + + // // Compute the root hash with the overlay + // let _root_result = tx.compute_root_with_overlay(&storage_overlay).unwrap(); + + // tx.commit().unwrap(); + // } + // }) + // }, + // ); + // }); + + group.finish(); +} + criterion_group!( benches, bench_account_reads, @@ -550,5 +684,6 @@ criterion_group!( bench_storage_inserts, bench_storage_updates, bench_storage_deletes, + bench_state_root_with_overlay, ); criterion_main!(benches); diff --git a/src/overlay.rs b/src/overlay.rs index 4f8dd442..1b004a5f 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -181,7 +181,7 @@ impl OverlayState { // Find the range of entries that start with the prefix after applying offset for (i, (path, _)) in slice.iter().enumerate() { if path.len() >= self.prefix_offset { - let adjusted_path = path.slice(self.prefix_offset..); + let adjusted_path = &path.as_slice()[self.prefix_offset..]; if adjusted_path.len() >= prefix.len() && adjusted_path[..prefix.len()] == *prefix { if start_idx.is_none() { start_idx = Some(i); diff --git a/src/page/manager/mmap.rs b/src/page/manager/mmap.rs index 41d62087..b68e9a69 100644 --- a/src/page/manager/mmap.rs +++ b/src/page/manager/mmap.rs @@ -268,7 +268,7 @@ impl PageManager { /// Syncs pages to the backing file. pub fn sync(&self) -> io::Result<()> { if cfg!(not(miri)) { - self.mmap.flush() + self.mmap.flush_async() } else { Ok(()) } diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 83824bec..51d8f310 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -4,12 +4,36 @@ use crate::{ node::{Node, NodeKind, TrieValue}, overlay::{OverlayState, OverlayValue}, page::{PageId, SlottedPage}, + pointer::Pointer, storage::engine::{Error, StorageEngine}, }; use alloy_primitives::{map::HashMap, B256, U256}; use alloy_rlp::encode; use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, TrieAccount}; +#[derive(Debug, Default)] +struct Stats { + account: AccountStats, + storage: StorageStats, +} + +#[derive(Debug, Default)] +struct AccountStats { + pointer_hashes: usize, + branch_nodes: usize, + leaves: usize, + overlay_branch_nodes: usize, + overlay_leaves: usize, +} + +#[derive(Debug, Default)] +struct StorageStats { + branch_nodes: usize, + leaves: usize, + overlay_branch_nodes: usize, + overlay_leaves: usize, +} + impl StorageEngine { /// Computes the root hash with overlay changes without persisting them. /// This uses alloy-trie's HashBuilder for efficient merkle root computation. @@ -24,15 +48,18 @@ impl StorageEngine { } let mut hash_builder = HashBuilder::default().with_updates(true); + let mut stats = Stats::default(); // Use proper trie traversal with overlay integration - self.traverse_with_overlay(context, overlay, &mut hash_builder)?; + self.traverse_with_overlay(context, overlay, &mut hash_builder, &mut stats)?; let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); // println!("Updated branch nodes: {:?}", updated_branch_nodes); let root = hash_builder.root(); // This will clear the hash builder // println!("Root: {:?}", root); + + println!("Stats: {:?}", stats); Ok((root, updated_branch_nodes)) } @@ -42,6 +69,7 @@ impl StorageEngine { context: &TransactionContext, overlay: &OverlayState, hash_builder: &mut HashBuilder, + stats: &mut Stats, ) -> Result<(), Error> { // First, collect all existing values from disk if let Some(root_page_id) = context.root_node_page_id { @@ -52,10 +80,11 @@ impl StorageEngine { root_page_id, &Nibbles::new(), Some(context.root_node_hash), + stats, )?; } else { // No root page, just add all overlay changes to the hash builder - self.process_overlay_state(context, overlay, hash_builder)?; + self.process_nonoverlapping_overlay_state(context, overlay, hash_builder, stats)?; } Ok(()) @@ -70,6 +99,7 @@ impl StorageEngine { page_id: PageId, path_prefix: &Nibbles, node_hash: Option, + stats: &mut Stats, ) -> Result<(), Error> { let page = self.get_page(context, page_id)?; let slotted_page = SlottedPage::try_from(page)?; @@ -83,6 +113,7 @@ impl StorageEngine { 0, // Start from root cell path_prefix, node_hash, + stats, ) } @@ -96,6 +127,7 @@ impl StorageEngine { cell_index: u8, path_prefix: &Nibbles, node_hash: Option, + stats: &mut Stats, ) -> Result<(), Error> { let node: Node = slotted_page.get_value(cell_index)?; let full_path = { @@ -114,7 +146,7 @@ impl StorageEngine { // println!("Post-overlay: {:?}", post_overlay.effective_slice()); // Add all pre-overlay changes - self.process_overlay_state(context, &pre_overlay, hash_builder)?; + self.process_nonoverlapping_overlay_state(context, &pre_overlay, hash_builder, stats)?; // Check if this node is affected by overlay if with_prefix.is_empty() { @@ -126,15 +158,17 @@ impl StorageEngine { let rlp_encoded = self.encode_trie_value(&node_value)?; // println!("Adding unaffected leaf: {:?}", full_path); hash_builder.add_leaf(full_path, &rlp_encoded); + stats.account.leaves += 1; } NodeKind::Branch { .. } => { // Branch node - since it's unaffected by overlay, we can use its existing - // hash This is much more efficient than recursing - // into all children println!("Adding branch node: - // {:?}", path_prefix); + // hash. This is much more efficient than recursing into all children + + // println!("Adding branch node: {:?}", path_prefix); let hash = node_hash.unwrap(); // println!("Adding unaffected branch: {:?}", path_prefix); hash_builder.add_branch(path_prefix.clone(), hash, true); + stats.account.branch_nodes += 1; } } break; @@ -143,14 +177,14 @@ impl StorageEngine { let next_overlay_length; if let Some((first_overlay_path, _)) = overlay.get(0) { // If the first overlay path is the same as the current path, we've found an overlay - if first_overlay_path == current_path { - // println!("Direct overlay: {:?}", first_overlay_path); - // Direct overlay - the overlay path exactly matches this node's path - self.process_overlay_state(context, &with_prefix, hash_builder)?; - break; - } else { + // if first_overlay_path == current_path { + // println!("Direct overlay: {:?}", first_overlay_path); + // // Direct overlay - the overlay path exactly matches this node's path + // self.process_nonoverlapping_overlay_state(context, &with_prefix, hash_builder, stats)?; + // break; + // } else { next_overlay_length = first_overlay_path.len(); - } + // } } else { panic!("Overlay path does not match node path {current_path:?}"); } @@ -165,6 +199,7 @@ impl StorageEngine { slotted_page, &node, &full_path, + stats, )?; break; } @@ -174,15 +209,16 @@ impl StorageEngine { current_path = full_path.slice(..path_len); } let post_overlay = overlay.sub_slice_after_prefix(¤t_path); - self.process_overlay_state(context, &post_overlay, hash_builder)?; + self.process_nonoverlapping_overlay_state(context, &post_overlay, hash_builder, stats)?; Ok(()) } - fn process_overlay_state( + fn process_nonoverlapping_overlay_state( &self, context: &TransactionContext, overlay: &OverlayState, hash_builder: &mut HashBuilder, + stats: &mut Stats, ) -> Result<(), Error> { let mut last_processed_path: Option = None; for (path, value) in overlay.iter() { @@ -194,18 +230,27 @@ impl StorageEngine { } match value { Some(OverlayValue::Account(account)) => { - self.handle_overlayed_account(context, overlay, hash_builder, &path, account)?; + self.handle_overlayed_account( + context, + overlay, + hash_builder, + &path, + account, + stats, + )?; last_processed_path = Some(path); } Some(OverlayValue::Storage(storage_value)) => { let rlp_encoded = self.encode_storage(storage_value)?; // println!("Adding overlayed storage leaf: {:?}", path); hash_builder.add_leaf(path.clone(), &rlp_encoded); + stats.storage.overlay_leaves += 1; last_processed_path = Some(path); } Some(OverlayValue::Hash(hash)) => { // println!("Adding overlayed branch: {:?}", path); hash_builder.add_branch(path.clone(), *hash, false); + stats.account.overlay_branch_nodes += 1; last_processed_path = Some(path); } None => { @@ -226,6 +271,7 @@ impl StorageEngine { slotted_page: &SlottedPage, node: &Node, full_path: &Nibbles, + stats: &mut Stats, ) -> Result<(), Error> { match node.kind() { NodeKind::Branch { .. } => { @@ -236,6 +282,7 @@ impl StorageEngine { slotted_page, node, full_path, + stats, )?; } NodeKind::AccountLeaf { .. } => { @@ -247,10 +294,15 @@ impl StorageEngine { slotted_page, Some(node), full_path, + stats, )?; } NodeKind::StorageLeaf { .. } => { - panic!("Storage leaf with descendant overlay: {full_path:?}"); + if let Some((_, Some(OverlayValue::Storage(value)))) = overlay.get(0) { + hash_builder.add_leaf(full_path.clone(), &self.encode_storage(value)?); + stats.storage.overlay_leaves += 1; + } + // panic!("Storage leaf with descendant overlay: {full_path:?}"); } } @@ -266,27 +318,102 @@ impl StorageEngine { slotted_page: &SlottedPage, node: &Node, full_path: &Nibbles, + stats: &mut Stats, ) -> Result<(), Error> { // For each child in the branch node (0-15), check if there are relevant overlay changes - // and recursively traverse child nodes - for child_index in 0..16 { - // Construct the path for this child + // and recursively traverse child nodes. + // As long as there will be at least 2 children, this branch will still exist, allowing us + // to add the children by hash instead of traversing them. + + let mut children: [(Nibbles, OverlayState, Option<&Pointer>); 16] = + std::array::from_fn(|_| (Nibbles::default(), OverlayState::empty(), None::<&Pointer>)); + for i in 0..16 { let mut child_path = full_path.clone(); - child_path.push(child_index); + child_path.push(i); // Create sub-slice of overlay for this child's subtree let child_overlay = overlay.sub_slice_for_prefix(&child_path); + children[i as usize] = (child_path, child_overlay, node.child(i as u8).unwrap()); + } + + let num_children = children + .iter() + .filter(|(_, child_overlay, child_pointer)| { + if let Some((_, value)) = child_overlay.get(0) { + return value.is_some(); + } else if child_pointer.is_some() { + return true; + } + + false + }) + .count(); + + if num_children > 1 { + // We have at least 2 children, so we can add non-overlayed children by hash instead of + // traversing them. This can save many lookups, hashes, and page reads. + for (child_path, child_overlay, child_pointer) in children.iter() { + if let Some((path, _)) = child_overlay.get(0) { + if &path == child_path { + self.process_nonoverlapping_overlay_state(context, &child_overlay, hash_builder, stats)?; + continue; + } + } + + if let Some(child_pointer) = child_pointer { + let child_hash = child_pointer.rlp().as_hash(); + if child_overlay.is_empty() && child_hash.is_some() { + hash_builder.add_branch(child_path.clone(), child_hash.unwrap(), true); + stats.account.pointer_hashes += 1; + } else { + let child_location = child_pointer.location(); + + if let Some(child_cell_index) = child_location.cell_index() { + // Child is in the same page + self.traverse_node_with_overlay( + context, + &child_overlay, + hash_builder, + slotted_page, + child_cell_index, + &child_path, + child_hash, + stats, + )?; + } else if let Some(child_page_id) = child_location.page_id() { + // Child is in a different page + self.traverse_page_with_overlay( + context, + &child_overlay, + hash_builder, + child_page_id, + &child_path, + child_hash, + stats, + )?; + } + } + } else { + // No child exists on disk, process any matching overlay changes + self.process_nonoverlapping_overlay_state(context, &child_overlay, hash_builder, stats)?; + } + } + return Ok(()); + } + + // Otherwise, this branch may no longer exist, so we need to traverse and add the children directly + for (child_path, child_overlay, child_pointer) in children.iter() { if let Some((path, _)) = child_overlay.get(0) { - if path == child_path { + if &path == child_path { // Direct overlay - the overlay path exactly matches this node's path // No need to traverse the child - self.process_overlay_state(context, &child_overlay, hash_builder)?; + self.process_nonoverlapping_overlay_state(context, &child_overlay, hash_builder, stats)?; continue; } } - if let Ok(Some(child_pointer)) = node.child(child_index) { + if let Some(child_pointer) = child_pointer { // Child exists on disk - traverse it with overlay integration let child_hash = child_pointer.rlp().as_hash(); let child_location = child_pointer.location(); @@ -301,6 +428,7 @@ impl StorageEngine { child_cell_index, &child_path, child_hash, + stats, )?; } else if let Some(child_page_id) = child_location.page_id() { // Child is in a different page @@ -311,11 +439,12 @@ impl StorageEngine { child_page_id, &child_path, child_hash, + stats, )?; } } else { // No child exists on disk, process any matching overlay changes - self.process_overlay_state(context, &child_overlay, hash_builder)?; + self.process_nonoverlapping_overlay_state(context, &child_overlay, hash_builder, stats)?; } } Ok(()) @@ -330,6 +459,7 @@ impl StorageEngine { slotted_page: &SlottedPage, node: Option<&Node>, full_path: &Nibbles, + stats: &mut Stats, ) -> Result<(), Error> { // Get the original account from the node let original_account = node.map(|n| match n.value().unwrap() { @@ -346,6 +476,7 @@ impl StorageEngine { // Even with account overlay, we might need to compute storage root for storage // overlays if let OverlayValue::Account(overlay_account) = trie_value { + stats.account.overlay_leaves += 1; Some(overlay_account) } else { panic!("Overlay value is not an account"); @@ -356,6 +487,7 @@ impl StorageEngine { None } } else { + stats.account.leaves += 1; original_account.as_ref() }; @@ -363,7 +495,7 @@ impl StorageEngine { // Check if there are any storage overlays for this account // Storage overlays have 128-nibble paths that start with this account's 64-nibble path let storage_overlays = overlay.find_prefix_range(full_path); - let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() == 128); + let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() > 64); if !has_storage_overlays { // No storage overlays for this account, use original account @@ -379,6 +511,7 @@ impl StorageEngine { node, &storage_overlays, slotted_page, + stats, )?; // Add the modified account to the main HashBuilder @@ -395,7 +528,10 @@ impl StorageEngine { hash_builder: &mut HashBuilder, full_path: &Nibbles, account: &Account, + stats: &mut Stats, ) -> Result<(), Error> { + println!("Handling overlayed account: {full_path:?}"); + println!("Account: {account:?}"); // Check if there are any storage overlays for this account // Storage overlays have 128-nibble paths that start with this account's 64-nibble path let storage_overlays = overlay.find_prefix_range(full_path); @@ -404,7 +540,9 @@ impl StorageEngine { if !has_storage_overlays { // No storage overlays for this account, use original account let rlp_encoded = self.encode_account(account)?; + // println!("Adding overlayed account leaf: {:?}", full_path); hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + stats.account.overlay_leaves += 1; return Ok(()); } @@ -413,13 +551,15 @@ impl StorageEngine { let mut storage_hash_builder = HashBuilder::default().with_updates(true); - self.process_overlay_state(context, &storage_overlay, &mut storage_hash_builder)?; + self.process_nonoverlapping_overlay_state(context, &storage_overlay, &mut storage_hash_builder, stats)?; let new_storage_root = storage_hash_builder.root(); + println!("New storage root: {new_storage_root:?}"); // Add the modified account to the main HashBuilder let rlp_encoded = self.encode_account_with_root(account, new_storage_root)?; hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + stats.account.overlay_leaves += 1; Ok(()) } @@ -432,6 +572,7 @@ impl StorageEngine { account_node: Option<&Node>, storage_overlays: &OverlayState, slotted_page: &SlottedPage, + stats: &mut Stats, ) -> Result { // Create a storage-specific overlay with 64-nibble prefix offset // This converts 128-nibble storage paths to 64-nibble storage-relative paths @@ -469,6 +610,7 @@ impl StorageEngine { storage_page_id, &Nibbles::new(), // Start with empty path for storage root storage_root_hash, + stats, )?; } else if let Some(storage_cell_index) = storage_root_location.cell_index() { // Storage root is in the same page as the account @@ -480,12 +622,18 @@ impl StorageEngine { storage_cell_index, &Nibbles::new(), // Start with empty path for storage root storage_root_hash, + stats, )?; } } } else { // No existing storage, just add overlay changes - self.process_overlay_state(context, &storage_overlay, &mut storage_hash_builder)?; + self.process_nonoverlapping_overlay_state( + context, + &storage_overlay, + &mut storage_hash_builder, + stats, + )?; } let (mut storage_hash_builder, _updated_branch_nodes) = storage_hash_builder.split(); @@ -545,6 +693,7 @@ mod tests { use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY}; use rand::Rng; use tempdir::TempDir; + use test_log::test; #[test] fn test_empty_overlay_root() { @@ -786,7 +935,7 @@ mod tests { let db = Database::create_new(file_path).unwrap(); // Step 1: Write account 1 and compute root - let mut context1 = db.storage_engine.write_context(); + 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, @@ -795,16 +944,7 @@ mod tests { ])); let account1 = Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY); - db.storage_engine - .set_values( - &mut context1, - &mut [(path1.clone().into(), Some(TrieValue::Account(account1.clone())))], - ) - .unwrap(); - let root_with_account1 = context1.root_node_hash; - // Step 2: Add account 2 - let mut context2 = db.storage_engine.write_context(); 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, @@ -813,39 +953,51 @@ mod tests { ])); 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 context2, + &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_both_accounts = context2.root_node_hash; - assert_ne!(root_with_both_accounts, root_with_account1); + let root_with_all_accounts = context.root_node_hash; + assert_ne!(root_with_all_accounts, root_without_account2); - // Step 3: Overlay a tombstone for account 2 and compute root + // 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_with_deletion, _) = - db.storage_engine.compute_root_with_overlay(&context2, &overlay).unwrap(); - assert_ne!(overlay_root_with_deletion, root_with_both_accounts); + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root_with_deletion, root_with_all_accounts); // Step 4: Verify by actually erasing account 2 and computing root - let mut context3 = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut context3, - &mut [ - (path1.clone().into(), Some(TrieValue::Account(account1))), - // path2 is omitted - effectively deleted - ], - ) - .unwrap(); - let root_after_deletion = context3.root_node_hash; + db.storage_engine.set_values(&mut context, &mut [(path2.clone().into(), None)]).unwrap(); + let root_after_deletion = context.root_node_hash; // The overlay root with tombstone should match the root after actual deletion assert_eq!( @@ -855,11 +1007,11 @@ mod tests { // Both should equal the original root with just account1 assert_eq!( - overlay_root_with_deletion, root_with_account1, + overlay_root_with_deletion, root_without_account2, "After deleting account2, root should match original single-account root" ); assert_eq!( - root_after_deletion, root_with_account1, + root_after_deletion, root_without_account2, "After deleting account2, committed root should match original single-account root" ); } @@ -1830,10 +1982,60 @@ mod tests { db.storage_engine.set_values(&mut context, &mut changes).unwrap(); let committed_root = context.root_node_hash; - // db.storage_engine - // .print_page(&context, std::io::stdout(), context.root_node_page_id) - // .unwrap(); - assert_eq!(overlay_root, committed_root, "Overlay new account with storage should match"); } + + #[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 initial_root = context.root_node_hash; + + 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(); + let (overlay_root, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); + + // Verify by committing the storage addition + 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(&mut context, &mut changes).unwrap(); + let committed_root = context.root_node_hash; + + assert_eq!(overlay_root, committed_root, "Overlay update account with storage should match"); + } } diff --git a/src/transaction.rs b/src/transaction.rs index 5005bb88..fd05c1c1 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, }; -use alloy_primitives::{StorageValue, B256}; -use alloy_trie::Nibbles; +use alloy_primitives::{map::HashMap, StorageValue, B256}; +use alloy_trie::{BranchNodeCompact, 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<(B256, HashMap), TransactionError> { + self.database + .storage_engine + .compute_root_with_overlay(&self.context, overlay_state) + .map_err(|_| TransactionError::Generic) + } + pub fn get_account_with_proof( &self, address_path: AddressPath, @@ -192,6 +203,7 @@ impl> Transaction { impl> Transaction { pub fn commit(mut self) -> Result<(), TransactionError> { + println!("transaction_metrics: {:?}", self.context.transaction_metrics); let mut transaction_manager = self.database.transaction_manager.lock(); transaction_manager.remove_tx(self.context.snapshot_id, false); From 5b885aa15e1577d1958f4d1d670b9f240665d375 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Tue, 5 Aug 2025 17:10:29 -0700 Subject: [PATCH 07/17] Track intermediate branch node updates in overlay root --- src/storage/overlay_root.rs | 443 ++++++++++++++++++++++++++---------- src/transaction.rs | 4 +- 2 files changed, 328 insertions(+), 119 deletions(-) diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 51d8f310..5a881e80 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -7,7 +7,10 @@ use crate::{ pointer::Pointer, storage::engine::{Error, StorageEngine}, }; -use alloy_primitives::{map::HashMap, B256, U256}; +use alloy_primitives::{ + map::{B256Map, HashMap}, + B256, U256, +}; use alloy_rlp::encode; use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, TrieAccount}; @@ -41,17 +44,27 @@ impl StorageEngine { &self, context: &TransactionContext, overlay: &OverlayState, - ) -> Result<(B256, HashMap), Error> { + ) -> Result< + (B256, HashMap, B256Map>), + Error, + > { if overlay.is_empty() { // No overlay changes, return current root - return Ok((context.root_node_hash, HashMap::default())); + return Ok((context.root_node_hash, HashMap::default(), B256Map::default())); } let mut hash_builder = HashBuilder::default().with_updates(true); let mut stats = Stats::default(); + let mut updated_storage_branch_nodes = B256Map::default(); // Use proper trie traversal with overlay integration - self.traverse_with_overlay(context, overlay, &mut hash_builder, &mut stats)?; + self.traverse_with_overlay( + context, + overlay, + &mut hash_builder, + &mut stats, + &mut updated_storage_branch_nodes, + )?; let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); // println!("Updated branch nodes: {:?}", updated_branch_nodes); @@ -60,7 +73,7 @@ impl StorageEngine { // println!("Root: {:?}", root); println!("Stats: {:?}", stats); - Ok((root, updated_branch_nodes)) + Ok((root, updated_branch_nodes, updated_storage_branch_nodes)) } /// Helper method to traverse the existing trie and integrate with overlay @@ -70,6 +83,7 @@ impl StorageEngine { overlay: &OverlayState, hash_builder: &mut HashBuilder, stats: &mut Stats, + updated_storage_branch_nodes: &mut B256Map>, ) -> Result<(), Error> { // First, collect all existing values from disk if let Some(root_page_id) = context.root_node_page_id { @@ -81,10 +95,17 @@ impl StorageEngine { &Nibbles::new(), Some(context.root_node_hash), stats, + updated_storage_branch_nodes, )?; } else { // No root page, just add all overlay changes to the hash builder - self.process_nonoverlapping_overlay_state(context, overlay, hash_builder, stats)?; + self.process_nonoverlapping_overlay_state( + context, + overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; } Ok(()) @@ -100,6 +121,7 @@ impl StorageEngine { path_prefix: &Nibbles, node_hash: Option, stats: &mut Stats, + updated_storage_branch_nodes: &mut B256Map>, ) -> Result<(), Error> { let page = self.get_page(context, page_id)?; let slotted_page = SlottedPage::try_from(page)?; @@ -114,10 +136,16 @@ impl StorageEngine { path_prefix, node_hash, stats, + updated_storage_branch_nodes, ) } - /// Traverse a specific node with overlay integration + /// Traverse a specific node with overlayed state. + /// If this node is a branch, we can add it by hash if there are no overlayed changes under its + /// path prefix. If this node is a leaf, we can add it if there are no overlayed changes + /// under its full path. Otherwise, we need to traverse through the node via its full path. + /// In all cases, we need to process the overlayed state before and after the added or traversed + /// node path. fn traverse_node_with_overlay( &self, context: &TransactionContext, @@ -128,6 +156,7 @@ impl StorageEngine { path_prefix: &Nibbles, node_hash: Option, stats: &mut Stats, + updated_storage_branch_nodes: &mut B256Map>, ) -> Result<(), Error> { let node: Node = slotted_page.get_value(cell_index)?; let full_path = { @@ -136,81 +165,104 @@ impl StorageEngine { full }; - let mut current_path = path_prefix.clone(); - while current_path.len() <= full_path.len() { - // println!("Current path: {:?}", current_path); - // println!("Full path: {:?}", full_path); - let (pre_overlay, with_prefix, _) = overlay.sub_slice_by_prefix(¤t_path); - // println!("Pre-overlay: {:?}", pre_overlay.effective_slice()); - // println!("With-prefix: {:?}", with_prefix.effective_slice()); - // println!("Post-overlay: {:?}", post_overlay.effective_slice()); - - // Add all pre-overlay changes - self.process_nonoverlapping_overlay_state(context, &pre_overlay, hash_builder, stats)?; - - // Check if this node is affected by overlay - if with_prefix.is_empty() { - // Node not affected by overlay, use disk value as-is - match node.kind() { - NodeKind::AccountLeaf { .. } | NodeKind::StorageLeaf { .. } => { - // Leaf node - use disk value directly - let node_value = node.value()?; - let rlp_encoded = self.encode_trie_value(&node_value)?; - // println!("Adding unaffected leaf: {:?}", full_path); - hash_builder.add_leaf(full_path, &rlp_encoded); - stats.account.leaves += 1; - } - NodeKind::Branch { .. } => { - // Branch node - since it's unaffected by overlay, we can use its existing - // hash. This is much more efficient than recursing into all children - - // println!("Adding branch node: {:?}", path_prefix); - let hash = node_hash.unwrap(); - // println!("Adding unaffected branch: {:?}", path_prefix); - hash_builder.add_branch(path_prefix.clone(), hash, true); - stats.account.branch_nodes += 1; - } + match node.kind() { + // if the node is a branch, we can only add it by hash if the path_prefix is completely + // unaffected by the overlay + NodeKind::Branch { .. } => { + // Filter the overlay based on the path_prefix + let (pre_overlay, with_prefix, post_overlay) = + overlay.sub_slice_by_prefix(&path_prefix); + if with_prefix.is_empty() { + self.process_nonoverlapping_overlay_state( + context, + &pre_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; + let hash = node_hash.unwrap(); + println!("Adding unaffected branch: {:?}", path_prefix); + hash_builder.add_branch(path_prefix.clone(), hash, true); + stats.account.branch_nodes += 1; + self.process_nonoverlapping_overlay_state( + context, + &post_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; + return Ok(()); + } else { + let (pre_overlay, with_prefix, post_overlay) = + overlay.sub_slice_by_prefix(&full_path); + self.process_nonoverlapping_overlay_state( + context, + &pre_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; + println!("Processing affected branch: {:?}", path_prefix); + self.handle_affected_node_with_overlay( + context, + &with_prefix, + hash_builder, + slotted_page, + &node, + &full_path, + stats, + updated_storage_branch_nodes, + )?; + self.process_nonoverlapping_overlay_state( + context, + &post_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; + return Ok(()); } - break; } - - let next_overlay_length; - if let Some((first_overlay_path, _)) = overlay.get(0) { - // If the first overlay path is the same as the current path, we've found an overlay - // if first_overlay_path == current_path { - // println!("Direct overlay: {:?}", first_overlay_path); - // // Direct overlay - the overlay path exactly matches this node's path - // self.process_nonoverlapping_overlay_state(context, &with_prefix, hash_builder, stats)?; - // break; - // } else { - next_overlay_length = first_overlay_path.len(); - // } - } else { - panic!("Overlay path does not match node path {current_path:?}"); - } - - // If we've reached the end of the full path, we've found the node - if current_path.len() == full_path.len() { - // println!("Found node: {:?}", full_path); - self.handle_affected_node_with_overlay( + NodeKind::AccountLeaf { .. } | NodeKind::StorageLeaf { .. } => { + // filter the overlay based on the full_path + let (pre_overlay, with_prefix, post_overlay) = + overlay.sub_slice_by_prefix(&full_path); + self.process_nonoverlapping_overlay_state( context, - overlay, + &pre_overlay, hash_builder, - slotted_page, - &node, - &full_path, stats, + updated_storage_branch_nodes, )?; - break; + if with_prefix.is_empty() { + println!("Adding unaffected leaf: {:?}", full_path); + let node_value = node.value()?; + let rlp_encoded = self.encode_trie_value(&node_value)?; + hash_builder.add_leaf(full_path, &rlp_encoded); + stats.account.leaves += 1; + } else { + println!("Processing affected leaf: {:?}", full_path); + self.handle_affected_node_with_overlay( + context, + &with_prefix, + hash_builder, + slotted_page, + &node, + &full_path, + stats, + updated_storage_branch_nodes, + )?; + } + self.process_nonoverlapping_overlay_state( + context, + &post_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; + return Ok(()); } - - // Fast forward the path length to the next length at which the overlay changes - let path_len = std::cmp::min(full_path.len(), next_overlay_length + 1); - current_path = full_path.slice(..path_len); } - let post_overlay = overlay.sub_slice_after_prefix(¤t_path); - self.process_nonoverlapping_overlay_state(context, &post_overlay, hash_builder, stats)?; - Ok(()) } fn process_nonoverlapping_overlay_state( @@ -219,6 +271,7 @@ impl StorageEngine { overlay: &OverlayState, hash_builder: &mut HashBuilder, stats: &mut Stats, + updated_storage_branch_nodes: &mut B256Map>, ) -> Result<(), Error> { let mut last_processed_path: Option = None; for (path, value) in overlay.iter() { @@ -237,18 +290,19 @@ impl StorageEngine { &path, account, stats, + updated_storage_branch_nodes, )?; last_processed_path = Some(path); } Some(OverlayValue::Storage(storage_value)) => { let rlp_encoded = self.encode_storage(storage_value)?; - // println!("Adding overlayed storage leaf: {:?}", path); + println!("Adding overlayed storage leaf: {:?}", path); hash_builder.add_leaf(path.clone(), &rlp_encoded); stats.storage.overlay_leaves += 1; last_processed_path = Some(path); } Some(OverlayValue::Hash(hash)) => { - // println!("Adding overlayed branch: {:?}", path); + println!("Adding overlayed branch: {:?}", path); hash_builder.add_branch(path.clone(), *hash, false); stats.account.overlay_branch_nodes += 1; last_processed_path = Some(path); @@ -272,6 +326,7 @@ impl StorageEngine { node: &Node, full_path: &Nibbles, stats: &mut Stats, + updated_storage_branch_nodes: &mut B256Map>, ) -> Result<(), Error> { match node.kind() { NodeKind::Branch { .. } => { @@ -283,6 +338,7 @@ impl StorageEngine { node, full_path, stats, + updated_storage_branch_nodes, )?; } NodeKind::AccountLeaf { .. } => { @@ -295,14 +351,16 @@ impl StorageEngine { Some(node), full_path, stats, + updated_storage_branch_nodes, )?; } NodeKind::StorageLeaf { .. } => { if let Some((_, Some(OverlayValue::Storage(value)))) = overlay.get(0) { hash_builder.add_leaf(full_path.clone(), &self.encode_storage(value)?); stats.storage.overlay_leaves += 1; + } else { + panic!("Storage leaf does not have a valid overlay: {full_path:?}"); } - // panic!("Storage leaf with descendant overlay: {full_path:?}"); } } @@ -319,6 +377,7 @@ impl StorageEngine { node: &Node, full_path: &Nibbles, stats: &mut Stats, + updated_storage_branch_nodes: &mut B256Map>, ) -> Result<(), Error> { // For each child in the branch node (0-15), check if there are relevant overlay changes // and recursively traverse child nodes. @@ -356,7 +415,13 @@ impl StorageEngine { for (child_path, child_overlay, child_pointer) in children.iter() { if let Some((path, _)) = child_overlay.get(0) { if &path == child_path { - self.process_nonoverlapping_overlay_state(context, &child_overlay, hash_builder, stats)?; + self.process_nonoverlapping_overlay_state( + context, + &child_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; continue; } } @@ -364,6 +429,7 @@ impl StorageEngine { if let Some(child_pointer) = child_pointer { let child_hash = child_pointer.rlp().as_hash(); if child_overlay.is_empty() && child_hash.is_some() { + println!("Adding branch node by pointer: {:?}", child_path); hash_builder.add_branch(child_path.clone(), child_hash.unwrap(), true); stats.account.pointer_hashes += 1; } else { @@ -380,6 +446,7 @@ impl StorageEngine { &child_path, child_hash, stats, + updated_storage_branch_nodes, )?; } else if let Some(child_page_id) = child_location.page_id() { // Child is in a different page @@ -391,24 +458,38 @@ impl StorageEngine { &child_path, child_hash, stats, + updated_storage_branch_nodes, )?; } } } else { // No child exists on disk, process any matching overlay changes - self.process_nonoverlapping_overlay_state(context, &child_overlay, hash_builder, stats)?; + self.process_nonoverlapping_overlay_state( + context, + &child_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; } } return Ok(()); } - // Otherwise, this branch may no longer exist, so we need to traverse and add the children directly + // Otherwise, this branch may no longer exist, so we need to traverse and add the children + // directly for (child_path, child_overlay, child_pointer) in children.iter() { if let Some((path, _)) = child_overlay.get(0) { if &path == child_path { // Direct overlay - the overlay path exactly matches this node's path // No need to traverse the child - self.process_nonoverlapping_overlay_state(context, &child_overlay, hash_builder, stats)?; + self.process_nonoverlapping_overlay_state( + context, + &child_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; continue; } } @@ -429,6 +510,7 @@ impl StorageEngine { &child_path, child_hash, stats, + updated_storage_branch_nodes, )?; } else if let Some(child_page_id) = child_location.page_id() { // Child is in a different page @@ -440,11 +522,18 @@ impl StorageEngine { &child_path, child_hash, stats, + updated_storage_branch_nodes, )?; } } else { // No child exists on disk, process any matching overlay changes - self.process_nonoverlapping_overlay_state(context, &child_overlay, hash_builder, stats)?; + self.process_nonoverlapping_overlay_state( + context, + &child_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; } } Ok(()) @@ -460,6 +549,7 @@ impl StorageEngine { node: Option<&Node>, full_path: &Nibbles, stats: &mut Stats, + updated_storage_branch_nodes: &mut B256Map>, ) -> Result<(), Error> { // Get the original account from the node let original_account = node.map(|n| match n.value().unwrap() { @@ -500,22 +590,26 @@ impl StorageEngine { if !has_storage_overlays { // No storage overlays for this account, use original account let rlp_encoded = self.encode_account(account)?; - // println!("Adding account leaf: {full_path:?}"); + println!("Adding account leaf: {full_path:?}"); hash_builder.add_leaf(full_path.clone(), &rlp_encoded); return Ok(()); } // We have storage overlays, need to compute a new storage root using iterative approach - let new_storage_root = self.compute_storage_root_iteratively( + let (new_storage_root, updated_branch_nodes) = self.compute_storage_root_with_overlay( context, node, &storage_overlays, slotted_page, stats, + updated_storage_branch_nodes, )?; + updated_storage_branch_nodes + .insert(B256::from_slice(&full_path.pack()), updated_branch_nodes); // Add the modified account to the main HashBuilder let rlp_encoded = self.encode_account_with_root(account, new_storage_root)?; + println!("Adding modified account leaf: {full_path:?}"); hash_builder.add_leaf(full_path.clone(), &rlp_encoded); } Ok(()) @@ -529,13 +623,14 @@ impl StorageEngine { full_path: &Nibbles, account: &Account, stats: &mut Stats, + updated_storage_branch_nodes: &mut B256Map>, ) -> Result<(), Error> { println!("Handling overlayed account: {full_path:?}"); println!("Account: {account:?}"); // Check if there are any storage overlays for this account // Storage overlays have 128-nibble paths that start with this account's 64-nibble path let storage_overlays = overlay.find_prefix_range(full_path); - let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() == 128); + let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() > 64); if !has_storage_overlays { // No storage overlays for this account, use original account @@ -551,10 +646,19 @@ impl StorageEngine { let mut storage_hash_builder = HashBuilder::default().with_updates(true); - self.process_nonoverlapping_overlay_state(context, &storage_overlay, &mut storage_hash_builder, stats)?; + self.process_nonoverlapping_overlay_state( + context, + &storage_overlay, + &mut storage_hash_builder, + stats, + updated_storage_branch_nodes, + )?; + let (mut storage_hash_builder, updated_branch_nodes) = storage_hash_builder.split(); let new_storage_root = storage_hash_builder.root(); println!("New storage root: {new_storage_root:?}"); + updated_storage_branch_nodes + .insert(B256::from_slice(&full_path.pack()), updated_branch_nodes); // Add the modified account to the main HashBuilder let rlp_encoded = self.encode_account_with_root(account, new_storage_root)?; @@ -566,14 +670,15 @@ impl StorageEngine { /// Compute the storage root for an account using iterative HashBuilder approach with prefix /// offset - fn compute_storage_root_iteratively( + fn compute_storage_root_with_overlay( &self, context: &TransactionContext, account_node: Option<&Node>, storage_overlays: &OverlayState, slotted_page: &SlottedPage, stats: &mut Stats, - ) -> Result { + updated_storage_branch_nodes: &mut B256Map>, + ) -> Result<(B256, HashMap), Error> { // Create a storage-specific overlay with 64-nibble prefix offset // This converts 128-nibble storage paths to 64-nibble storage-relative paths let mut storage_overlay = storage_overlays.with_prefix_offset(64); @@ -611,6 +716,7 @@ impl StorageEngine { &Nibbles::new(), // Start with empty path for storage root storage_root_hash, stats, + updated_storage_branch_nodes, )?; } else if let Some(storage_cell_index) = storage_root_location.cell_index() { // Storage root is in the same page as the account @@ -623,6 +729,7 @@ impl StorageEngine { &Nibbles::new(), // Start with empty path for storage root storage_root_hash, stats, + updated_storage_branch_nodes, )?; } } @@ -633,15 +740,16 @@ impl StorageEngine { &storage_overlay, &mut storage_hash_builder, stats, + updated_storage_branch_nodes, )?; } - let (mut storage_hash_builder, _updated_branch_nodes) = storage_hash_builder.split(); + let (mut storage_hash_builder, updated_branch_nodes) = storage_hash_builder.split(); // println!("Updated storage branch nodes: {updated_branch_nodes:?}"); let computed_root = storage_hash_builder.root(); // println!("Computed storage root: {computed_root:?}"); - Ok(computed_root) + Ok((computed_root, updated_branch_nodes)) } /// Helper to encode TrieValue as RLP @@ -704,7 +812,7 @@ mod tests { let context = db.storage_engine.read_context(); let empty_overlay = OverlayStateMut::new().freeze(); - let (root, _) = + let (root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &empty_overlay).unwrap(); assert_eq!(root, context.root_node_hash); } @@ -727,7 +835,7 @@ mod tests { let overlay = overlay_mut.freeze(); // Compute root with overlay - let (root, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let (root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); // The root should be different from the empty root (since we have changes) assert_ne!(root, EMPTY_ROOT_HASH); @@ -758,7 +866,7 @@ mod tests { // Test that we can compute root with empty overlay first (should match initial_root) let empty_overlay = OverlayStateMut::new().freeze(); - let (root_with_empty_overlay, _) = + let (root_with_empty_overlay, _, _) = db.storage_engine.compute_root_with_overlay(&context, &empty_overlay).unwrap(); assert_eq!(root_with_empty_overlay, initial_root); @@ -770,7 +878,7 @@ mod tests { .insert(address_path.clone().into(), Some(OverlayValue::Account(account2.clone()))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -839,7 +947,7 @@ mod tests { .insert(path1.clone().into(), Some(OverlayValue::Account(account1_updated.clone()))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -873,7 +981,7 @@ mod tests { overlay_mut2.insert(path3.clone().into(), Some(OverlayValue::Account(account3.clone()))); let overlay2 = overlay_mut2.freeze(); - let (overlay_root2, _) = + let (overlay_root2, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay2).unwrap(); assert_ne!(overlay_root2, initial_root); assert_ne!(overlay_root2, overlay_root); @@ -904,7 +1012,7 @@ mod tests { // path2 is left unaffected - should use add_branch optimization let overlay3 = overlay_mut3.freeze(); - let (overlay_root3, _) = + let (overlay_root3, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay3).unwrap(); assert_ne!(overlay_root3, initial_root); assert_ne!(overlay_root3, overlay_root); @@ -991,7 +1099,7 @@ mod tests { overlay_mut.insert(path2.clone().into(), None); // Delete account2 let overlay = overlay_mut.freeze(); - let (overlay_root_with_deletion, _) = + let (overlay_root_with_deletion, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root_with_deletion, root_with_all_accounts); @@ -1062,7 +1170,7 @@ mod tests { .insert(storage_path1.full_path(), Some(OverlayValue::Storage(new_storage_value1))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1094,7 +1202,7 @@ mod tests { overlay_mut2.insert(storage_path3.full_path(), Some(OverlayValue::Storage(storage_value3))); let overlay2 = overlay_mut2.freeze(); - let (overlay_root2, _) = + let (overlay_root2, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay2).unwrap(); assert_ne!(overlay_root2, initial_root); assert_ne!(overlay_root2, overlay_root); @@ -1124,7 +1232,7 @@ mod tests { overlay_mut3.insert(storage_path2.full_path(), None); // Delete storage slot let overlay3 = overlay_mut3.freeze(); - let (overlay_root3, _) = + let (overlay_root3, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay3).unwrap(); assert_ne!(overlay_root3, initial_root); @@ -1158,7 +1266,7 @@ mod tests { .insert(storage_path1.full_path(), Some(OverlayValue::Storage(new_storage_value1))); let overlay4 = overlay_mut4.freeze(); - let (overlay_root4, _) = + let (overlay_root4, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay4).unwrap(); assert_ne!(overlay_root4, initial_root); @@ -1210,7 +1318,7 @@ mod tests { ); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1283,7 +1391,7 @@ mod tests { overlay_mut.insert(storage_path2.full_path(), Some(OverlayValue::Storage(U256::from(222)))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1346,7 +1454,7 @@ mod tests { overlay_mut.insert(storage2_path.full_path(), Some(OverlayValue::Storage(U256::from(888)))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1420,7 +1528,7 @@ mod tests { let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1478,7 +1586,7 @@ mod tests { overlay_mut.insert(storage_path1.full_path(), Some(OverlayValue::Storage(U256::from(999)))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1564,7 +1672,7 @@ mod tests { let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1637,7 +1745,7 @@ mod tests { overlay_mut.insert(storage1_path.full_path(), Some(OverlayValue::Storage(U256::from(999)))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1704,7 +1812,7 @@ mod tests { overlay_mut.insert(storage2_path.full_path(), Some(OverlayValue::Storage(U256::from(888)))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1771,7 +1879,7 @@ mod tests { .insert(storage1_path2.full_path(), Some(OverlayValue::Storage(U256::from(444)))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1820,7 +1928,7 @@ mod tests { overlay_mut.insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1890,7 +1998,7 @@ mod tests { let mut overlay_mut_with_branches = overlay_mut.clone(); let overlay = overlay_mut.freeze(); - let (overlay_root, overlay_branches) = + let (overlay_root, overlay_branches, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -1914,7 +2022,7 @@ mod tests { } let overlay_with_branches = overlay_mut_with_branches.freeze(); // println!("Overlay with branches: {:?}", overlay_with_branches); - let (overlay_root_with_branches, _) = + let (overlay_root_with_branches, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay_with_branches).unwrap(); assert_eq!(overlay_root_with_branches, overlay_root); @@ -1969,7 +2077,7 @@ mod tests { overlay_mut.insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -2023,7 +2131,7 @@ mod tests { overlay_mut.insert(storage_path1.full_path(), Some(OverlayValue::Storage(U256::from(333)))); let overlay = overlay_mut.freeze(); - let (overlay_root, _) = + let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); assert_ne!(overlay_root, initial_root); @@ -2036,6 +2144,107 @@ mod tests { db.storage_engine.set_values(&mut context, &mut changes).unwrap(); let committed_root = context.root_node_hash; - assert_eq!(overlay_root, committed_root, "Overlay update account with storage should match"); + assert_eq!( + overlay_root, committed_root, + "Overlay update account with storage should match" + ); + } + + #[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 + let result = db.storage_engine.compute_root_with_overlay(&context, &overlay); + let overlay_root = result.unwrap().0; + + // commit the overlay, ensure that the root is the same as the initial 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(&mut context, &mut changes).unwrap(); + let committed_root = context.root_node_hash; + assert_eq!(committed_root, overlay_root); } } diff --git a/src/transaction.rs b/src/transaction.rs index fd05c1c1..f3b7535a 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -10,7 +10,7 @@ use crate::{ path::{AddressPath, StoragePath}, storage::proofs::AccountProof, }; -use alloy_primitives::{map::HashMap, StorageValue, B256}; +use alloy_primitives::{map::{B256Map, HashMap}, StorageValue, B256}; use alloy_trie::{BranchNodeCompact, Nibbles}; pub use error::TransactionError; pub use manager::TransactionManager; @@ -88,7 +88,7 @@ impl, K: TransactionKind> Transaction { pub fn compute_root_with_overlay( &self, overlay_state: &OverlayState, - ) -> Result<(B256, HashMap), TransactionError> { + ) -> Result<(B256, HashMap, B256Map>), TransactionError> { self.database .storage_engine .compute_root_with_overlay(&self.context, overlay_state) From ff2e8d786c6cee2c1a9c535d4427f4b7d84320fe Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Thu, 7 Aug 2025 17:54:23 -0700 Subject: [PATCH 08/17] Handle edge case of extension with modified prefix --- benches/crud_benchmarks.rs | 10 +- src/overlay.rs | 6 +- src/storage/overlay_root.rs | 255 +++++++++++++++++++++--------------- src/transaction.rs | 10 +- 4 files changed, 163 insertions(+), 118 deletions(-) diff --git a/benches/crud_benchmarks.rs b/benches/crud_benchmarks.rs index 35e38b89..7454e2ac 100644 --- a/benches/crud_benchmarks.rs +++ b/benches/crud_benchmarks.rs @@ -603,7 +603,8 @@ fn bench_state_root_with_overlay(c: &mut Criterion) { let total_addresses = BATCH_SIZE; let addresses: Vec = (0..total_addresses).map(|_| generate_random_address(&mut rng)).collect(); - // let storage_paths_values = generate_storage_paths_values(&addresses, total_storage_per_address); + // let storage_paths_values = generate_storage_paths_values(&addresses, + // total_storage_per_address); let mut account_overlay_mut = OverlayStateMut::new(); addresses.iter().enumerate().for_each(|(i, addr)| { @@ -635,7 +636,7 @@ fn bench_state_root_with_overlay(c: &mut Criterion) { let db_path = dir.path().join(&file_name); Database::open(db_path).unwrap() }, - |db| { + |db| { let tx = db.begin_ro().unwrap(); // Compute the root hash with the overlay @@ -660,8 +661,9 @@ fn bench_state_root_with_overlay(c: &mut Criterion) { // let tx = db.begin_ro().unwrap(); // // Compute the root hash with the overlay - // let _root_result = tx.compute_root_with_overlay(&storage_overlay).unwrap(); - + // let _root_result = + // tx.compute_root_with_overlay(&storage_overlay).unwrap(); + // tx.commit().unwrap(); // } // }) diff --git a/src/overlay.rs b/src/overlay.rs index 1b004a5f..a65d570f 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -776,10 +776,8 @@ mod tests { ]; for path in &paths { - mutable.insert( - Nibbles::from_nibbles(path.clone()), - Some(OverlayValue::Account(account.clone())), - ); + mutable + .insert(Nibbles::from_nibbles(*path), Some(OverlayValue::Account(account.clone()))); } let frozen = mutable.freeze(); diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 5a881e80..281a1304 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -171,7 +171,7 @@ impl StorageEngine { NodeKind::Branch { .. } => { // Filter the overlay based on the path_prefix let (pre_overlay, with_prefix, post_overlay) = - overlay.sub_slice_by_prefix(&path_prefix); + overlay.sub_slice_by_prefix(path_prefix); if with_prefix.is_empty() { self.process_nonoverlapping_overlay_state( context, @@ -180,18 +180,40 @@ impl StorageEngine { stats, updated_storage_branch_nodes, )?; - let hash = node_hash.unwrap(); - println!("Adding unaffected branch: {:?}", path_prefix); - hash_builder.add_branch(path_prefix.clone(), hash, true); - stats.account.branch_nodes += 1; - self.process_nonoverlapping_overlay_state( - context, - &post_overlay, - hash_builder, - stats, - updated_storage_branch_nodes, - )?; - return Ok(()); + if !node.prefix().is_empty() { + // The node is actually an extension + branch node. We cannot add it by hash + // as the extension prefix may have changed, which + // would in turn change the hash. We need to + // traverse the node via its full path. + println!( + "Processing extension + branch: {:?} / {:?}", + path_prefix, full_path + ); + self.handle_affected_node_with_overlay( + context, + &with_prefix, + hash_builder, + slotted_page, + &node, + &full_path, + stats, + updated_storage_branch_nodes, + )?; + } else { + // This is a true branch node, so we can add it by hash. + let hash = node_hash.unwrap(); + println!("Adding unaffected branch: {:?}", path_prefix); + hash_builder.add_branch(path_prefix.clone(), hash, true); + stats.account.branch_nodes += 1; + self.process_nonoverlapping_overlay_state( + context, + &post_overlay, + hash_builder, + stats, + updated_storage_branch_nodes, + )?; + } + Ok(()) } else { let (pre_overlay, with_prefix, post_overlay) = overlay.sub_slice_by_prefix(&full_path); @@ -220,7 +242,7 @@ impl StorageEngine { stats, updated_storage_branch_nodes, )?; - return Ok(()); + Ok(()) } } NodeKind::AccountLeaf { .. } | NodeKind::StorageLeaf { .. } => { @@ -260,7 +282,7 @@ impl StorageEngine { stats, updated_storage_branch_nodes, )?; - return Ok(()); + Ok(()) } } } @@ -309,6 +331,7 @@ impl StorageEngine { } None => { // Tombstone - skip + println!("Skipping tombstone: {:?}", path); last_processed_path = Some(path); } } @@ -354,14 +377,19 @@ impl StorageEngine { updated_storage_branch_nodes, )?; } - NodeKind::StorageLeaf { .. } => { - if let Some((_, Some(OverlayValue::Storage(value)))) = overlay.get(0) { + NodeKind::StorageLeaf { .. } => match overlay.get(0) { + Some((_, Some(OverlayValue::Storage(value)))) => { hash_builder.add_leaf(full_path.clone(), &self.encode_storage(value)?); stats.storage.overlay_leaves += 1; - } else { + } + Some((_, None)) => { + println!("Skipping deleted storage leaf: {:?}", full_path); + stats.storage.overlay_leaves += 1; + } + _ => { panic!("Storage leaf does not have a valid overlay: {full_path:?}"); } - } + }, } Ok(()) @@ -393,15 +421,20 @@ impl StorageEngine { // Create sub-slice of overlay for this child's subtree let child_overlay = overlay.sub_slice_for_prefix(&child_path); - children[i as usize] = (child_path, child_overlay, node.child(i as u8).unwrap()); + children[i as usize] = (child_path, child_overlay, node.child(i).unwrap()); } - let num_children = children + // This conservatively counts the minimum number of children that will exist. It may + // undercount. + let min_num_children = children .iter() .filter(|(_, child_overlay, child_pointer)| { if let Some((_, value)) = child_overlay.get(0) { + // If the overlayed value is none, this branch path might be removed. return value.is_some(); } else if child_pointer.is_some() { + // If there is no overlayed value and an existing child, then this path will + // definitely exist. return true; } @@ -409,15 +442,16 @@ impl StorageEngine { }) .count(); - if num_children > 1 { - // We have at least 2 children, so we can add non-overlayed children by hash instead of - // traversing them. This can save many lookups, hashes, and page reads. + if min_num_children > 1 { + // We are guaranteed to have at least 2 children, keeping this branch node alive. + // We can add non-overlayed children by hash instead of traversing them. This can save + // many lookups, hashes, and page reads. for (child_path, child_overlay, child_pointer) in children.iter() { if let Some((path, _)) = child_overlay.get(0) { if &path == child_path { self.process_nonoverlapping_overlay_state( context, - &child_overlay, + child_overlay, hash_builder, stats, updated_storage_branch_nodes, @@ -439,11 +473,11 @@ impl StorageEngine { // Child is in the same page self.traverse_node_with_overlay( context, - &child_overlay, + child_overlay, hash_builder, slotted_page, child_cell_index, - &child_path, + child_path, child_hash, stats, updated_storage_branch_nodes, @@ -452,10 +486,10 @@ impl StorageEngine { // Child is in a different page self.traverse_page_with_overlay( context, - &child_overlay, + child_overlay, hash_builder, child_page_id, - &child_path, + child_path, child_hash, stats, updated_storage_branch_nodes, @@ -466,7 +500,7 @@ impl StorageEngine { // No child exists on disk, process any matching overlay changes self.process_nonoverlapping_overlay_state( context, - &child_overlay, + child_overlay, hash_builder, stats, updated_storage_branch_nodes, @@ -476,8 +510,8 @@ impl StorageEngine { return Ok(()); } - // Otherwise, this branch may no longer exist, so we need to traverse and add the children - // directly + // Otherwise, it is possible that this branch might be removed. + // We need to traverse and add the children directly in order to be safe. for (child_path, child_overlay, child_pointer) in children.iter() { if let Some((path, _)) = child_overlay.get(0) { if &path == child_path { @@ -485,7 +519,7 @@ impl StorageEngine { // No need to traverse the child self.process_nonoverlapping_overlay_state( context, - &child_overlay, + child_overlay, hash_builder, stats, updated_storage_branch_nodes, @@ -503,11 +537,11 @@ impl StorageEngine { // Child is in the same page self.traverse_node_with_overlay( context, - &child_overlay, + child_overlay, hash_builder, slotted_page, child_cell_index, - &child_path, + child_path, child_hash, stats, updated_storage_branch_nodes, @@ -516,10 +550,10 @@ impl StorageEngine { // Child is in a different page self.traverse_page_with_overlay( context, - &child_overlay, + child_overlay, hash_builder, child_page_id, - &child_path, + child_path, child_hash, stats, updated_storage_branch_nodes, @@ -529,7 +563,7 @@ impl StorageEngine { // No child exists on disk, process any matching overlay changes self.process_nonoverlapping_overlay_state( context, - &child_overlay, + child_overlay, hash_builder, stats, updated_storage_branch_nodes, @@ -573,7 +607,7 @@ impl StorageEngine { } } else { // If overlay_value is None, it's a deletion - skip - // println!("Skipping deletion: {full_path:?}"); + println!("Skipping deleted account: {full_path:?}"); None } } else { @@ -1953,93 +1987,98 @@ mod tests { #[test] fn test_1000_accounts_with_10_overlay() { - 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(); + for _ in 0..100 { + 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 context = db.storage_engine.write_context(); + let mut rng = rand::rng(); - let mut changes: Vec<(Nibbles, Option)> = Vec::with_capacity(1000); + 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); + for i in 0..10000 { + 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.push((account_path.into(), Some(TrieValue::Account(account)))); + } - changes.sort_by(|a, b| a.0.cmp(&b.0)); + changes.sort_by(|a, b| a.0.cmp(&b.0)); - db.storage_engine.set_values(&mut context, &mut changes).unwrap(); + db.storage_engine.set_values(&mut context, &mut changes).unwrap(); - let initial_root = context.root_node_hash; + let initial_root = context.root_node_hash; - // Create overlay with modifications to every 100th account - let mut overlay_mut = OverlayStateMut::new(); + // 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); + // 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 mut overlay_mut_with_branches = overlay_mut.clone(); - let overlay = overlay_mut.freeze(); - - let (overlay_root, overlay_branches, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); + let mut overlay_mut_with_branches = overlay_mut.clone(); + let overlay = overlay_mut.freeze(); - // println!("Overlay branches: {:?}", overlay_branches); + let (overlay_root, overlay_branches, _) = + db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + assert_ne!(overlay_root, initial_root); - for (path, branch) in overlay_branches.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); + for (path, branch) in overlay_branches.iter() { + if let Some(root_hash) = branch.root_hash { overlay_mut_with_branches - .insert(path.clone(), Some(OverlayValue::Hash(branch.hashes[hash_idx]))); - hash_idx += 1; - path.pop(); + .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(); + } } } + let overlay_with_branches = overlay_mut_with_branches.freeze(); + // println!("Overlay with branches: {:?}", overlay_with_branches); + let (overlay_root_with_branches, _, _) = db + .storage_engine + .compute_root_with_overlay(&context, &overlay_with_branches) + .unwrap(); + assert_eq!(overlay_root_with_branches, overlay_root); + + // Verify by committing the storage addition + 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(&mut context, &mut changes).unwrap(); + let committed_root = context.root_node_hash; + + // db.storage_engine + // .print_page(&context, std::io::stdout(), context.root_node_page_id) + // .unwrap(); + + assert_eq!(overlay_root, committed_root, "1000 accounts with 10 overlay should match"); } - let overlay_with_branches = overlay_mut_with_branches.freeze(); - // println!("Overlay with branches: {:?}", overlay_with_branches); - let (overlay_root_with_branches, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay_with_branches).unwrap(); - assert_eq!(overlay_root_with_branches, overlay_root); - - // Verify by committing the storage addition - 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(&mut context, &mut changes).unwrap(); - let committed_root = context.root_node_hash; - - // db.storage_engine - // .print_page(&context, std::io::stdout(), context.root_node_page_id) - // .unwrap(); - - assert_eq!(overlay_root, committed_root, "1000 accounts with 10 overlay should match"); } #[test] diff --git a/src/transaction.rs b/src/transaction.rs index f3b7535a..886a895b 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -10,7 +10,10 @@ use crate::{ path::{AddressPath, StoragePath}, storage::proofs::AccountProof, }; -use alloy_primitives::{map::{B256Map, HashMap}, StorageValue, B256}; +use alloy_primitives::{ + map::{B256Map, HashMap}, + StorageValue, B256, +}; use alloy_trie::{BranchNodeCompact, Nibbles}; pub use error::TransactionError; pub use manager::TransactionManager; @@ -88,7 +91,10 @@ impl, K: TransactionKind> Transaction { pub fn compute_root_with_overlay( &self, overlay_state: &OverlayState, - ) -> Result<(B256, HashMap, B256Map>), TransactionError> { + ) -> Result< + (B256, HashMap, B256Map>), + TransactionError, + > { self.database .storage_engine .compute_root_with_overlay(&self.context, overlay_state) From 82d7fe0200fed55c99a0ff6a69802eb22f8d38d9 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Fri, 15 Aug 2025 08:33:39 -0700 Subject: [PATCH 09/17] Iterative traversal for overlay root, fix zero storage edge case --- src/overlay.rs | 152 +-- src/storage.rs | 1 + src/storage/overlay_root.rs | 303 ++++- src/storage/overlay_root_iterative.rs | 1665 +++++++++++++++++++++++++ src/transaction.rs | 4 +- 5 files changed, 2005 insertions(+), 120 deletions(-) create mode 100644 src/storage/overlay_root_iterative.rs diff --git a/src/overlay.rs b/src/overlay.rs index a65d570f..23f259f5 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -3,7 +3,7 @@ use crate::{ node::{Node, TrieValue}, pointer::Pointer, }; -use alloy_primitives::{StorageValue, B256}; +use alloy_primitives::{StorageValue, B256, U256}; use alloy_trie::Nibbles; use std::sync::Arc; @@ -49,7 +49,16 @@ impl OverlayStateMut { 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. - self.changes.push((path, value)); + + 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. @@ -127,14 +136,15 @@ impl OverlayState { 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<(Nibbles, &Option)> { + pub fn get(&self, index: usize) -> Option<(&[u8], &Option)> { let slice = self.effective_slice(); if index < slice.len() { - Some((slice[index].0.slice(self.prefix_offset..), &slice[index].1)) + Some((&slice[index].0[self.prefix_offset..], &slice[index].1)) } else { None } @@ -158,7 +168,7 @@ impl OverlayState { /// 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: &Nibbles) -> Option<&Option> { + 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), @@ -169,7 +179,7 @@ impl OverlayState { /// Finds all entries that have the given prefix. /// Returns a zero-copy sub-slice of the overlay containing only matching entries. /// The prefix is compared against offset-adjusted paths. - pub fn find_prefix_range(&self, prefix: &Nibbles) -> OverlayState { + pub fn find_prefix_range(&self, prefix: &[u8]) -> OverlayState { if self.is_empty() { return OverlayState::empty(); } @@ -182,7 +192,9 @@ impl OverlayState { for (i, (path, _)) in slice.iter().enumerate() { if path.len() >= self.prefix_offset { let adjusted_path = &path.as_slice()[self.prefix_offset..]; - if adjusted_path.len() >= prefix.len() && adjusted_path[..prefix.len()] == *prefix { + if adjusted_path.len() >= prefix.len() && + adjusted_path[..prefix.len()] == prefix[..] + { if start_idx.is_none() { start_idx = Some(i); } @@ -253,15 +265,11 @@ impl OverlayState { /// 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: &Nibbles, - target_path: &Nibbles, - ) -> std::cmp::Ordering { + 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.slice(self.prefix_offset..); + let adjusted_path = &stored_path[self.prefix_offset..]; adjusted_path.cmp(target_path) } } @@ -317,48 +325,19 @@ impl OverlayState { (before, with_prefix, after) } - /// Creates a zero-copy sub-slice containing only changes that come before the given prefix. - /// The prefix comparison takes prefix_offset into account. - pub fn sub_slice_before_prefix(&self, prefix: &Nibbles) -> OverlayState { - let slice = self.effective_slice(); - let index = slice - .binary_search_by(|(p, _)| self.compare_with_offset(p, prefix)) - .unwrap_or_else(|i| i); // Insert position is exactly what we want for "before" - self.sub_slice(0, index) - } - - /// Creates a zero-copy sub-slice containing only changes that come strictly after the given - /// prefix. The prefix comparison takes prefix_offset into account. - pub fn sub_slice_after_prefix(&self, prefix: &Nibbles) -> OverlayState { - // Find the next key after prefix by incrementing prefix - let next_prefix = match prefix.increment() { - Some(next) => next, - None => { - // Prefix is all 0xF's, nothing can come after it - return OverlayState::empty(); - } - }; - - let slice = self.effective_slice(); - let index = slice - .binary_search_by(|(p, _)| self.compare_with_offset(p, &next_prefix)) - .unwrap_or_else(|i| i); - self.sub_slice(index, slice.len()) - } - /// 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: &Nibbles) -> OverlayState { + pub fn sub_slice_for_prefix(&self, prefix: &[u8]) -> OverlayState { self.find_prefix_range(prefix) } /// 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)> { + 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.slice(self.prefix_offset..); + let adjusted_path = &path.as_slice()[self.prefix_offset..]; Some((adjusted_path, value)) } else { None @@ -366,12 +345,9 @@ impl OverlayState { }) } - /// Checks if the overlay contains any changes that would affect the given path. - /// This includes exact matches and any changes to descendants of the path. - /// The path comparison takes prefix_offset into account. - pub fn affects_path(&self, path: &Nibbles) -> bool { + 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, _)| overlay_path.starts_with(path)) + self.iter().any(|(overlay_path, _)| path.starts_with(overlay_path)) } } @@ -577,7 +553,7 @@ mod tests { } #[test] - fn test_affects_path() { + fn test_contains_prefix_of() { let mut mutable = OverlayStateMut::new(); let account = test_account(); @@ -593,13 +569,20 @@ mod tests { let frozen = mutable.freeze(); // Test exact match - assert!(frozen.affects_path(&Nibbles::from_nibbles([1, 2, 3, 4]))); + 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.affects_path(&Nibbles::from_nibbles([1, 2]))); + 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.affects_path(&Nibbles::from_nibbles([7, 8, 9, 10]))); + assert!(!frozen.contains_prefix_of(&[7, 8, 9, 10])); } #[test] @@ -612,7 +595,7 @@ mod tests { let path = Nibbles::from_nibbles([1, 2, 3, 4]); assert!(frozen.lookup(&path).is_none()); - assert!(!frozen.affects_path(&path)); + assert!(!frozen.contains_prefix_of(&path)); let subset = frozen.find_prefix_range(&path); assert!(subset.is_empty()); @@ -756,7 +739,7 @@ mod tests { ]; for (i, (path, _)) in iter_results.iter().enumerate() { - assert_eq!(*path, expected_paths[i]); + assert_eq!(*path, expected_paths[i].as_slice()); } } @@ -785,20 +768,25 @@ mod tests { // Test sub_slice_before_prefix let split_point = Nibbles::from_nibbles([5, 0]); // Storage key [5, 0] - let before = storage_overlay.sub_slice_before_prefix(&split_point); - let after = storage_overlay.sub_slice_after_prefix(&split_point); + + 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(&Nibbles::from_nibbles([2, 0]))); + 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(&Nibbles::from_nibbles([5, 0]))); // Strictly after, so [5, 0] not included - assert!(after_paths.contains(&Nibbles::from_nibbles([8, 0]))); - assert!(after_paths.contains(&Nibbles::from_nibbles([9, 0]))); + 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] @@ -834,15 +822,11 @@ mod tests { // Now test with prefix_offset - should convert to storage-relative paths let storage_overlay = account1_storage.with_prefix_offset(4); - let storage_with_prefix5: Vec<_> = storage_overlay - .find_prefix_range(&Nibbles::from_nibbles([5])) - .iter() - .map(|(path, _)| path) - .collect(); + let storage_with_prefix5 = storage_overlay.find_prefix_range(&Nibbles::from_nibbles([5])); assert_eq!(storage_with_prefix5.len(), 2); - assert!(storage_with_prefix5.contains(&Nibbles::from_nibbles([5, 5]))); - assert!(storage_with_prefix5.contains(&Nibbles::from_nibbles([5, 6]))); + 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] @@ -967,24 +951,24 @@ mod tests { // 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(&Nibbles::from_nibbles([0, 5, 6, 7]))); - assert!(before_paths.contains(&Nibbles::from_nibbles([1, 0, 0, 0]))); - assert!(before_paths.contains(&Nibbles::from_nibbles([1, 1, 9, 9]))); + 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(&Nibbles::from_nibbles([1, 2]))); - assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([1, 2, 3, 4]))); - assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([1, 2, 5, 6]))); - assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([1, 2, 7, 8]))); + 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(&Nibbles::from_nibbles([1, 3, 0, 0]))); - assert!(after_paths.contains(&Nibbles::from_nibbles([2, 0, 0, 0]))); - assert!(after_paths.contains(&Nibbles::from_nibbles([5, 5, 5, 5]))); + 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()); @@ -1035,7 +1019,7 @@ mod tests { // 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)); + assert!(with_prefix_paths.contains(&prefix.as_slice())); } #[test] @@ -1071,17 +1055,17 @@ mod tests { // 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(&Nibbles::from_nibbles([2, 0, 0, 0]))); + 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(&Nibbles::from_nibbles([5, 0, 0, 0]))); - assert!(with_prefix_paths.contains(&Nibbles::from_nibbles([5, 1, 0, 0]))); + 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(&Nibbles::from_nibbles([8, 0, 0, 0]))); + assert!(after_paths.contains(&[8, 0, 0, 0].as_slice())); } } diff --git a/src/storage.rs b/src/storage.rs index 2638784e..a579ff53 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -1,6 +1,7 @@ pub mod debug; pub mod engine; pub mod overlay_root; +pub mod overlay_root_iterative; pub mod proofs; mod test_utils; pub mod value; diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 281a1304..1517c86c 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -70,7 +70,7 @@ impl StorageEngine { // println!("Updated branch nodes: {:?}", updated_branch_nodes); let root = hash_builder.root(); // This will clear the hash builder - // println!("Root: {:?}", root); + println!("Root: {:?}", root); println!("Stats: {:?}", stats); Ok((root, updated_branch_nodes, updated_storage_branch_nodes)) @@ -171,7 +171,7 @@ impl StorageEngine { NodeKind::Branch { .. } => { // Filter the overlay based on the path_prefix let (pre_overlay, with_prefix, post_overlay) = - overlay.sub_slice_by_prefix(path_prefix); + overlay.sub_slice_by_prefix(&path_prefix); if with_prefix.is_empty() { self.process_nonoverlapping_overlay_state( context, @@ -213,7 +213,7 @@ impl StorageEngine { updated_storage_branch_nodes, )?; } - Ok(()) + return Ok(()); } else { let (pre_overlay, with_prefix, post_overlay) = overlay.sub_slice_by_prefix(&full_path); @@ -242,7 +242,7 @@ impl StorageEngine { stats, updated_storage_branch_nodes, )?; - Ok(()) + return Ok(()); } } NodeKind::AccountLeaf { .. } | NodeKind::StorageLeaf { .. } => { @@ -282,7 +282,7 @@ impl StorageEngine { stats, updated_storage_branch_nodes, )?; - Ok(()) + return Ok(()); } } } @@ -295,10 +295,10 @@ impl StorageEngine { stats: &mut Stats, updated_storage_branch_nodes: &mut B256Map>, ) -> Result<(), Error> { - let mut last_processed_path: Option = None; + let mut last_processed_path: Option<&[u8]> = None; for (path, value) in overlay.iter() { if let Some(ref last_processed_path) = last_processed_path { - if path.has_prefix(last_processed_path) { + if path.starts_with(last_processed_path) { // skip over all descendants of a processed path continue; } @@ -319,13 +319,13 @@ impl StorageEngine { Some(OverlayValue::Storage(storage_value)) => { let rlp_encoded = self.encode_storage(storage_value)?; println!("Adding overlayed storage leaf: {:?}", path); - hash_builder.add_leaf(path.clone(), &rlp_encoded); + hash_builder.add_leaf(Nibbles::from_nibbles(path), &rlp_encoded); stats.storage.overlay_leaves += 1; last_processed_path = Some(path); } Some(OverlayValue::Hash(hash)) => { println!("Adding overlayed branch: {:?}", path); - hash_builder.add_branch(path.clone(), *hash, false); + hash_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); stats.account.overlay_branch_nodes += 1; last_processed_path = Some(path); } @@ -421,7 +421,7 @@ impl StorageEngine { // Create sub-slice of overlay for this child's subtree let child_overlay = overlay.sub_slice_for_prefix(&child_path); - children[i as usize] = (child_path, child_overlay, node.child(i).unwrap()); + children[i as usize] = (child_path, child_overlay, node.child(i as u8).unwrap()); } // This conservatively counts the minimum number of children that will exist. It may @@ -448,10 +448,10 @@ impl StorageEngine { // many lookups, hashes, and page reads. for (child_path, child_overlay, child_pointer) in children.iter() { if let Some((path, _)) = child_overlay.get(0) { - if &path == child_path { + if path == child_path.as_slice() { self.process_nonoverlapping_overlay_state( context, - child_overlay, + &child_overlay, hash_builder, stats, updated_storage_branch_nodes, @@ -473,11 +473,11 @@ impl StorageEngine { // Child is in the same page self.traverse_node_with_overlay( context, - child_overlay, + &child_overlay, hash_builder, slotted_page, child_cell_index, - child_path, + &child_path, child_hash, stats, updated_storage_branch_nodes, @@ -486,10 +486,10 @@ impl StorageEngine { // Child is in a different page self.traverse_page_with_overlay( context, - child_overlay, + &child_overlay, hash_builder, child_page_id, - child_path, + &child_path, child_hash, stats, updated_storage_branch_nodes, @@ -500,7 +500,7 @@ impl StorageEngine { // No child exists on disk, process any matching overlay changes self.process_nonoverlapping_overlay_state( context, - child_overlay, + &child_overlay, hash_builder, stats, updated_storage_branch_nodes, @@ -514,12 +514,12 @@ impl StorageEngine { // We need to traverse and add the children directly in order to be safe. for (child_path, child_overlay, child_pointer) in children.iter() { if let Some((path, _)) = child_overlay.get(0) { - if &path == child_path { + if path == child_path.as_slice() { // Direct overlay - the overlay path exactly matches this node's path // No need to traverse the child self.process_nonoverlapping_overlay_state( context, - child_overlay, + &child_overlay, hash_builder, stats, updated_storage_branch_nodes, @@ -537,11 +537,11 @@ impl StorageEngine { // Child is in the same page self.traverse_node_with_overlay( context, - child_overlay, + &child_overlay, hash_builder, slotted_page, child_cell_index, - child_path, + &child_path, child_hash, stats, updated_storage_branch_nodes, @@ -550,10 +550,10 @@ impl StorageEngine { // Child is in a different page self.traverse_page_with_overlay( context, - child_overlay, + &child_overlay, hash_builder, child_page_id, - child_path, + &child_path, child_hash, stats, updated_storage_branch_nodes, @@ -563,7 +563,7 @@ impl StorageEngine { // No child exists on disk, process any matching overlay changes self.process_nonoverlapping_overlay_state( context, - child_overlay, + &child_overlay, hash_builder, stats, updated_storage_branch_nodes, @@ -642,8 +642,8 @@ impl StorageEngine { .insert(B256::from_slice(&full_path.pack()), updated_branch_nodes); // Add the modified account to the main HashBuilder + println!("Adding modified account leaf: {full_path:?} with storage root: {new_storage_root:?}"); let rlp_encoded = self.encode_account_with_root(account, new_storage_root)?; - println!("Adding modified account leaf: {full_path:?}"); hash_builder.add_leaf(full_path.clone(), &rlp_encoded); } Ok(()) @@ -654,7 +654,7 @@ impl StorageEngine { context: &TransactionContext, overlay: &OverlayState, hash_builder: &mut HashBuilder, - full_path: &Nibbles, + full_path: &[u8], account: &Account, stats: &mut Stats, updated_storage_branch_nodes: &mut B256Map>, @@ -670,7 +670,7 @@ impl StorageEngine { // No storage overlays for this account, use original account let rlp_encoded = self.encode_account(account)?; // println!("Adding overlayed account leaf: {:?}", full_path); - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + hash_builder.add_leaf(Nibbles::from_nibbles(full_path), &rlp_encoded); stats.account.overlay_leaves += 1; return Ok(()); } @@ -691,12 +691,14 @@ impl StorageEngine { let (mut storage_hash_builder, updated_branch_nodes) = storage_hash_builder.split(); let new_storage_root = storage_hash_builder.root(); println!("New storage root: {new_storage_root:?}"); - updated_storage_branch_nodes - .insert(B256::from_slice(&full_path.pack()), updated_branch_nodes); + updated_storage_branch_nodes.insert( + B256::from_slice(&Nibbles::from_nibbles(full_path).pack()), + updated_branch_nodes, + ); // Add the modified account to the main HashBuilder let rlp_encoded = self.encode_account_with_root(account, new_storage_root)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + hash_builder.add_leaf(Nibbles::from_nibbles(full_path), &rlp_encoded); stats.account.overlay_leaves += 1; Ok(()) @@ -788,7 +790,7 @@ impl StorageEngine { /// Helper to encode TrieValue as RLP #[inline] - fn encode_trie_value(&self, trie_value: &TrieValue) -> Result, Error> { + pub fn encode_trie_value(&self, trie_value: &TrieValue) -> Result, Error> { let rlp_encoded = match trie_value { TrieValue::Account(account) => self.encode_account(account)?, TrieValue::Storage(storage_value) => self.encode_storage(storage_value)?, @@ -797,7 +799,7 @@ impl StorageEngine { } #[inline] - fn encode_account(&self, account: &Account) -> Result, Error> { + pub fn encode_account(&self, account: &Account) -> Result, Error> { let trie_account = TrieAccount { nonce: account.nonce, balance: account.balance, @@ -808,7 +810,11 @@ impl StorageEngine { } #[inline] - fn encode_account_with_root(&self, account: &Account, root: B256) -> Result, Error> { + pub fn encode_account_with_root( + &self, + account: &Account, + root: B256, + ) -> Result, Error> { let trie_account = TrieAccount { nonce: account.nonce, balance: account.balance, @@ -819,7 +825,7 @@ impl StorageEngine { } #[inline] - fn encode_storage(&self, storage_value: &U256) -> Result, Error> { + pub fn encode_storage(&self, storage_value: &U256) -> Result, Error> { Ok(encode(storage_value)) } } @@ -2286,4 +2292,233 @@ mod tests { let committed_root = context.root_node_hash; assert_eq!(committed_root, overlay_root); } + + #[test] + fn test_complex_case() {} + + // #[test] + // fn test_overlay_state_root_fuzz_rounds() { + // use crate::{path::AddressPath, path::StoragePath, Database}; + // use alloy_primitives::{Address, B256, U256}; + // use rand::{rngs::StdRng, Rng, SeedableRng}; + // use std::collections::{HashMap, HashSet}; + // use tempfile::TempDir; + + // // Deterministic RNG for reproducibility + // let mut rng = StdRng::seed_from_u64(0xDEADBEEFCAFEBABE); + + // // Create an empty DB + // let tmp_dir = TempDir::new().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(); + + // // In-memory expected state: AddressPath -> (Account, slot_hash -> value) + // let mut state: HashMap)> = HashMap::new(); + // let mut all_paths: Vec = Vec::new(); + // let mut used_paths: HashSet = HashSet::new(); + + // // Pre-seed 10,000 accounts + // let mut initial_changes: Vec<(Nibbles, Option)> = Vec::with_capacity(10_000); + // for i in 0..10_000u64 { + // let addr = Address::random(); + // let ap = AddressPath::for_address(addr); + // used_paths.insert(ap.clone()); + // all_paths.push(ap.clone()); + // let account = Account::new(i, U256::from(rng.random::()), EMPTY_ROOT_HASH, + // KECCAK_EMPTY); state.insert(ap.clone(), (account.clone(), HashMap::new())); + // initial_changes.push((ap.into(), Some(TrieValue::Account(account)))); + // } + + // // Of these, 1,000 should have storage: 100 with 10 slots, 900 with 1,000 slots + // let storage_accounts_small: Vec = + // all_paths.iter().cloned().take(100).collect(); let storage_accounts_large: + // Vec = all_paths.iter().cloned().skip(100).take(900).collect(); + + // // Helper to insert unique slots for an address + // let mut insert_slots = |ap: &AddressPath, num_slots: usize| { + // let (_, ref mut slots) = state.get_mut(ap).unwrap(); + // for _ in 0..num_slots { + // // Ensure unique slot per account + // let mut key: B256; + // loop { + // key = B256::from(rng.random::<[u8; 32]>()); + // if !slots.contains_key(&key) { + // break; + // } + // } + // let value = U256::from(rng.random::()); + // slots.insert(key, value); + // let storage_path = StoragePath::for_address_path_and_slot(ap.clone(), key); + // initial_changes.push((storage_path.full_path(), + // Some(TrieValue::Storage(value)))); } + // }; + + // for ap in storage_accounts_small.iter() { + // insert_slots(ap, 10); + // } + // for ap in storage_accounts_large.iter() { + // insert_slots(ap, 1_000); + // } + + // // Sort and commit initial state + // initial_changes.sort_by(|a, b| a.0.cmp(&b.0)); + // db.storage_engine.set_values(&mut context, &mut initial_changes).unwrap(); + + // // Accounts considered to have storage (non-empty map) + // let mut accounts_with_storage: Vec = state + // .iter() + // .filter_map(|(ap, (_acc, slots))| if !slots.is_empty() { Some(ap.clone()) } else { + // None }) .collect(); + + // // Helper to sample n unique elements from a candidate list without replacement + // let mut sample_n = |cands_len: usize, n: usize| -> Vec { + // let n = n.min(cands_len); + // let mut idxs: Vec = (0..cands_len).collect(); + // for i in 0..n { + // let rand_u = rng.random::() as usize; + // let j = i + (rand_u % (cands_len - i)); + // idxs.swap(i, j); + // } + // idxs.truncate(n); + // idxs + // }; + + // // 100 rounds of randomized operations + // for _round in 0..100 { + // let mut overlay_mut = OverlayStateMut::new(); + + // // Insert 50 new accounts + // let mut new_accounts: Vec = Vec::with_capacity(50); + // for i in 0..50u64 { + // let ap = loop { + // let a = Address::random(); + // let candidate = AddressPath::for_address(a); + // if !used_paths.contains(&candidate) { + // break candidate; + // } + // }; + // used_paths.insert(ap.clone()); + // new_accounts.push(ap.clone()); + // all_paths.push(ap.clone()); + // let account = Account::new(i, U256::from(rng.random::()), EMPTY_ROOT_HASH, + // KECCAK_EMPTY); state.insert(ap.clone(), (account.clone(), HashMap::new())); + // overlay_mut.insert(ap.into(), Some(OverlayValue::Account(account))); + // } + + // // Update 50 existing accounts + // let existing_paths: Vec = state.keys().cloned().collect(); + // let filtered: Vec = existing_paths + // .into_iter() + // .filter(|a| !new_accounts.contains(a)) + // .collect(); + // let idxs = sample_n(filtered.len(), 50); + // let update_targets: Vec = idxs.into_iter().map(|i| + // filtered[i].clone()).collect(); for ap in update_targets.iter().cloned() { + // if let Some((acc, _)) = state.get_mut(&ap) { + // acc.balance = U256::from(rng.random::()); + // acc.nonce = acc.nonce.wrapping_add(1); + // overlay_mut.insert(ap.into(), Some(OverlayValue::Account(acc.clone()))); + // } + // } + + // // Delete 10 accounts (avoid newly inserted and updated to keep sets disjoint) + // let delete_candidates: Vec = state + // .keys() + // .filter(|a| !new_accounts.contains(a) && !update_targets.contains(a)) + // .cloned() + // .collect(); + // let idxs = sample_n(delete_candidates.len(), 10); + // let delete_targets: Vec = idxs.into_iter().map(|i| + // delete_candidates[i].clone()).collect(); for ap in delete_targets.iter().cloned() + // { overlay_mut.insert(ap.clone().into(), None); + // state.remove(&ap); + // accounts_with_storage.retain(|a| *a != ap); + // } + + // // Storage operations + // // Helper: sample N accounts with at least K storage slots + // let mut sample_accounts_with_min_slots = |n: usize, k: usize| -> Vec { + // let candidates: Vec = accounts_with_storage + // .iter() + // .filter(|ap| state.get(*ap).map(|(_, m)| m.len()).unwrap_or(0) >= k) + // .cloned() + // .collect(); + // let idxs = sample_n(candidates.len(), n); + // idxs.into_iter().map(|i| candidates[i].clone()).collect() + // }; + + // // Insert 100 new slots into 10 random storage accounts + // let insert_slot_accounts = sample_accounts_with_min_slots(10, 0); + // for ap in insert_slot_accounts.iter().cloned() { + // if let Some((_, slots)) = state.get_mut(&ap) { + // for _ in 0..100 { + // let key = loop { + // let k = B256::from(rng.random::<[u8; 32]>()); + // if !slots.contains_key(&k) { + // break k; + // } + // }; + // let value = U256::from(rng.random::()); + // slots.insert(key, value); + // let storage_path = StoragePath::for_address_path_and_slot(ap.clone(), + // key); overlay_mut.insert(storage_path.full_path(), + // Some(OverlayValue::Storage(value))); } + // if !accounts_with_storage.contains(&ap) { + // accounts_with_storage.push(ap); + // } + // } + // } + + // // Update 10 slots in 10 random storage accounts (10 slots each) + // let update_slot_accounts = sample_accounts_with_min_slots(10, 10); + // for ap in update_slot_accounts.iter().cloned() { + // if let Some((_, map)) = state.get_mut(&ap) { + // let keys: Vec = map.keys().cloned().collect(); + // let idxs = sample_n(keys.len(), 10); + // for i in idxs { + // let key = keys[i]; + // let new_value = U256::from(rng.random::()); + // map.insert(key, new_value); + // let storage_path = StoragePath::for_address_path_and_slot(ap.clone(), + // key); overlay_mut + // .insert(storage_path.full_path(), + // Some(OverlayValue::Storage(new_value))); } + // } + // } + + // // Delete 10 slots in 10 random storage accounts (10 slots each) + // let delete_slot_accounts = sample_accounts_with_min_slots(10, 10); + // for ap in delete_slot_accounts.iter().cloned() { + // if let Some((_, map)) = state.get_mut(&ap) { + // let keys: Vec = map.keys().cloned().collect(); + // let idxs = sample_n(keys.len(), 10); + // for i in idxs { + // let key = keys[i]; + // map.remove(&key); + // let storage_path = StoragePath::for_address_path_and_slot(ap.clone(), + // key); overlay_mut.insert(storage_path.full_path(), None); + // } + // if map.is_empty() { + // accounts_with_storage.retain(|a| *a != ap); + // } + // } + // } + + // // Compute overlay root, apply, and compare + // let overlay = overlay_mut.freeze(); + // let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, + // &overlay).unwrap(); + + // let mut commit_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(&mut context, &mut commit_changes).unwrap(); + // let committed_root = context.root_node_hash; + + // assert_eq!(overlay_root, committed_root); + // } + // } } diff --git a/src/storage/overlay_root_iterative.rs b/src/storage/overlay_root_iterative.rs new file mode 100644 index 00000000..6c950041 --- /dev/null +++ b/src/storage/overlay_root_iterative.rs @@ -0,0 +1,1665 @@ +use std::sync::Arc; + +use crate::{ + account::Account, + context::TransactionContext, + node::{Node, NodeKind, TrieValue}, + overlay::{OverlayState, OverlayValue}, + page::SlottedPage, + pointer::Pointer, + storage::engine::{Error, StorageEngine}, +}; +use alloy_primitives::{ + map::{B256Map, HashMap}, + B256, +}; +use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles}; + +#[derive(Debug)] +enum TriePosition<'a> { + Node(Nibbles, Arc>, Node), + Pointer(Nibbles, Arc>, 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: Arc>, + overlay: OverlayState, + ) { + self.push(TriePosition::Node(path, page, node), overlay); + } + + fn push_pointer( + &mut self, + path: Nibbles, + pointer: Pointer, + page: Arc>, + can_add_by_hash: bool, + overlay: OverlayState, + ) { + self.push(TriePosition::Pointer(path, page, pointer, can_add_by_hash), overlay); + } + + fn push_none(&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() + } + + fn is_empty(&self) -> bool { + self.stack.is_empty() + } +} + +impl StorageEngine { + pub fn compute_state_root_with_overlay_iterative( + &self, + context: &TransactionContext, + overlay: OverlayState, + ) -> Result< + (B256, HashMap, B256Map>), + Error, + > { + if overlay.is_empty() { + return Ok((context.root_node_hash, HashMap::default(), B256Map::default())); + } + + let mut hash_builder = HashBuilder::default().with_updates(true); + let mut storage_branch_updates = B256Map::default(); + + 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_hash_builder( + &mut hash_builder, + &overlay, + &mut storage_branch_updates, + ); + let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); + return Ok((hash_builder.root(), updated_branch_nodes, B256Map::default())); + }; + + let root_node: Node = root_page.get_value(0)?; + let mut stack = TraversalStack::new(); + stack.push_node(root_node.prefix().clone(), root_node, Arc::new(root_page), overlay); + + self.compute_root_with_overlay_iterative( + context, + &mut stack, + &mut hash_builder, + &mut storage_branch_updates, + )?; + + let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); + Ok((hash_builder.root(), updated_branch_nodes, storage_branch_updates)) + } + + fn compute_root_with_overlay_iterative<'a>( + &'a self, + context: &TransactionContext, + stack: &mut TraversalStack<'a>, + hash_builder: &mut HashBuilder, + storage_branch_updates: &mut B256Map>, + ) -> 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 !stack.is_empty() { + let (position, overlay) = stack.pop().unwrap(); + println!("Processing position: {:?}", position); + + match position { + TriePosition::None => { + // No trie position, process whatever is in the overlay + self.add_overlay_to_hash_builder( + hash_builder, + &overlay, + storage_branch_updates, + ); + } + TriePosition::Pointer(path, page, pointer, can_add_by_hash) => { + if overlay.is_empty() && can_add_by_hash && pointer.rlp().as_hash().is_some() { + // No overlay, just add the pointer by hash + println!("Adding pointer: {:?}", path); + hash_builder.add_branch(path, pointer.rlp().as_hash().unwrap(), true); + } else { + // We have an overlay, need to process the child + self.process_overlayed_child( + context, + overlay, + hash_builder, + &path, + &pointer, + page, + stack, + storage_branch_updates, + )?; + } + } + TriePosition::Node(path, page, node) => { + // First check if the node is invalidated by the overlay + let first_overlay_item = overlay.iter().next(); + if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = first_overlay_item { + if path.has_prefix(overlay_path) { + // the overlay invalidates the current node, so just add this and skip + // the rest of the db traversal + self.add_overlay_to_hash_builder( + hash_builder, + &overlay, + storage_branch_updates, + ); + continue; + } + } + + // Otherwise, process the overlay and node + let (pre_overlay, overlay, post_overlay) = overlay.sub_slice_by_prefix(&path); + self.add_overlay_to_hash_builder( + hash_builder, + &pre_overlay, + storage_branch_updates, + ); + // Defer the post_overlay to be processed after the node is traversed + stack.push_none(post_overlay); + + if pre_overlay.contains_prefix_of(&path) { + // A prefix of the node has already been processed, so we can skip the rest + // of the db traversal + continue; + } + + match node.kind() { + NodeKind::Branch { .. } => { + let first_overlay_item = overlay.iter().next(); + if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = + first_overlay_item + { + 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_hash_builder( + hash_builder, + &overlay, + storage_branch_updates, + ); + continue; + } + } + self.process_branch_node_with_overlay( + &overlay, &path, &node, page, stack, + )?; + } + NodeKind::AccountLeaf { .. } => { + self.process_account_leaf_with_overlay( + context, + &overlay, + hash_builder, + path, + &node, + page, + storage_branch_updates, + )?; + } + NodeKind::StorageLeaf { .. } => { + let first_overlay_item = overlay.iter().next(); + if let Some((overlay_path, _)) = first_overlay_item { + 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_hash_builder( + hash_builder, + &overlay, + storage_branch_updates, + ); + continue; + } + } + // Leaf node, add it to the hash builder + let encoded = self.encode_trie_value(&node.value().unwrap())?; + println!("Adding storage leaf: {:?}", path); + hash_builder.add_leaf(path, &encoded); + } + } + } + } + } + Ok(()) + } + + fn process_branch_node_with_overlay<'a>( + &'a self, + overlay: &OverlayState, + path: &Nibbles, + node: &Node, + current_page: Arc>, + stack: &mut TraversalStack<'a>, + ) -> Result<(), Error> { + let mut child_data = Vec::with_capacity(16); + + let mut minimum_possible_child_count = 0; + for i in 0..16 { + let mut child_path = path.clone(); + child_path.push(i); + let child_overlay = overlay.sub_slice_for_prefix(&child_path); + let child_pointer = node.child(i as u8).unwrap(); + + if child_pointer.is_some() && child_overlay.is_empty() { + minimum_possible_child_count += 1; + } else { + match child_overlay.get(0) { + Some((_, Some(_))) => { + // 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)); + } + let can_add_by_hash = minimum_possible_child_count > 1; + + println!("Iterating over children of branch: {:?}", path); + for (child_path, child_pointer, child_overlay) in child_data.into_iter().rev() { + match child_pointer { + Some(pointer) => { + stack.push_pointer( + child_path, + pointer.clone(), + 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_none(child_overlay); + } + } + } + } + Ok(()) + } + + fn process_account_leaf_with_overlay<'a>( + &'a self, + context: &TransactionContext, + overlay: &OverlayState, + hash_builder: &mut HashBuilder, + path: Nibbles, + node: &Node, + current_page: Arc>, + storage_branch_updates: &mut B256Map>, + ) -> Result<(), Error> { + let original_account = match node.value() { + Ok(TrieValue::Account(account)) => account, + _ => panic!("node is not an account leaf"), + }; + let storage_root = match node.kind() { + NodeKind::AccountLeaf { storage_root, .. } => storage_root, + _ => panic!("node is not an account leaf"), + }; + + let overlayed_account = overlay.lookup(&path); + let account = match overlayed_account { + Some(None) => { + // The account is removed in the overlay + println!("Not adding removed account: {:?}", path); + return Ok(()); + } + Some(Some(OverlayValue::Account(overlayed_account))) => { + // The account is updated in the overlay + overlayed_account + } + _ => { + // The account is not updated in the overlay + &original_account + } + }; + + let has_storage_overlays = overlay.iter().any(|(path, _)| path.len() > 64); + if !has_storage_overlays { + let rlp_encoded = + self.encode_account_with_root(account, original_account.storage_root)?; + println!("Adding account leaf with no storage overlays: {:?}", path); + hash_builder.add_leaf(path, &rlp_encoded); + return Ok(()); + } + + let mut storage_hash_builder = HashBuilder::default().with_updates(true); + + // We have storage overlays, need to compute a new storage root + let storage_overlay = overlay.with_prefix_offset(64); + + match storage_root { + Some(pointer) => { + println!("Processing overlayed storage root for: {:?}", path); + 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_iterative( + context, + &mut storage_stack, + &mut storage_hash_builder, + storage_branch_updates, + )? + } 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, + Arc::new(slotted_page), + storage_overlay, + ); + self.compute_root_with_overlay_iterative( + context, + &mut storage_stack, + &mut storage_hash_builder, + storage_branch_updates, + )?; + } + } + None => { + // No existing storage root, just add the overlay + self.add_overlay_to_hash_builder( + &mut storage_hash_builder, + &storage_overlay, + storage_branch_updates, + ); + } + }; + let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_hash_builder.split(); + let new_root = storage_hash_builder.root(); + println!("New root: {:?}", new_root); + + storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); + + let encoded = self.encode_account_with_root(account, new_root).unwrap(); + println!("Adding overlayed account leaf: {:?}", path); + + hash_builder.add_leaf(path, &encoded); + + Ok(()) + } + + fn process_overlayed_child<'a>( + &'a self, + context: &TransactionContext, + overlay: OverlayState, + hash_builder: &mut HashBuilder, + child_path: &Nibbles, + child: &Pointer, + current_page: Arc>, + stack: &mut TraversalStack<'a>, + storage_branch_updates: &mut B256Map>, + ) -> 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. + let first_overlay_item = overlay.iter().next(); + if let Some((overlay_path, overlay_value)) = first_overlay_item { + if child_path == overlay_path && + (matches!(overlay_value, Some(OverlayValue::Hash(_))) || overlay_value.is_none()) + { + // the child path is directly overlayed, so only use the overlay state + self.add_overlay_to_hash_builder(hash_builder, &overlay, storage_branch_updates); + return Ok(()); + } + } + + if let Some(child_cell) = child.location().cell_index() { + let child_node: Node = current_page.get_value(child_cell)?; + stack.push_node( + child_path.join(child_node.prefix()), + 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)?; + stack.push_node( + child_path.join(child_node.prefix()), + child_node, + Arc::new(child_slotted_page), + overlay, + ); + } + Ok(()) + } + + fn process_overlayed_account( + &self, + hash_builder: &mut HashBuilder, + path: Nibbles, + account: &Account, + storage_overlay: OverlayState, + storage_branch_updates: &mut B256Map>, + ) -> Result<(), Error> { + if storage_overlay.is_empty() { + let encoded = self.encode_account(account).unwrap(); + println!("Adding overlayed account leaf with no storage overlays: {:?}", path); + hash_builder.add_leaf(path, &encoded); + return Ok(()); + } + + let mut storage_hash_builder = HashBuilder::default().with_updates(true); + self.add_overlay_to_hash_builder( + &mut storage_hash_builder, + &storage_overlay, + storage_branch_updates, + ); + + let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_hash_builder.split(); + let storage_root = storage_hash_builder.root(); + + println!("Updated storage branch nodes: {:?}", updated_storage_branch_nodes); + + storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); + + let encoded = self.encode_account_with_root(account, storage_root).unwrap(); + println!("Adding overlayed account leaf with storage overlays: {:?}", path); + hash_builder.add_leaf(path, &encoded); + Ok(()) + } + + fn add_overlay_to_hash_builder( + &self, + hash_builder: &mut HashBuilder, + overlay: &OverlayState, + storage_branch_updates: &mut B256Map>, + ) { + let mut last_processed_path: Option<&[u8]> = None; + for (path, value) in overlay.iter() { + if let Some(ref 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(&Nibbles::from_nibbles(path)) + .with_prefix_offset(64); + self.process_overlayed_account( + hash_builder, + Nibbles::from_nibbles(path), + &account, + storage_overlay, + storage_branch_updates, + ) + .unwrap(); + last_processed_path = Some(path); + } + Some(OverlayValue::Storage(storage_value)) => { + let encoded = self.encode_storage(&storage_value).unwrap(); + println!("Adding overlayed storage leaf: {:?}", path); + hash_builder.add_leaf(Nibbles::from_nibbles(path), &encoded); + } + Some(OverlayValue::Hash(hash)) => { + println!("Adding overlayed branch node: {:?}", path); + hash_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); + last_processed_path = Some(path); + } + None => { + // Tombstone - skip + println!("Skipping tombstone: {:?}", path); + last_processed_path = Some(path); + } + } + } + } +} + +#[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 (overlay_root, account_branch_updates, storage_branch_updates) = db + .storage_engine + .compute_state_root_with_overlay_iterative(context, overlay.clone()) + .unwrap(); + assert_ne!(overlay_root, initial_root, "Overlay should not match initial root"); + + println!("Account branch updates: {:?}", account_branch_updates); + println!("Storage branch updates: {:?}", storage_branch_updates); + + 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 (overlay_root_with_branches, _, _) = db + .storage_engine + .compute_state_root_with_overlay_iterative(context, overlay_with_branches.clone()) + .unwrap(); + 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 (overlay_root_after_commit, _, _) = db + .storage_engine + .compute_state_root_with_overlay_iterative(context, overlay_with_branches) + .unwrap(); + 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 (root, _, _) = db + .storage_engine + .compute_state_root_with_overlay_iterative(&context, empty_overlay) + .unwrap(); + assert_eq!(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 (root, _, _) = + db.storage_engine.compute_state_root_with_overlay_iterative(&context, overlay).unwrap(); + + // The root should be different from the empty root (since we have changes) + assert_ne!(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()); + + // println!("Storage path 1: {:?}", storage_path1.full_path()); + + // 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()); + + // println!("Storage path 2: {:?}", storage_path2.full_path()); + + 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); + } +} diff --git a/src/transaction.rs b/src/transaction.rs index 886a895b..28686578 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -90,14 +90,14 @@ impl, K: TransactionKind> Transaction { pub fn compute_root_with_overlay( &self, - overlay_state: &OverlayState, + overlay_state: OverlayState, ) -> Result< (B256, HashMap, B256Map>), TransactionError, > { self.database .storage_engine - .compute_root_with_overlay(&self.context, overlay_state) + .compute_state_root_with_overlay_iterative(&self.context, overlay_state) .map_err(|_| TransactionError::Generic) } From 241257a1cbc972ecd076cf4d0ef29b8482c13963 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Tue, 5 Aug 2025 17:12:17 -0700 Subject: [PATCH 10/17] Faster overlay root, less RLP overhead --- src/account.rs | 13 ++- src/node.rs | 46 ++++++-- src/overlay.rs | 5 + src/page/manager/mmap.rs | 2 +- src/storage/overlay_root.rs | 23 ++-- src/storage/overlay_root_iterative.rs | 162 ++++++++++++++------------ 6 files changed, 151 insertions(+), 100 deletions(-) diff --git a/src/account.rs b/src/account.rs index e4937ba3..44f683b2 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/node.rs b/src/node.rs index 86e7d52d..30dc6772 100644 --- a/src/node.rs +++ b/src/node.rs @@ -152,6 +152,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 +571,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 +638,26 @@ impl Encodable for Node { } } +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 +} + 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 index 23f259f5..089db40e 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -150,6 +150,11 @@ impl OverlayState { } } + #[inline] + pub fn first(&self) -> Option<(&[u8], &Option)> { + self.get(0) + } + /// 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] diff --git a/src/page/manager/mmap.rs b/src/page/manager/mmap.rs index b68e9a69..41d62087 100644 --- a/src/page/manager/mmap.rs +++ b/src/page/manager/mmap.rs @@ -268,7 +268,7 @@ impl PageManager { /// Syncs pages to the backing file. pub fn sync(&self) -> io::Result<()> { if cfg!(not(miri)) { - self.mmap.flush_async() + self.mmap.flush() } else { Ok(()) } diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 1517c86c..924489bd 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -11,8 +11,9 @@ use alloy_primitives::{ map::{B256Map, HashMap}, B256, U256, }; -use alloy_rlp::encode; +use alloy_rlp::{encode, encode_fixed_size}; use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, TrieAccount}; +use arrayvec::ArrayVec; #[derive(Debug, Default)] struct Stats { @@ -792,21 +793,21 @@ impl StorageEngine { #[inline] pub fn encode_trie_value(&self, trie_value: &TrieValue) -> Result, Error> { let rlp_encoded = match trie_value { - TrieValue::Account(account) => self.encode_account(account)?, - TrieValue::Storage(storage_value) => self.encode_storage(storage_value)?, + TrieValue::Account(account) => self.encode_account(account)?.to_vec(), + TrieValue::Storage(storage_value) => self.encode_storage(storage_value)?.to_vec(), }; Ok(rlp_encoded) } #[inline] - pub fn encode_account(&self, account: &Account) -> Result, Error> { - let trie_account = TrieAccount { + pub fn encode_account(&self, account: &Account) -> Result, Error> { + let trie_account = Account { nonce: account.nonce, balance: account.balance, storage_root: account.storage_root, code_hash: account.code_hash, }; - Ok(encode(trie_account)) + Ok(encode_fixed_size(&trie_account)) } #[inline] @@ -814,19 +815,19 @@ impl StorageEngine { &self, account: &Account, root: B256, - ) -> Result, Error> { - let trie_account = TrieAccount { + ) -> Result, Error> { + let trie_account = Account { nonce: account.nonce, balance: account.balance, storage_root: root, code_hash: account.code_hash, }; - Ok(encode(trie_account)) + Ok(encode_fixed_size(&trie_account)) } #[inline] - pub fn encode_storage(&self, storage_value: &U256) -> Result, Error> { - Ok(encode(storage_value)) + pub fn encode_storage(&self, storage_value: &U256) -> Result, Error> { + Ok(encode_fixed_size(storage_value)) } } diff --git a/src/storage/overlay_root_iterative.rs b/src/storage/overlay_root_iterative.rs index 6c950041..94aabf71 100644 --- a/src/storage/overlay_root_iterative.rs +++ b/src/storage/overlay_root_iterative.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use crate::{ account::Account, context::TransactionContext, - node::{Node, NodeKind, TrieValue}, + node::{encode_account_leaf, Node, NodeKind, TrieValue}, overlay::{OverlayState, OverlayValue}, page::SlottedPage, pointer::Pointer, @@ -13,7 +13,8 @@ use alloy_primitives::{ map::{B256Map, HashMap}, B256, }; -use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles}; +use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, EMPTY_ROOT_HASH}; +use arrayvec::ArrayVec; #[derive(Debug)] enum TriePosition<'a> { @@ -124,10 +125,7 @@ impl StorageEngine { // 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 !stack.is_empty() { - let (position, overlay) = stack.pop().unwrap(); - println!("Processing position: {:?}", position); - + while let Some((position, overlay)) = stack.pop() { match position { TriePosition::None => { // No trie position, process whatever is in the overlay @@ -140,7 +138,7 @@ impl StorageEngine { TriePosition::Pointer(path, page, pointer, can_add_by_hash) => { if overlay.is_empty() && can_add_by_hash && pointer.rlp().as_hash().is_some() { // No overlay, just add the pointer by hash - println!("Adding pointer: {:?}", path); + // println!("Adding pointer: {:?}", path); hash_builder.add_branch(path, pointer.rlp().as_hash().unwrap(), true); } else { // We have an overlay, need to process the child @@ -157,22 +155,6 @@ impl StorageEngine { } } TriePosition::Node(path, page, node) => { - // First check if the node is invalidated by the overlay - let first_overlay_item = overlay.iter().next(); - if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = first_overlay_item { - if path.has_prefix(overlay_path) { - // the overlay invalidates the current node, so just add this and skip - // the rest of the db traversal - self.add_overlay_to_hash_builder( - hash_builder, - &overlay, - storage_branch_updates, - ); - continue; - } - } - - // Otherwise, process the overlay and node let (pre_overlay, overlay, post_overlay) = overlay.sub_slice_by_prefix(&path); self.add_overlay_to_hash_builder( hash_builder, @@ -188,12 +170,9 @@ impl StorageEngine { continue; } - match node.kind() { - NodeKind::Branch { .. } => { - let first_overlay_item = overlay.iter().next(); - if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = - first_overlay_item - { + match node.into_kind() { + NodeKind::Branch { children } => { + if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = overlay.first() { if overlay_path == &path { // the overlay invalidates the current node, so just add this // and skip the rest of the db traversal @@ -206,23 +185,25 @@ impl StorageEngine { } } self.process_branch_node_with_overlay( - &overlay, &path, &node, page, stack, + &overlay, &path, children, page, stack, )?; } - NodeKind::AccountLeaf { .. } => { + NodeKind::AccountLeaf { nonce_rlp, balance_rlp, code_hash, storage_root } => { self.process_account_leaf_with_overlay( context, &overlay, hash_builder, path, - &node, page, storage_branch_updates, + nonce_rlp, + balance_rlp, + code_hash, + storage_root, )?; } - NodeKind::StorageLeaf { .. } => { - let first_overlay_item = overlay.iter().next(); - if let Some((overlay_path, _)) = first_overlay_item { + NodeKind::StorageLeaf { value_rlp } => { + if let Some((overlay_path, _)) = overlay.first() { if overlay_path == &path { // the overlay invalidates the current node, so just add this // and skip the rest of the db traversal @@ -235,9 +216,8 @@ impl StorageEngine { } } // Leaf node, add it to the hash builder - let encoded = self.encode_trie_value(&node.value().unwrap())?; - println!("Adding storage leaf: {:?}", path); - hash_builder.add_leaf(path, &encoded); + // println!("Adding storage leaf: {:?}", path); + hash_builder.add_leaf(path, &value_rlp); } } } @@ -246,11 +226,12 @@ impl StorageEngine { Ok(()) } + #[inline(never)] fn process_branch_node_with_overlay<'a>( &'a self, overlay: &OverlayState, path: &Nibbles, - node: &Node, + mut children: [Option; 16], current_page: Arc>, stack: &mut TraversalStack<'a>, ) -> Result<(), Error> { @@ -261,7 +242,7 @@ impl StorageEngine { let mut child_path = path.clone(); child_path.push(i); let child_overlay = overlay.sub_slice_for_prefix(&child_path); - let child_pointer = node.child(i as u8).unwrap(); + let child_pointer = children[i as usize].take(); if child_pointer.is_some() && child_overlay.is_empty() { minimum_possible_child_count += 1; @@ -280,13 +261,13 @@ impl StorageEngine { } let can_add_by_hash = minimum_possible_child_count > 1; - println!("Iterating over children of branch: {:?}", path); + // println!("Iterating over children of branch: {:?}", path); for (child_path, child_pointer, child_overlay) in child_data.into_iter().rev() { match child_pointer { Some(pointer) => { stack.push_pointer( child_path, - pointer.clone(), + pointer, current_page.clone(), can_add_by_hash, child_overlay, @@ -305,48 +286,51 @@ impl StorageEngine { Ok(()) } + #[inline(never)] fn process_account_leaf_with_overlay<'a>( &'a self, context: &TransactionContext, overlay: &OverlayState, hash_builder: &mut HashBuilder, path: Nibbles, - node: &Node, current_page: Arc>, storage_branch_updates: &mut B256Map>, + mut nonce_rlp: ArrayVec, + mut balance_rlp: ArrayVec, + mut code_hash: B256, + storage_root: Option, ) -> Result<(), Error> { - let original_account = match node.value() { - Ok(TrieValue::Account(account)) => account, - _ => panic!("node is not an account leaf"), - }; - let storage_root = match node.kind() { - NodeKind::AccountLeaf { storage_root, .. } => storage_root, - _ => panic!("node is not an account leaf"), - }; - let overlayed_account = overlay.lookup(&path); - let account = match overlayed_account { + match overlayed_account { Some(None) => { // The account is removed in the overlay - println!("Not adding removed account: {:?}", path); + // println!("Not adding removed account: {:?}", path); return Ok(()); } Some(Some(OverlayValue::Account(overlayed_account))) => { // The account is updated in the overlay - overlayed_account + 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 - &original_account } }; let has_storage_overlays = overlay.iter().any(|(path, _)| path.len() > 64); if !has_storage_overlays { - let rlp_encoded = - self.encode_account_with_root(account, original_account.storage_root)?; - println!("Adding account leaf with no storage overlays: {:?}", path); - hash_builder.add_leaf(path, &rlp_encoded); + // println!("Adding account leaf with no storage overlays: {:?}", path); + 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_hash_builder( + hash_builder, + path, + &nonce_rlp, + &balance_rlp, + &code_hash, + &storage_root_hash, + ); return Ok(()); } @@ -357,7 +341,7 @@ impl StorageEngine { match storage_root { Some(pointer) => { - println!("Processing overlayed storage root for: {:?}", path); + // println!("Processing overlayed storage root for: {:?}", path); let mut storage_stack = TraversalStack::new(); // load the root storage node @@ -405,18 +389,42 @@ impl StorageEngine { }; let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_hash_builder.split(); let new_root = storage_hash_builder.root(); - println!("New root: {:?}", new_root); + // println!("New root: {:?}", new_root); storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); - let encoded = self.encode_account_with_root(account, new_root).unwrap(); - println!("Adding overlayed account leaf: {:?}", path); + // println!("Adding overlayed account leaf: {:?}", path); - hash_builder.add_leaf(path, &encoded); + self.add_account_leaf_to_hash_builder( + hash_builder, + path, + &nonce_rlp, + &balance_rlp, + &code_hash, + &new_root, + ); Ok(()) } + #[inline(never)] + fn add_account_leaf_to_hash_builder( + &self, + hash_builder: &mut HashBuilder, + 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); + hash_builder.add_leaf(path, &buf[..account_rlp_length]); + } + + #[inline(never)] fn process_overlayed_child<'a>( &'a self, context: &TransactionContext, @@ -431,10 +439,10 @@ impl StorageEngine { // 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. - let first_overlay_item = overlay.iter().next(); - if let Some((overlay_path, overlay_value)) = first_overlay_item { - if child_path == overlay_path && - (matches!(overlay_value, Some(OverlayValue::Hash(_))) || overlay_value.is_none()) + // 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_hash_builder(hash_builder, &overlay, storage_branch_updates); @@ -465,6 +473,7 @@ impl StorageEngine { Ok(()) } + #[inline(never)] fn process_overlayed_account( &self, hash_builder: &mut HashBuilder, @@ -475,7 +484,7 @@ impl StorageEngine { ) -> Result<(), Error> { if storage_overlay.is_empty() { let encoded = self.encode_account(account).unwrap(); - println!("Adding overlayed account leaf with no storage overlays: {:?}", path); + // println!("Adding overlayed account leaf with no storage overlays: {:?}", path); hash_builder.add_leaf(path, &encoded); return Ok(()); } @@ -490,16 +499,17 @@ impl StorageEngine { let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_hash_builder.split(); let storage_root = storage_hash_builder.root(); - println!("Updated storage branch nodes: {:?}", updated_storage_branch_nodes); + // println!("Updated storage branch nodes: {:?}", updated_storage_branch_nodes); storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); let encoded = self.encode_account_with_root(account, storage_root).unwrap(); - println!("Adding overlayed account leaf with storage overlays: {:?}", path); + // println!("Adding overlayed account leaf with storage overlays: {:?}", path); hash_builder.add_leaf(path, &encoded); Ok(()) } + #[inline(never)] fn add_overlay_to_hash_builder( &self, hash_builder: &mut HashBuilder, @@ -518,7 +528,7 @@ impl StorageEngine { match value { Some(OverlayValue::Account(account)) => { let storage_overlay = overlay - .sub_slice_for_prefix(&Nibbles::from_nibbles(path)) + .sub_slice_for_prefix(path) .with_prefix_offset(64); self.process_overlayed_account( hash_builder, @@ -532,17 +542,17 @@ impl StorageEngine { } Some(OverlayValue::Storage(storage_value)) => { let encoded = self.encode_storage(&storage_value).unwrap(); - println!("Adding overlayed storage leaf: {:?}", path); + // println!("Adding overlayed storage leaf: {:?}", path); hash_builder.add_leaf(Nibbles::from_nibbles(path), &encoded); } Some(OverlayValue::Hash(hash)) => { - println!("Adding overlayed branch node: {:?}", path); + // println!("Adding overlayed branch node: {:?}", path); hash_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); last_processed_path = Some(path); } None => { // Tombstone - skip - println!("Skipping tombstone: {:?}", path); + // println!("Skipping tombstone: {:?}", path); last_processed_path = Some(path); } } From 2194a8b33d7a328b182ecd7ce421c9a342aae4fd Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Fri, 15 Aug 2025 16:05:42 -0700 Subject: [PATCH 11/17] Improve overlay performance, particularly for prefix search --- benches/crud_benchmarks.rs | 26 +-- src/overlay.rs | 291 +++++--------------------- src/storage/overlay_root.rs | 4 +- src/storage/overlay_root_iterative.rs | 113 +++++----- 4 files changed, 119 insertions(+), 315 deletions(-) diff --git a/benches/crud_benchmarks.rs b/benches/crud_benchmarks.rs index 7454e2ac..4273ebc4 100644 --- a/benches/crud_benchmarks.rs +++ b/benches/crud_benchmarks.rs @@ -242,7 +242,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(); @@ -596,15 +596,13 @@ fn bench_state_root_with_overlay(c: &mut Criterion) { 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 total_storage_per_address = DEFAULT_SETUP_DB_STORAGE_PER_CONTRACT; - let total_addresses = BATCH_SIZE; let addresses: Vec = - (0..total_addresses).map(|_| generate_random_address(&mut rng)).collect(); - // let storage_paths_values = generate_storage_paths_values(&addresses, - // total_storage_per_address); + (0..BATCH_SIZE).map(|_| generate_random_address(&mut rng)).collect(); let mut account_overlay_mut = OverlayStateMut::new(); addresses.iter().enumerate().for_each(|(i, addr)| { @@ -614,25 +612,11 @@ fn bench_state_root_with_overlay(c: &mut Criterion) { }); let account_overlay = account_overlay_mut.freeze(); - // Build overlay state from storage paths and values - // let mut storage_overlay_mut = OverlayStateMut::new(); - - // for (storage_path, storage_value) in &storage_paths_values { - // // Convert storage path to nibbles for overlay - // let nibbles = storage_path.full_path(); - // storage_overlay_mut.insert(nibbles, Some(OverlayValue::Storage(*storage_value))); - // } - - // // Freeze the mutable overlay to get an immutable one - // let storage_overlay = storage_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 dir = TempDir::new("triedb_bench_state_root_with_account_overlay").unwrap(); - copy_files(&base_dir, dir.path()).unwrap(); let db_path = dir.path().join(&file_name); Database::open(db_path).unwrap() }, @@ -640,7 +624,7 @@ fn bench_state_root_with_overlay(c: &mut Criterion) { let tx = db.begin_ro().unwrap(); // Compute the root hash with the overlay - let _root_result = tx.compute_root_with_overlay(&account_overlay).unwrap(); + let _root_result = tx.compute_root_with_overlay(account_overlay.clone()).unwrap(); tx.commit().unwrap(); }, diff --git a/src/overlay.rs b/src/overlay.rs index 089db40e..3a02e71e 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -1,8 +1,4 @@ -use crate::{ - account::Account, - node::{Node, TrieValue}, - pointer::Pointer, -}; +use crate::{account::Account, node::TrieValue}; use alloy_primitives::{StorageValue, B256, U256}; use alloy_trie::Nibbles; use std::sync::Arc; @@ -152,7 +148,12 @@ impl OverlayState { #[inline] pub fn first(&self) -> Option<(&[u8], &Option)> { - self.get(0) + if self.len() == 0 { + 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. @@ -181,47 +182,6 @@ impl OverlayState { } } - /// Finds all entries that have the given prefix. - /// Returns a zero-copy sub-slice of the overlay containing only matching entries. - /// The prefix is compared against offset-adjusted paths. - pub fn find_prefix_range(&self, prefix: &[u8]) -> OverlayState { - if self.is_empty() { - return OverlayState::empty(); - } - - let slice = self.effective_slice(); - let mut start_idx = None; - let mut end_idx = slice.len(); - - // Find the range of entries that start with the prefix after applying offset - for (i, (path, _)) in slice.iter().enumerate() { - if path.len() >= self.prefix_offset { - let adjusted_path = &path.as_slice()[self.prefix_offset..]; - if adjusted_path.len() >= prefix.len() && - adjusted_path[..prefix.len()] == prefix[..] - { - if start_idx.is_none() { - start_idx = Some(i); - } - } else if start_idx.is_some() { - // We've found the end of the matching range - end_idx = i; - break; - } - } - } - - match start_idx { - Some(start) => OverlayState { - data: Arc::clone(&self.data), - start_idx: self.start_idx + start, - end_idx: self.start_idx + end_idx, - prefix_offset: self.prefix_offset, - }, - None => OverlayState::empty(), - } - } - /// 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 @@ -268,6 +228,34 @@ impl OverlayState { } } + /// 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[..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 { @@ -305,24 +293,11 @@ impl OverlayState { /// - 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: &Nibbles, - ) -> (OverlayState, OverlayState, OverlayState) { + pub fn sub_slice_by_prefix(&self, prefix: &[u8]) -> (OverlayState, OverlayState, OverlayState) { let slice = self.effective_slice(); - let start_index = slice - .binary_search_by(|(p, _)| self.compare_with_offset(p, prefix)) - .unwrap_or_else(|i| i); - let end_index = match prefix.increment() { - Some(next) => slice - .binary_search_by(|(p, _)| self.compare_with_offset(p, &next)) - .unwrap_or_else(|i| i), - None => { - // Prefix is all 0xF's, nothing can come after it - slice.len() - } - }; + 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); @@ -334,7 +309,16 @@ impl OverlayState { /// 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 { - self.find_prefix_range(prefix) + 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. @@ -356,97 +340,10 @@ impl OverlayState { } } -/// Pointer that can reference either on-disk nodes or in-memory overlay nodes. -/// This allows the trie traversal to seamlessly handle both persisted and overlayed state. -#[derive(Debug, Clone)] -pub enum OverlayPointer { - /// Reference to a node stored on disk - OnDisk(Pointer), - /// Reference to a node stored in memory as part of the overlay - InMemory(Box), -} - -impl OverlayPointer { - /// Creates a new on-disk overlay pointer. - pub fn on_disk(pointer: Pointer) -> Self { - Self::OnDisk(pointer) - } - - /// Creates a new in-memory overlay pointer. - pub fn in_memory(node: Node) -> Self { - Self::InMemory(Box::new(node)) - } - - /// Returns true if this pointer references an on-disk node. - pub fn is_on_disk(&self) -> bool { - matches!(self, Self::OnDisk(_)) - } - - /// Returns true if this pointer references an in-memory node. - pub fn is_in_memory(&self) -> bool { - matches!(self, Self::InMemory(_)) - } - - /// Returns the underlying disk pointer if this is an on-disk reference. - pub fn as_disk_pointer(&self) -> Option<&Pointer> { - match self { - Self::OnDisk(pointer) => Some(pointer), - Self::InMemory(_) => None, - } - } - - /// Returns the underlying in-memory node if this is an in-memory reference. - pub fn as_memory_node(&self) -> Option<&Node> { - match self { - Self::OnDisk(_) => None, - Self::InMemory(node) => Some(node), - } - } - - /// Converts this overlay pointer to an owned Node. - /// For on-disk pointers, this would require loading the node from storage. - /// For in-memory pointers, this clones the existing node. - pub fn into_node(self) -> Result { - match self { - Self::OnDisk(_) => Err(OverlayError::DiskNodeNotLoaded), - Self::InMemory(node) => Ok(*node), - } - } -} - -impl From for OverlayPointer { - fn from(pointer: Pointer) -> Self { - Self::OnDisk(pointer) - } -} - -impl From for OverlayPointer { - fn from(node: Node) -> Self { - Self::InMemory(Box::new(node)) - } -} - -/// Errors that can occur when working with overlay pointers. -#[derive(Debug)] -pub enum OverlayError { - /// Attempted to access a disk node without loading it first - DiskNodeNotLoaded, -} - -impl std::fmt::Display for OverlayError { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self { - Self::DiskNodeNotLoaded => write!(f, "disk node not loaded"), - } - } -} - -impl std::error::Error for OverlayError {} - #[cfg(test)] mod tests { use super::*; - use crate::{account::Account, node::TrieValue}; + use crate::account::Account; use alloy_primitives::U256; use alloy_trie::{Nibbles, EMPTY_ROOT_HASH, KECCAK_EMPTY}; @@ -546,13 +443,13 @@ mod tests { // Find entries with prefix [1, 2] let prefix = Nibbles::from_nibbles([1, 2]); - let subset = frozen.find_prefix_range(&prefix); + 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.find_prefix_range(&prefix); + let subset = frozen.sub_slice_for_prefix(&prefix); assert_eq!(subset.len(), 3); // Should match first three entries } @@ -602,84 +499,10 @@ mod tests { assert!(frozen.lookup(&path).is_none()); assert!(!frozen.contains_prefix_of(&path)); - let subset = frozen.find_prefix_range(&path); + let subset = frozen.sub_slice_for_prefix(&path); assert!(subset.is_empty()); } - #[test] - fn test_overlay_pointer_creation() { - use crate::{location::Location, node::Node, pointer::Pointer}; - use alloy_trie::{nodes::RlpNode, Nibbles, EMPTY_ROOT_HASH}; - - // Test on-disk pointer - let location = Location::for_cell(42); - let rlp = RlpNode::word_rlp(&EMPTY_ROOT_HASH); - let disk_pointer = Pointer::new(location, rlp); - let overlay_pointer = OverlayPointer::on_disk(disk_pointer.clone()); - - assert!(overlay_pointer.is_on_disk()); - assert!(!overlay_pointer.is_in_memory()); - assert_eq!(overlay_pointer.as_disk_pointer(), Some(&disk_pointer)); - assert!(overlay_pointer.as_memory_node().is_none()); - - // Test in-memory pointer - let path = Nibbles::from_nibbles([1, 2, 3, 4]); - let account = test_account(); - let node = Node::new_leaf(path, &TrieValue::Account(account)).unwrap(); - let overlay_pointer = OverlayPointer::in_memory(node.clone()); - - assert!(!overlay_pointer.is_on_disk()); - assert!(overlay_pointer.is_in_memory()); - assert!(overlay_pointer.as_disk_pointer().is_none()); - assert_eq!(overlay_pointer.as_memory_node(), Some(&node)); - } - - #[test] - fn test_overlay_pointer_conversion() { - use crate::{location::Location, pointer::Pointer}; - use alloy_trie::{nodes::RlpNode, Nibbles, EMPTY_ROOT_HASH}; - - // Test From - let location = Location::for_cell(42); - let rlp = RlpNode::word_rlp(&EMPTY_ROOT_HASH); - let disk_pointer = Pointer::new(location, rlp); - let overlay_pointer: OverlayPointer = disk_pointer.into(); - assert!(overlay_pointer.is_on_disk()); - - // Test From - let path = Nibbles::from_nibbles([1, 2, 3, 4]); - let account = test_account(); - let node = Node::new_leaf(path, &TrieValue::Account(account)).unwrap(); - let overlay_pointer: OverlayPointer = node.into(); - assert!(overlay_pointer.is_in_memory()); - } - - #[test] - fn test_overlay_pointer_into_node() { - use alloy_trie::Nibbles; - - // Test in-memory pointer - let path = Nibbles::from_nibbles([1, 2, 3, 4]); - let account = test_account(); - let original_node = Node::new_leaf(path, &TrieValue::Account(account)).unwrap(); - let overlay_pointer = OverlayPointer::in_memory(original_node.clone()); - - let extracted_node = overlay_pointer.into_node().unwrap(); - assert_eq!(extracted_node.prefix(), original_node.prefix()); - assert_eq!(extracted_node.value().unwrap(), original_node.value().unwrap()); - - // Test on-disk pointer (should fail) - use crate::{location::Location, pointer::Pointer}; - use alloy_trie::{nodes::RlpNode, EMPTY_ROOT_HASH}; - - let location = Location::for_cell(42); - let rlp = RlpNode::word_rlp(&EMPTY_ROOT_HASH); - let disk_pointer = Pointer::new(location, rlp); - let overlay_pointer = OverlayPointer::on_disk(disk_pointer); - - assert!(overlay_pointer.into_node().is_err()); - } - #[test] fn test_prefix_offset_functionality() { let mut mutable = OverlayStateMut::new(); @@ -772,7 +595,7 @@ mod tests { let storage_overlay = frozen.with_prefix_offset(4); // Strip 4-nibble account prefix // Test sub_slice_before_prefix - let split_point = Nibbles::from_nibbles([5, 0]); // Storage key [5, 0] + let split_point = [5, 0]; // Storage key [5, 0] let (before, with_prefix, after) = storage_overlay.sub_slice_by_prefix(&split_point); @@ -822,12 +645,12 @@ mod tests { let frozen = mutable.freeze(); // Test finding storage for account 1 using find_prefix_range on original overlay - let account1_storage = frozen.find_prefix_range(&Nibbles::from_nibbles(account1_prefix)); + 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.find_prefix_range(&Nibbles::from_nibbles([5])); + 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())); @@ -948,7 +771,7 @@ mod tests { } let frozen = mutable.freeze(); - let prefix = Nibbles::from_nibbles([1, 2]); + let prefix = [1, 2]; // Test sub_slice_by_prefix let (before, with_prefix, after) = frozen.sub_slice_by_prefix(&prefix); @@ -1054,7 +877,7 @@ mod tests { let storage_overlay = frozen.with_prefix_offset(4); // Test partitioning by storage prefix [5] - let storage_prefix = Nibbles::from_nibbles([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] diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 924489bd..74e1ca27 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -619,7 +619,7 @@ impl StorageEngine { if let Some(account) = account { // Check if there are any storage overlays for this account // Storage overlays have 128-nibble paths that start with this account's 64-nibble path - let storage_overlays = overlay.find_prefix_range(full_path); + let storage_overlays = overlay.sub_slice_for_prefix(full_path); let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() > 64); if !has_storage_overlays { @@ -664,7 +664,7 @@ impl StorageEngine { println!("Account: {account:?}"); // Check if there are any storage overlays for this account // Storage overlays have 128-nibble paths that start with this account's 64-nibble path - let storage_overlays = overlay.find_prefix_range(full_path); + let storage_overlays = overlay.sub_slice_for_prefix(full_path); let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() > 64); if !has_storage_overlays { diff --git a/src/storage/overlay_root_iterative.rs b/src/storage/overlay_root_iterative.rs index 94aabf71..d1031227 100644 --- a/src/storage/overlay_root_iterative.rs +++ b/src/storage/overlay_root_iterative.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use crate::{ account::Account, context::TransactionContext, - node::{encode_account_leaf, Node, NodeKind, TrieValue}, + node::{encode_account_leaf, Node, NodeKind}, overlay::{OverlayState, OverlayValue}, page::SlottedPage, pointer::Pointer, @@ -64,10 +64,6 @@ impl<'a> TraversalStack<'a> { fn pop(&mut self) -> Option<(TriePosition<'a>, OverlayState)> { self.stack.pop() } - - fn is_empty(&self) -> bool { - self.stack.is_empty() - } } impl StorageEngine { @@ -136,23 +132,25 @@ impl StorageEngine { ); } TriePosition::Pointer(path, page, pointer, can_add_by_hash) => { - if overlay.is_empty() && can_add_by_hash && pointer.rlp().as_hash().is_some() { - // No overlay, just add the pointer by hash - // println!("Adding pointer: {:?}", path); - hash_builder.add_branch(path, pointer.rlp().as_hash().unwrap(), true); - } else { - // We have an overlay, need to process the child - self.process_overlayed_child( - context, - overlay, - hash_builder, - &path, - &pointer, - page, - stack, - storage_branch_updates, - )?; + if overlay.is_empty() && can_add_by_hash { + if let Some(hash) = pointer.rlp().as_hash() { + // No overlay, just add the pointer by hash + hash_builder.add_branch(path, hash, true); + continue; + } } + // println!("Adding pointer: {:?}", path); + // We have an overlay, need to process the child + self.process_overlayed_child( + context, + overlay, + hash_builder, + path, + &pointer, + page, + stack, + storage_branch_updates, + )?; } TriePosition::Node(path, page, node) => { let (pre_overlay, overlay, post_overlay) = overlay.sub_slice_by_prefix(&path); @@ -172,7 +170,9 @@ impl StorageEngine { match node.into_kind() { NodeKind::Branch { children } => { - if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = overlay.first() { + if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = + overlay.first() + { if overlay_path == &path { // the overlay invalidates the current node, so just add this // and skip the rest of the db traversal @@ -185,10 +185,15 @@ impl StorageEngine { } } self.process_branch_node_with_overlay( - &overlay, &path, children, page, stack, + overlay, &path, children, page, stack, )?; } - NodeKind::AccountLeaf { nonce_rlp, balance_rlp, code_hash, storage_root } => { + NodeKind::AccountLeaf { + nonce_rlp, + balance_rlp, + code_hash, + storage_root, + } => { self.process_account_leaf_with_overlay( context, &overlay, @@ -226,28 +231,31 @@ impl StorageEngine { Ok(()) } - #[inline(never)] fn process_branch_node_with_overlay<'a>( &'a self, - overlay: &OverlayState, + mut overlay: OverlayState, path: &Nibbles, mut children: [Option; 16], current_page: Arc>, stack: &mut TraversalStack<'a>, ) -> Result<(), Error> { - let mut child_data = Vec::with_capacity(16); + let mut child_data = ArrayVec::<_, 16>::new(); let mut minimum_possible_child_count = 0; - for i in 0..16 { + 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(i); - let child_overlay = overlay.sub_slice_for_prefix(&child_path); - let child_pointer = children[i as usize].take(); + 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 { - match child_overlay.get(0) { + match child_overlay.first() { Some((_, Some(_))) => { // we have a non-tombstone overlay, so there must be at least one descendant // in this child index @@ -258,10 +266,10 @@ impl StorageEngine { } child_data.push((child_path, child_pointer, child_overlay)); + overlay = overlay_after_child; } let can_add_by_hash = minimum_possible_child_count > 1; - // println!("Iterating over children of branch: {:?}", path); for (child_path, child_pointer, child_overlay) in child_data.into_iter().rev() { match child_pointer { Some(pointer) => { @@ -286,7 +294,6 @@ impl StorageEngine { Ok(()) } - #[inline(never)] fn process_account_leaf_with_overlay<'a>( &'a self, context: &TransactionContext, @@ -321,7 +328,9 @@ impl StorageEngine { let has_storage_overlays = overlay.iter().any(|(path, _)| path.len() > 64); if !has_storage_overlays { // println!("Adding account leaf with no storage overlays: {:?}", path); - let storage_root_hash = storage_root.as_ref().map_or(EMPTY_ROOT_HASH, |p| p.rlp().as_hash().unwrap_or(EMPTY_ROOT_HASH)); + 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_hash_builder( hash_builder, @@ -407,7 +416,6 @@ impl StorageEngine { Ok(()) } - #[inline(never)] fn add_account_leaf_to_hash_builder( &self, hash_builder: &mut HashBuilder, @@ -418,19 +426,19 @@ impl StorageEngine { 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 + // 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); + let account_rlp_length = + encode_account_leaf(nonce_rlp, balance_rlp, code_hash, storage_root, &mut value_rlp); hash_builder.add_leaf(path, &buf[..account_rlp_length]); } - #[inline(never)] fn process_overlayed_child<'a>( &'a self, context: &TransactionContext, overlay: OverlayState, hash_builder: &mut HashBuilder, - child_path: &Nibbles, + mut child_path: Nibbles, child: &Pointer, current_page: Arc>, stack: &mut TraversalStack<'a>, @@ -442,7 +450,8 @@ impl StorageEngine { // 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(_))) + 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_hash_builder(hash_builder, &overlay, storage_branch_updates); @@ -452,28 +461,19 @@ impl StorageEngine { if let Some(child_cell) = child.location().cell_index() { let child_node: Node = current_page.get_value(child_cell)?; - stack.push_node( - child_path.join(child_node.prefix()), - child_node, - current_page, - overlay, - ); + 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)?; - stack.push_node( - child_path.join(child_node.prefix()), - child_node, - Arc::new(child_slotted_page), - overlay, - ); + child_path.extend_from_slice(child_node.prefix()); + stack.push_node(child_path, child_node, Arc::new(child_slotted_page), overlay); } Ok(()) } - #[inline(never)] fn process_overlayed_account( &self, hash_builder: &mut HashBuilder, @@ -509,7 +509,6 @@ impl StorageEngine { Ok(()) } - #[inline(never)] fn add_overlay_to_hash_builder( &self, hash_builder: &mut HashBuilder, @@ -527,9 +526,7 @@ impl StorageEngine { match value { Some(OverlayValue::Account(account)) => { - let storage_overlay = overlay - .sub_slice_for_prefix(path) - .with_prefix_offset(64); + let storage_overlay = overlay.sub_slice_for_prefix(path).with_prefix_offset(64); self.process_overlayed_account( hash_builder, Nibbles::from_nibbles(path), From 767d1f61d94697aab72bd41b7058cbca6d270e05 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Fri, 15 Aug 2025 16:39:07 -0700 Subject: [PATCH 12/17] Consolidate to iterative overlay root --- src/account.rs | 2 +- src/node.rs | 7 +- src/overlay.rs | 2 +- src/storage.rs | 1 - src/storage/overlay_root.rs | 2178 ++++++++----------------- src/storage/overlay_root_iterative.rs | 1672 ------------------- src/transaction.rs | 15 +- 7 files changed, 697 insertions(+), 3180 deletions(-) delete mode 100644 src/storage/overlay_root_iterative.rs diff --git a/src/account.rs b/src/account.rs index 44f683b2..b127d58f 100644 --- a/src/account.rs +++ b/src/account.rs @@ -26,7 +26,7 @@ impl Account { /// 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; +const MAX_RLP_ENCODED_LEN: usize = 2 + 3 * 33 + 9; unsafe impl MaxEncodedLen for Account {} diff --git a/src/node.rs b/src/node.rs index 30dc6772..339d34a2 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. diff --git a/src/overlay.rs b/src/overlay.rs index 3a02e71e..0c80f4f7 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -148,7 +148,7 @@ impl OverlayState { #[inline] pub fn first(&self) -> Option<(&[u8], &Option)> { - if self.len() == 0 { + if self.is_empty() { None } else { let (path, value) = &self.data[self.start_idx]; diff --git a/src/storage.rs b/src/storage.rs index a579ff53..2638784e 100644 --- a/src/storage.rs +++ b/src/storage.rs @@ -1,7 +1,6 @@ pub mod debug; pub mod engine; pub mod overlay_root; -pub mod overlay_root_iterative; pub mod proofs; mod test_utils; pub mod value; diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 74e1ca27..5e8133d8 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -1,9 +1,11 @@ +use std::rc::Rc; + use crate::{ account::Account, context::TransactionContext, - node::{Node, NodeKind, TrieValue}, + node::{encode_account_leaf, Node, NodeKind}, overlay::{OverlayState, OverlayValue}, - page::{PageId, SlottedPage}, + page::SlottedPage, pointer::Pointer, storage::engine::{Error, StorageEngine}, }; @@ -11,838 +13,711 @@ use alloy_primitives::{ map::{B256Map, HashMap}, B256, U256, }; -use alloy_rlp::{encode, encode_fixed_size}; -use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, TrieAccount}; +use alloy_rlp::encode_fixed_size; +use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, EMPTY_ROOT_HASH}; use arrayvec::ArrayVec; -#[derive(Debug, Default)] -struct Stats { - account: AccountStats, - storage: StorageStats, +#[derive(Debug)] +enum TriePosition<'a> { + Node(Nibbles, Rc>, Node), + Pointer(Nibbles, Rc>, Pointer, bool), + None, } -#[derive(Debug, Default)] -struct AccountStats { - pointer_hashes: usize, - branch_nodes: usize, - leaves: usize, - overlay_branch_nodes: usize, - overlay_leaves: usize, +struct TraversalStack<'a> { + stack: Vec<(TriePosition<'a>, OverlayState)>, } -#[derive(Debug, Default)] -struct StorageStats { - branch_nodes: usize, - leaves: usize, - overlay_branch_nodes: usize, - overlay_leaves: usize, -} +impl<'a> TraversalStack<'a> { + fn new() -> Self { + Self { stack: vec![] } + } -impl StorageEngine { - /// Computes the root hash with overlay changes without persisting them. - /// This uses alloy-trie's HashBuilder for efficient merkle root computation. - pub fn compute_root_with_overlay( - &self, - context: &TransactionContext, - overlay: &OverlayState, - ) -> Result< - (B256, HashMap, B256Map>), - Error, - > { - if overlay.is_empty() { - // No overlay changes, return current root - return Ok((context.root_node_hash, HashMap::default(), B256Map::default())); - } + fn push_node( + &mut self, + path: Nibbles, + node: Node, + page: Rc>, + overlay: OverlayState, + ) { + self.push(TriePosition::Node(path, page, node), overlay); + } - let mut hash_builder = HashBuilder::default().with_updates(true); - let mut stats = Stats::default(); - let mut updated_storage_branch_nodes = B256Map::default(); + 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); + } - // Use proper trie traversal with overlay integration - self.traverse_with_overlay( - context, - overlay, - &mut hash_builder, - &mut stats, - &mut updated_storage_branch_nodes, - )?; + fn push_none(&mut self, overlay: OverlayState) { + self.push(TriePosition::None, overlay); + } - let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); - // println!("Updated branch nodes: {:?}", updated_branch_nodes); + fn push(&mut self, position: TriePosition<'a>, overlay: OverlayState) { + self.stack.push((position, overlay)); + } - let root = hash_builder.root(); // This will clear the hash builder - println!("Root: {:?}", root); + 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>, +} - println!("Stats: {:?}", stats); - Ok((root, updated_branch_nodes, updated_storage_branch_nodes)) +impl OverlayedRoot { + pub fn new( + root: B256, + updated_branch_nodes: HashMap, + storage_branch_updates: B256Map>, + ) -> Self { + Self { root, updated_branch_nodes, storage_branch_updates } } - /// Helper method to traverse the existing trie and integrate with overlay - fn traverse_with_overlay( - &self, - context: &TransactionContext, - overlay: &OverlayState, - hash_builder: &mut HashBuilder, - stats: &mut Stats, - updated_storage_branch_nodes: &mut B256Map>, - ) -> Result<(), Error> { - // First, collect all existing values from disk - if let Some(root_page_id) = context.root_node_page_id { - self.traverse_page_with_overlay( - context, - overlay, - hash_builder, - root_page_id, - &Nibbles::new(), - Some(context.root_node_hash), - stats, - updated_storage_branch_nodes, - )?; - } else { - // No root page, just add all overlay changes to the hash builder - self.process_nonoverlapping_overlay_state( - context, - overlay, - hash_builder, - stats, - updated_storage_branch_nodes, - )?; + pub fn new_hash(root: B256) -> Self { + Self { + root, + updated_branch_nodes: HashMap::default(), + storage_branch_updates: B256Map::default(), } - - Ok(()) } +} - /// Traverse a specific page with overlay integration - fn traverse_page_with_overlay( +impl StorageEngine { + pub fn compute_state_root_with_overlay_iterative( &self, context: &TransactionContext, - overlay: &OverlayState, - hash_builder: &mut HashBuilder, - page_id: PageId, - path_prefix: &Nibbles, - node_hash: Option, - stats: &mut Stats, - updated_storage_branch_nodes: &mut B256Map>, - ) -> Result<(), Error> { - let page = self.get_page(context, page_id)?; - let slotted_page = SlottedPage::try_from(page)?; + overlay: OverlayState, + ) -> Result { + if overlay.is_empty() { + return Ok(OverlayedRoot::new_hash(context.root_node_hash)); + } - // Start at the root cell (cell_index = 0) and recurse properly - self.traverse_node_with_overlay( + let mut hash_builder = HashBuilder::default().with_updates(true); + let mut storage_branch_updates = B256Map::default(); + + 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_hash_builder( + &mut hash_builder, + &overlay, + &mut storage_branch_updates, + ); + let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); + return Ok(OverlayedRoot::new( + hash_builder.root(), + updated_branch_nodes, + B256Map::default(), + )); + }; + + 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_iterative( context, - overlay, - hash_builder, - &slotted_page, - 0, // Start from root cell - path_prefix, - node_hash, - stats, - updated_storage_branch_nodes, - ) + &mut stack, + &mut hash_builder, + &mut storage_branch_updates, + )?; + + let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); + Ok(OverlayedRoot::new(hash_builder.root(), updated_branch_nodes, storage_branch_updates)) } - /// Traverse a specific node with overlayed state. - /// If this node is a branch, we can add it by hash if there are no overlayed changes under its - /// path prefix. If this node is a leaf, we can add it if there are no overlayed changes - /// under its full path. Otherwise, we need to traverse through the node via its full path. - /// In all cases, we need to process the overlayed state before and after the added or traversed - /// node path. - fn traverse_node_with_overlay( - &self, + fn compute_root_with_overlay_iterative<'a>( + &'a self, context: &TransactionContext, - overlay: &OverlayState, + stack: &mut TraversalStack<'a>, hash_builder: &mut HashBuilder, - slotted_page: &SlottedPage, - cell_index: u8, - path_prefix: &Nibbles, - node_hash: Option, - stats: &mut Stats, - updated_storage_branch_nodes: &mut B256Map>, + storage_branch_updates: &mut B256Map>, ) -> Result<(), Error> { - let node: Node = slotted_page.get_value(cell_index)?; - let full_path = { - let mut full = path_prefix.clone(); - full.extend_from_slice_unchecked(node.prefix()); - full - }; - - match node.kind() { - // if the node is a branch, we can only add it by hash if the path_prefix is completely - // unaffected by the overlay - NodeKind::Branch { .. } => { - // Filter the overlay based on the path_prefix - let (pre_overlay, with_prefix, post_overlay) = - overlay.sub_slice_by_prefix(&path_prefix); - if with_prefix.is_empty() { - self.process_nonoverlapping_overlay_state( - context, - &pre_overlay, + // 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_hash_builder( hash_builder, - stats, - updated_storage_branch_nodes, - )?; - if !node.prefix().is_empty() { - // The node is actually an extension + branch node. We cannot add it by hash - // as the extension prefix may have changed, which - // would in turn change the hash. We need to - // traverse the node via its full path. - println!( - "Processing extension + branch: {:?} / {:?}", - path_prefix, full_path - ); - self.handle_affected_node_with_overlay( - context, - &with_prefix, - hash_builder, - slotted_page, - &node, - &full_path, - stats, - updated_storage_branch_nodes, - )?; - } else { - // This is a true branch node, so we can add it by hash. - let hash = node_hash.unwrap(); - println!("Adding unaffected branch: {:?}", path_prefix); - hash_builder.add_branch(path_prefix.clone(), hash, true); - stats.account.branch_nodes += 1; - self.process_nonoverlapping_overlay_state( - context, - &post_overlay, - hash_builder, - stats, - updated_storage_branch_nodes, - )?; + &overlay, + storage_branch_updates, + ); + } + 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 + hash_builder.add_branch(path, hash, true); + continue; + } } - return Ok(()); - } else { - let (pre_overlay, with_prefix, post_overlay) = - overlay.sub_slice_by_prefix(&full_path); - self.process_nonoverlapping_overlay_state( - context, - &pre_overlay, - hash_builder, - stats, - updated_storage_branch_nodes, - )?; - println!("Processing affected branch: {:?}", path_prefix); - self.handle_affected_node_with_overlay( - context, - &with_prefix, - hash_builder, - slotted_page, - &node, - &full_path, - stats, - updated_storage_branch_nodes, - )?; - self.process_nonoverlapping_overlay_state( + // println!("Adding pointer: {:?}", path); + // We have an overlay, need to process the child + self.process_overlayed_child( context, - &post_overlay, + overlay, hash_builder, - stats, - updated_storage_branch_nodes, + path, + &pointer, + page, + stack, + storage_branch_updates, )?; - return Ok(()); } - } - NodeKind::AccountLeaf { .. } | NodeKind::StorageLeaf { .. } => { - // filter the overlay based on the full_path - let (pre_overlay, with_prefix, post_overlay) = - overlay.sub_slice_by_prefix(&full_path); - self.process_nonoverlapping_overlay_state( - context, - &pre_overlay, - hash_builder, - stats, - updated_storage_branch_nodes, - )?; - if with_prefix.is_empty() { - println!("Adding unaffected leaf: {:?}", full_path); - let node_value = node.value()?; - let rlp_encoded = self.encode_trie_value(&node_value)?; - hash_builder.add_leaf(full_path, &rlp_encoded); - stats.account.leaves += 1; - } else { - println!("Processing affected leaf: {:?}", full_path); - self.handle_affected_node_with_overlay( - context, - &with_prefix, + TriePosition::Node(path, page, node) => { + let (pre_overlay, overlay, post_overlay) = overlay.sub_slice_by_prefix(&path); + self.add_overlay_to_hash_builder( hash_builder, - slotted_page, - &node, - &full_path, - stats, - updated_storage_branch_nodes, - )?; + &pre_overlay, + storage_branch_updates, + ); + // Defer the post_overlay to be processed after the node is traversed + stack.push_none(post_overlay); + + if pre_overlay.contains_prefix_of(&path) { + // A prefix of the node has already been processed, so we can skip the rest + // of the db traversal + continue; + } + + match node.into_kind() { + NodeKind::Branch { children } => { + if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = + 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_hash_builder( + hash_builder, + &overlay, + storage_branch_updates, + ); + continue; + } + } + self.process_branch_node_with_overlay( + overlay, &path, children, page, stack, + )?; + } + NodeKind::AccountLeaf { + nonce_rlp, + balance_rlp, + code_hash, + storage_root, + } => { + self.process_account_leaf_with_overlay( + context, + &overlay, + hash_builder, + path, + page, + storage_branch_updates, + nonce_rlp, + balance_rlp, + code_hash, + storage_root, + )?; + } + NodeKind::StorageLeaf { value_rlp } => { + if let Some((overlay_path, _)) = 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_hash_builder( + hash_builder, + &overlay, + storage_branch_updates, + ); + continue; + } + } + // Leaf node, add it to the hash builder + // println!("Adding storage leaf: {:?}", path); + hash_builder.add_leaf(path, &value_rlp); + } + } } - self.process_nonoverlapping_overlay_state( - context, - &post_overlay, - hash_builder, - stats, - updated_storage_branch_nodes, - )?; - return Ok(()); } } + Ok(()) } - fn process_nonoverlapping_overlay_state( - &self, - context: &TransactionContext, - overlay: &OverlayState, - hash_builder: &mut HashBuilder, - stats: &mut Stats, - updated_storage_branch_nodes: &mut B256Map>, + 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 last_processed_path: Option<&[u8]> = None; - for (path, value) in overlay.iter() { - if let Some(ref last_processed_path) = last_processed_path { - if path.starts_with(last_processed_path) { - // skip over all descendants of a processed path - continue; - } + 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; } - match value { - Some(OverlayValue::Account(account)) => { - self.handle_overlayed_account( - context, - overlay, - hash_builder, - &path, - account, - stats, - updated_storage_branch_nodes, - )?; - last_processed_path = Some(path); - } - Some(OverlayValue::Storage(storage_value)) => { - let rlp_encoded = self.encode_storage(storage_value)?; - println!("Adding overlayed storage leaf: {:?}", path); - hash_builder.add_leaf(Nibbles::from_nibbles(path), &rlp_encoded); - stats.storage.overlay_leaves += 1; - last_processed_path = Some(path); - } - Some(OverlayValue::Hash(hash)) => { - println!("Adding overlayed branch: {:?}", path); - hash_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); - stats.account.overlay_branch_nodes += 1; - last_processed_path = Some(path); + + 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 => { - // Tombstone - skip - println!("Skipping tombstone: {:?}", path); - last_processed_path = Some(path); + 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_none(child_overlay); + } } } } Ok(()) } - /// Handle a node that is affected by overlay changes - fn handle_affected_node_with_overlay( - &self, + fn process_account_leaf_with_overlay<'a>( + &'a self, context: &TransactionContext, overlay: &OverlayState, hash_builder: &mut HashBuilder, - slotted_page: &SlottedPage, - node: &Node, - full_path: &Nibbles, - stats: &mut Stats, - updated_storage_branch_nodes: &mut B256Map>, + path: Nibbles, + current_page: Rc>, + storage_branch_updates: &mut B256Map>, + mut nonce_rlp: ArrayVec, + mut balance_rlp: ArrayVec, + mut code_hash: B256, + storage_root: Option, ) -> Result<(), Error> { - match node.kind() { - NodeKind::Branch { .. } => { - self.traverse_branch_node_with_overlay( - context, - overlay, - hash_builder, - slotted_page, - node, - full_path, - stats, - updated_storage_branch_nodes, - )?; + let overlayed_account = overlay.lookup(&path); + match overlayed_account { + Some(None) => { + // The account is removed in the overlay + // println!("Not adding removed account: {:?}", path); + return Ok(()); } - NodeKind::AccountLeaf { .. } => { - // Account leaf with descendant storage overlays - self.traverse_account_leaf_with_overlayed_storage( - context, - overlay, - hash_builder, - slotted_page, - Some(node), - full_path, - stats, - updated_storage_branch_nodes, - )?; + 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; } - NodeKind::StorageLeaf { .. } => match overlay.get(0) { - Some((_, Some(OverlayValue::Storage(value)))) => { - hash_builder.add_leaf(full_path.clone(), &self.encode_storage(value)?); - stats.storage.overlay_leaves += 1; - } - Some((_, None)) => { - println!("Skipping deleted storage leaf: {:?}", full_path); - stats.storage.overlay_leaves += 1; - } - _ => { - panic!("Storage leaf does not have a valid overlay: {full_path:?}"); - } - }, - } - - Ok(()) - } - - /// Handle traversal of branch nodes with overlay integration - fn traverse_branch_node_with_overlay( - &self, - context: &TransactionContext, - overlay: &OverlayState, - hash_builder: &mut HashBuilder, - slotted_page: &SlottedPage, - node: &Node, - full_path: &Nibbles, - stats: &mut Stats, - updated_storage_branch_nodes: &mut B256Map>, - ) -> Result<(), Error> { - // For each child in the branch node (0-15), check if there are relevant overlay changes - // and recursively traverse child nodes. - // As long as there will be at least 2 children, this branch will still exist, allowing us - // to add the children by hash instead of traversing them. - - let mut children: [(Nibbles, OverlayState, Option<&Pointer>); 16] = - std::array::from_fn(|_| (Nibbles::default(), OverlayState::empty(), None::<&Pointer>)); - for i in 0..16 { - let mut child_path = full_path.clone(); - child_path.push(i); + _ => { + // The account is not updated in the overlay + } + }; - // Create sub-slice of overlay for this child's subtree - let child_overlay = overlay.sub_slice_for_prefix(&child_path); + let has_storage_overlays = overlay.iter().any(|(path, _)| path.len() > 64); + if !has_storage_overlays { + // println!("Adding account leaf with no storage overlays: {:?}", path); + let storage_root_hash = storage_root + .as_ref() + .map_or(EMPTY_ROOT_HASH, |p| p.rlp().as_hash().unwrap_or(EMPTY_ROOT_HASH)); - children[i as usize] = (child_path, child_overlay, node.child(i as u8).unwrap()); + self.add_account_leaf_to_hash_builder( + hash_builder, + path, + &nonce_rlp, + &balance_rlp, + &code_hash, + &storage_root_hash, + ); + return Ok(()); } - // This conservatively counts the minimum number of children that will exist. It may - // undercount. - let min_num_children = children - .iter() - .filter(|(_, child_overlay, child_pointer)| { - if let Some((_, value)) = child_overlay.get(0) { - // If the overlayed value is none, this branch path might be removed. - return value.is_some(); - } else if child_pointer.is_some() { - // If there is no overlayed value and an existing child, then this path will - // definitely exist. - return true; - } - - false - }) - .count(); - - if min_num_children > 1 { - // We are guaranteed to have at least 2 children, keeping this branch node alive. - // We can add non-overlayed children by hash instead of traversing them. This can save - // many lookups, hashes, and page reads. - for (child_path, child_overlay, child_pointer) in children.iter() { - if let Some((path, _)) = child_overlay.get(0) { - if path == child_path.as_slice() { - self.process_nonoverlapping_overlay_state( - context, - &child_overlay, - hash_builder, - stats, - updated_storage_branch_nodes, - )?; - continue; - } - } - - if let Some(child_pointer) = child_pointer { - let child_hash = child_pointer.rlp().as_hash(); - if child_overlay.is_empty() && child_hash.is_some() { - println!("Adding branch node by pointer: {:?}", child_path); - hash_builder.add_branch(child_path.clone(), child_hash.unwrap(), true); - stats.account.pointer_hashes += 1; - } else { - let child_location = child_pointer.location(); + let mut storage_hash_builder = HashBuilder::default().with_updates(true); - if let Some(child_cell_index) = child_location.cell_index() { - // Child is in the same page - self.traverse_node_with_overlay( - context, - &child_overlay, - hash_builder, - slotted_page, - child_cell_index, - &child_path, - child_hash, - stats, - updated_storage_branch_nodes, - )?; - } else if let Some(child_page_id) = child_location.page_id() { - // Child is in a different page - self.traverse_page_with_overlay( - context, - &child_overlay, - hash_builder, - child_page_id, - &child_path, - child_hash, - stats, - updated_storage_branch_nodes, - )?; - } - } + // We have storage overlays, need to compute a new storage root + let storage_overlay = overlay.with_prefix_offset(64); + + match storage_root { + Some(pointer) => { + // println!("Processing overlayed storage root for: {:?}", path); + 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_iterative( + context, + &mut storage_stack, + &mut storage_hash_builder, + storage_branch_updates, + )? } else { - // No child exists on disk, process any matching overlay changes - self.process_nonoverlapping_overlay_state( + 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_iterative( context, - &child_overlay, - hash_builder, - stats, - updated_storage_branch_nodes, + &mut storage_stack, + &mut storage_hash_builder, + storage_branch_updates, )?; } } - return Ok(()); - } - - // Otherwise, it is possible that this branch might be removed. - // We need to traverse and add the children directly in order to be safe. - for (child_path, child_overlay, child_pointer) in children.iter() { - if let Some((path, _)) = child_overlay.get(0) { - if path == child_path.as_slice() { - // Direct overlay - the overlay path exactly matches this node's path - // No need to traverse the child - self.process_nonoverlapping_overlay_state( - context, - &child_overlay, - hash_builder, - stats, - updated_storage_branch_nodes, - )?; - continue; - } + None => { + // No existing storage root, just add the overlay + self.add_overlay_to_hash_builder( + &mut storage_hash_builder, + &storage_overlay, + storage_branch_updates, + ); } + }; + let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_hash_builder.split(); + let new_root = storage_hash_builder.root(); + // println!("New root: {:?}", new_root); - if let Some(child_pointer) = child_pointer { - // Child exists on disk - traverse it with overlay integration - let child_hash = child_pointer.rlp().as_hash(); - let child_location = child_pointer.location(); + storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); + + // println!("Adding overlayed account leaf: {:?}", path); + + self.add_account_leaf_to_hash_builder( + hash_builder, + path, + &nonce_rlp, + &balance_rlp, + &code_hash, + &new_root, + ); - if let Some(child_cell_index) = child_location.cell_index() { - // Child is in the same page - self.traverse_node_with_overlay( - context, - &child_overlay, - hash_builder, - slotted_page, - child_cell_index, - &child_path, - child_hash, - stats, - updated_storage_branch_nodes, - )?; - } else if let Some(child_page_id) = child_location.page_id() { - // Child is in a different page - self.traverse_page_with_overlay( - context, - &child_overlay, - hash_builder, - child_page_id, - &child_path, - child_hash, - stats, - updated_storage_branch_nodes, - )?; - } - } else { - // No child exists on disk, process any matching overlay changes - self.process_nonoverlapping_overlay_state( - context, - &child_overlay, - hash_builder, - stats, - updated_storage_branch_nodes, - )?; - } - } Ok(()) } - /// Handle an account leaf with potential storage overlays using iterative HashBuilder approach - fn traverse_account_leaf_with_overlayed_storage( + fn add_account_leaf_to_hash_builder( &self, + hash_builder: &mut HashBuilder, + 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); + hash_builder.add_leaf(path, &buf[..account_rlp_length]); + } + + fn process_overlayed_child<'a>( + &'a self, context: &TransactionContext, - overlay: &OverlayState, + overlay: OverlayState, hash_builder: &mut HashBuilder, - slotted_page: &SlottedPage, - node: Option<&Node>, - full_path: &Nibbles, - stats: &mut Stats, - updated_storage_branch_nodes: &mut B256Map>, + mut child_path: Nibbles, + child: &Pointer, + current_page: Rc>, + stack: &mut TraversalStack<'a>, + storage_branch_updates: &mut B256Map>, ) -> Result<(), Error> { - // Get the original account from the node - let original_account = node.map(|n| match n.value().unwrap() { - TrieValue::Account(account) => account, - _ => { - panic!("Node is not an account leaf"); - } - }); - - // Check if this account is directly overlaid (account-level changes) - let account = if let Some(overlay_value) = overlay.lookup(full_path) { - // Direct account overlay - use overlay value, but still check for storage overlays - if let Some(trie_value) = overlay_value { - // Even with account overlay, we might need to compute storage root for storage - // overlays - if let OverlayValue::Account(overlay_account) = trie_value { - stats.account.overlay_leaves += 1; - Some(overlay_account) - } else { - panic!("Overlay value is not an account"); - } - } else { - // If overlay_value is None, it's a deletion - skip - println!("Skipping deleted account: {full_path:?}"); - None - } - } else { - stats.account.leaves += 1; - original_account.as_ref() - }; - - if let Some(account) = account { - // Check if there are any storage overlays for this account - // Storage overlays have 128-nibble paths that start with this account's 64-nibble path - let storage_overlays = overlay.sub_slice_for_prefix(full_path); - let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() > 64); - - if !has_storage_overlays { - // No storage overlays for this account, use original account - let rlp_encoded = self.encode_account(account)?; - println!("Adding account leaf: {full_path:?}"); - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + // 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_hash_builder(hash_builder, &overlay, storage_branch_updates); return Ok(()); } + } - // We have storage overlays, need to compute a new storage root using iterative approach - let (new_storage_root, updated_branch_nodes) = self.compute_storage_root_with_overlay( - context, - node, - &storage_overlays, - slotted_page, - stats, - updated_storage_branch_nodes, - )?; - updated_storage_branch_nodes - .insert(B256::from_slice(&full_path.pack()), updated_branch_nodes); - - // Add the modified account to the main HashBuilder - println!("Adding modified account leaf: {full_path:?} with storage root: {new_storage_root:?}"); - let rlp_encoded = self.encode_account_with_root(account, new_storage_root)?; - hash_builder.add_leaf(full_path.clone(), &rlp_encoded); + 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 handle_overlayed_account( + fn process_overlayed_account( &self, - context: &TransactionContext, - overlay: &OverlayState, hash_builder: &mut HashBuilder, - full_path: &[u8], + path: Nibbles, account: &Account, - stats: &mut Stats, - updated_storage_branch_nodes: &mut B256Map>, + storage_overlay: OverlayState, + storage_branch_updates: &mut B256Map>, ) -> Result<(), Error> { - println!("Handling overlayed account: {full_path:?}"); - println!("Account: {account:?}"); - // Check if there are any storage overlays for this account - // Storage overlays have 128-nibble paths that start with this account's 64-nibble path - let storage_overlays = overlay.sub_slice_for_prefix(full_path); - let has_storage_overlays = storage_overlays.iter().any(|(path, _)| path.len() > 64); - - if !has_storage_overlays { - // No storage overlays for this account, use original account - let rlp_encoded = self.encode_account(account)?; - // println!("Adding overlayed account leaf: {:?}", full_path); - hash_builder.add_leaf(Nibbles::from_nibbles(full_path), &rlp_encoded); - stats.account.overlay_leaves += 1; + if storage_overlay.is_empty() { + let encoded = self.encode_account(account); + // println!("Adding overlayed account leaf with no storage overlays: {:?}", path); + hash_builder.add_leaf(path, &encoded); return Ok(()); } - // We have storage overlays, need to compute a new storage root - let storage_overlay = storage_overlays.with_prefix_offset(64); - let mut storage_hash_builder = HashBuilder::default().with_updates(true); - - self.process_nonoverlapping_overlay_state( - context, - &storage_overlay, + self.add_overlay_to_hash_builder( &mut storage_hash_builder, - stats, - updated_storage_branch_nodes, - )?; - - let (mut storage_hash_builder, updated_branch_nodes) = storage_hash_builder.split(); - let new_storage_root = storage_hash_builder.root(); - println!("New storage root: {new_storage_root:?}"); - updated_storage_branch_nodes.insert( - B256::from_slice(&Nibbles::from_nibbles(full_path).pack()), - updated_branch_nodes, + &storage_overlay, + storage_branch_updates, ); - // Add the modified account to the main HashBuilder - let rlp_encoded = self.encode_account_with_root(account, new_storage_root)?; - hash_builder.add_leaf(Nibbles::from_nibbles(full_path), &rlp_encoded); - stats.account.overlay_leaves += 1; + let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_hash_builder.split(); + let storage_root = storage_hash_builder.root(); + + // println!("Updated storage branch nodes: {:?}", updated_storage_branch_nodes); + storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); + + let encoded = self.encode_account_with_root(account, storage_root); + // println!("Adding overlayed account leaf with storage overlays: {:?}", path); + hash_builder.add_leaf(path, &encoded); Ok(()) } - /// Compute the storage root for an account using iterative HashBuilder approach with prefix - /// offset - fn compute_storage_root_with_overlay( + fn add_overlay_to_hash_builder( &self, - context: &TransactionContext, - account_node: Option<&Node>, - storage_overlays: &OverlayState, - slotted_page: &SlottedPage, - stats: &mut Stats, - updated_storage_branch_nodes: &mut B256Map>, - ) -> Result<(B256, HashMap), Error> { - // Create a storage-specific overlay with 64-nibble prefix offset - // This converts 128-nibble storage paths to 64-nibble storage-relative paths - let mut storage_overlay = storage_overlays.with_prefix_offset(64); - if let Some((_, Some(OverlayValue::Account(_)))) = storage_overlay.get(0) { - // println!("Skipping account overlay: {:?}", path); - storage_overlay = storage_overlay.sub_slice(1, storage_overlay.len()); - } - - let mut storage_hash_builder = HashBuilder::default().with_updates(true); + hash_builder: &mut HashBuilder, + overlay: &OverlayState, + storage_branch_updates: &mut B256Map>, + ) { + 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; + } + } - // Get the original account's storage root to determine if we need to traverse existing - // storage - let original_storage_root = account_node - .map(|n| match n.value().unwrap() { - TrieValue::Account(account) => account.storage_root, - _ => panic!("Node is not an account leaf"), /* This shouldn't happen for account - * nodes */ - }) - .unwrap_or(alloy_trie::EMPTY_ROOT_HASH); - - // If the account has existing storage, traverse the existing storage trie - if original_storage_root != alloy_trie::EMPTY_ROOT_HASH { - // Try to find the storage root page from the account node's direct child - if let Ok(Some(storage_root_pointer)) = account_node.unwrap().direct_child() { - let storage_root_location = storage_root_pointer.location(); - let storage_root_hash = storage_root_pointer.rlp().as_hash(); - - if let Some(storage_page_id) = storage_root_location.page_id() { - // Traverse the existing storage trie with storage overlays - self.traverse_page_with_overlay( - context, - &storage_overlay, - &mut storage_hash_builder, - storage_page_id, - &Nibbles::new(), // Start with empty path for storage root - storage_root_hash, - stats, - updated_storage_branch_nodes, - )?; - } else if let Some(storage_cell_index) = storage_root_location.cell_index() { - // Storage root is in the same page as the account - self.traverse_node_with_overlay( - context, - &storage_overlay, - &mut storage_hash_builder, - slotted_page, - storage_cell_index, - &Nibbles::new(), // Start with empty path for storage root - storage_root_hash, - stats, - updated_storage_branch_nodes, - )?; + match value { + Some(OverlayValue::Account(account)) => { + let storage_overlay = overlay.sub_slice_for_prefix(path).with_prefix_offset(64); + self.process_overlayed_account( + hash_builder, + Nibbles::from_nibbles(path), + account, + storage_overlay, + storage_branch_updates, + ) + .unwrap(); + last_processed_path = Some(path); + } + Some(OverlayValue::Storage(storage_value)) => { + let encoded = self.encode_storage(storage_value); + // println!("Adding overlayed storage leaf: {:?}", path); + hash_builder.add_leaf(Nibbles::from_nibbles(path), &encoded); + } + Some(OverlayValue::Hash(hash)) => { + // println!("Adding overlayed branch node: {:?}", path); + hash_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); + last_processed_path = Some(path); + } + None => { + // Tombstone - skip + // println!("Skipping tombstone: {:?}", path); + last_processed_path = Some(path); } } - } else { - // No existing storage, just add overlay changes - self.process_nonoverlapping_overlay_state( - context, - &storage_overlay, - &mut storage_hash_builder, - stats, - updated_storage_branch_nodes, - )?; } - - let (mut storage_hash_builder, updated_branch_nodes) = storage_hash_builder.split(); - // println!("Updated storage branch nodes: {updated_branch_nodes:?}"); - - let computed_root = storage_hash_builder.root(); - // println!("Computed storage root: {computed_root:?}"); - Ok((computed_root, updated_branch_nodes)) - } - - /// Helper to encode TrieValue as RLP - #[inline] - pub fn encode_trie_value(&self, trie_value: &TrieValue) -> Result, Error> { - let rlp_encoded = match trie_value { - TrieValue::Account(account) => self.encode_account(account)?.to_vec(), - TrieValue::Storage(storage_value) => self.encode_storage(storage_value)?.to_vec(), - }; - Ok(rlp_encoded) } #[inline] - pub fn encode_account(&self, account: &Account) -> Result, Error> { + 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, }; - Ok(encode_fixed_size(&trie_account)) + encode_fixed_size(&trie_account) } #[inline] - pub fn encode_account_with_root( - &self, - account: &Account, - root: B256, - ) -> Result, Error> { + 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, }; - Ok(encode_fixed_size(&trie_account)) + encode_fixed_size(&trie_account) } #[inline] - pub fn encode_storage(&self, storage_value: &U256) -> Result, Error> { - Ok(encode_fixed_size(storage_value)) + pub fn encode_storage(&self, storage_value: &U256) -> ArrayVec { + encode_fixed_size(storage_value) } } #[cfg(test)] mod tests { - use super::*; - use crate::{ - account::Account, database::Database, node::TrieValue, overlay::OverlayStateMut, - path::AddressPath, - }; use alloy_primitives::{address, Address, U256}; use alloy_trie::{EMPTY_ROOT_HASH, KECCAK_EMPTY}; use rand::Rng; use tempdir::TempDir; - use test_log::test; + + 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_iterative(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"); + + println!("Account branch updates: {:?}", account_branch_updates); + println!("Storage branch updates: {:?}", storage_branch_updates); + + 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_iterative(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_iterative(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() { @@ -853,9 +728,11 @@ mod tests { let context = db.storage_engine.read_context(); let empty_overlay = OverlayStateMut::new().freeze(); - let (root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &empty_overlay).unwrap(); - assert_eq!(root, context.root_node_hash); + let output = db + .storage_engine + .compute_state_root_with_overlay_iterative(&context, empty_overlay) + .unwrap(); + assert_eq!(output.root, context.root_node_hash); } #[test] @@ -876,10 +753,11 @@ mod tests { let overlay = overlay_mut.freeze(); // Compute root with overlay - let (root, _, _) = db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); + let output = + db.storage_engine.compute_state_root_with_overlay_iterative(&context, overlay).unwrap(); // The root should be different from the empty root (since we have changes) - assert_ne!(root, EMPTY_ROOT_HASH); + assert_ne!(output.root, EMPTY_ROOT_HASH); } #[test] @@ -905,12 +783,6 @@ mod tests { let initial_root = context.root_node_hash; assert_ne!(initial_root, EMPTY_ROOT_HASH); - // Test that we can compute root with empty overlay first (should match initial_root) - let empty_overlay = OverlayStateMut::new().freeze(); - let (root_with_empty_overlay, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &empty_overlay).unwrap(); - assert_eq!(root_with_empty_overlay, initial_root); - // 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); @@ -919,21 +791,7 @@ mod tests { .insert(address_path.clone().into(), Some(OverlayValue::Account(account2.clone()))); let overlay = overlay_mut.freeze(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify: commit the overlay changes and compare roots - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context, - &mut [(address_path.into(), Some(TrieValue::Account(account2)))], - ) - .unwrap(); - let committed_root = verification_context.root_node_hash; - - assert_eq!(overlay_root, committed_root, "Overlay root should match committed root"); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -988,28 +846,11 @@ mod tests { .insert(path1.clone().into(), Some(OverlayValue::Account(account1_updated.clone()))); let overlay = overlay_mut.freeze(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing the change - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context, - &mut [ - (path1.clone().into(), Some(TrieValue::Account(account1_updated.clone()))), - (path2.clone().into(), Some(TrieValue::Account(account2.clone()))), - ], - ) - .unwrap(); - assert_eq!( - overlay_root, verification_context.root_node_hash, - "Case 1: Overlay root should match committed root" - ); - - // Test Case 2: Overlay that creates a new account in an empty subtree (None child case) - // Path 3: starts with 0x2... (first nibble = 2) - this child doesn't exist on disk + 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, @@ -1019,62 +860,11 @@ mod tests { 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(); - let (overlay_root2, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay2).unwrap(); - assert_ne!(overlay_root2, initial_root); - assert_ne!(overlay_root2, overlay_root); - - // Verify by committing the change - let mut verification_context2 = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context2, - &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(); - assert_eq!( - overlay_root2, verification_context2.root_node_hash, - "Case 2: Overlay root should match committed root" - ); - - // Test Case 3: Mixed overlay - affects one subtree, creates new subtree, leaves one - // unaffected - let mut overlay_mut3 = OverlayStateMut::new(); - overlay_mut3 - .insert(path1.clone().into(), Some(OverlayValue::Account(account1_updated.clone()))); // Modify existing - overlay_mut3.insert(path3.clone().into(), Some(OverlayValue::Account(account3.clone()))); // Create new - // path2 is left unaffected - should use add_branch optimization - let overlay3 = overlay_mut3.freeze(); - - let (overlay_root3, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay3).unwrap(); - assert_ne!(overlay_root3, initial_root); - assert_ne!(overlay_root3, overlay_root); - assert_ne!(overlay_root3, overlay_root2); - - // Verify by committing the changes - let mut verification_context3 = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context3, - &mut [ - (path1.clone().into(), Some(TrieValue::Account(account1_updated))), - (path2.clone().into(), Some(TrieValue::Account(account2))), - (path3.clone().into(), Some(TrieValue::Account(account3))), - ], - ) - .unwrap(); - assert_eq!( - overlay_root3, verification_context3.root_node_hash, - "Case 3: Overlay root should match committed root" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay2); } #[test] @@ -1140,27 +930,10 @@ mod tests { overlay_mut.insert(path2.clone().into(), None); // Delete account2 let overlay = overlay_mut.freeze(); - let (overlay_root_with_deletion, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root_with_deletion, root_with_all_accounts); - - // Step 4: Verify by actually erasing account 2 and computing root - db.storage_engine.set_values(&mut context, &mut [(path2.clone().into(), None)]).unwrap(); - let root_after_deletion = context.root_node_hash; - - // The overlay root with tombstone should match the root after actual deletion - assert_eq!( - overlay_root_with_deletion, root_after_deletion, - "Tombstone overlay root should match actual deletion root" - ); + let overlay_root = compare_overlay_with_committed_root(&db, &mut context, &overlay); - // Both should equal the original root with just account1 - assert_eq!( - overlay_root_with_deletion, root_without_account2, - "After deleting account2, root should match original single-account root" - ); assert_eq!( - root_after_deletion, root_without_account2, + overlay_root, root_without_account2, "After deleting account2, committed root should match original single-account root" ); } @@ -1201,9 +974,6 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - assert_ne!(initial_root, EMPTY_ROOT_HASH); - // Test Case 1: Overlay that modifies existing storage let mut overlay_mut = OverlayStateMut::new(); let new_storage_value1 = U256::from(999); @@ -1211,28 +981,7 @@ mod tests { .insert(storage_path1.full_path(), Some(OverlayValue::Storage(new_storage_value1))); let overlay = overlay_mut.freeze(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing the storage change - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context, - &mut [ - (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), - (storage_path1.full_path(), Some(TrieValue::Storage(new_storage_value1))), - (storage_path2.full_path(), Some(TrieValue::Storage(storage_value2))), - ], - ) - .unwrap(); - let committed_root = verification_context.root_node_hash; - - assert_eq!( - overlay_root, committed_root, - "Storage overlay root should match committed root" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay); // Test Case 2: Overlay that adds new storage let mut overlay_mut2 = OverlayStateMut::new(); @@ -1243,58 +992,14 @@ mod tests { overlay_mut2.insert(storage_path3.full_path(), Some(OverlayValue::Storage(storage_value3))); let overlay2 = overlay_mut2.freeze(); - let (overlay_root2, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay2).unwrap(); - assert_ne!(overlay_root2, initial_root); - assert_ne!(overlay_root2, overlay_root); - - // Verify by committing the new storage - let mut verification_context2 = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context2, - &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))), - (storage_path3.full_path(), Some(TrieValue::Storage(storage_value3))), - ], - ) - .unwrap(); - let committed_root2 = verification_context2.root_node_hash; - - assert_eq!( - overlay_root2, committed_root2, - "New storage overlay root should match committed root" - ); + 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(); - let (overlay_root3, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay3).unwrap(); - assert_ne!(overlay_root3, initial_root); - - // Verify by committing the deletion - let mut verification_context3 = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context3, - &mut [ - (account_path.clone().into(), Some(TrieValue::Account(account.clone()))), - (storage_path1.full_path(), Some(TrieValue::Storage(storage_value1))), - // storage_path2 is omitted - effectively deleted - ], - ) - .unwrap(); - let committed_root3 = verification_context3.root_node_hash; - - assert_eq!( - overlay_root3, committed_root3, - "Storage deletion overlay root should match committed root" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay3); // Test Case 4: Combined account and storage changes let mut overlay_mut4 = OverlayStateMut::new(); @@ -1307,86 +1012,14 @@ mod tests { .insert(storage_path1.full_path(), Some(OverlayValue::Storage(new_storage_value1))); let overlay4 = overlay_mut4.freeze(); - let (overlay_root4, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay4).unwrap(); - assert_ne!(overlay_root4, initial_root); - - // Note: For combined account+storage changes, the account overlay takes precedence - // and storage overlays should still be applied to compute the new storage root - // This is a complex case that exercises our account overlay + storage overlay logic - } - - #[test] - fn test_overlay_performance_cases() { - let tmp_dir = TempDir::new("test_overlay_performance_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 a larger trie with multiple levels to test add_branch optimization - let accounts_data = [ - ([0x0, 0x0], Account::new(1, U256::from(100), EMPTY_ROOT_HASH, KECCAK_EMPTY)), - ([0x0, 0x1], Account::new(2, U256::from(200), EMPTY_ROOT_HASH, KECCAK_EMPTY)), - ([0x1, 0x0], Account::new(3, U256::from(300), EMPTY_ROOT_HASH, KECCAK_EMPTY)), - ([0x1, 0x1], Account::new(4, U256::from(400), EMPTY_ROOT_HASH, KECCAK_EMPTY)), - ([0x2, 0x0], Account::new(5, U256::from(500), EMPTY_ROOT_HASH, KECCAK_EMPTY)), - ]; - - let mut changes = Vec::new(); - for (prefix, account) in accounts_data.iter() { - let mut nibbles = vec![prefix[0], prefix[1]]; - // Pad to 64 nibbles for address path - nibbles.extend(vec![0x0; 62]); - let path = AddressPath::new(Nibbles::from_nibbles(nibbles)); - changes.push((path.into(), Some(TrieValue::Account(account.clone())))); - } - - db.storage_engine.set_values(&mut context, &mut changes).unwrap(); - let initial_root = context.root_node_hash; - - // Test: Modify only one leaf in a large subtree - // This should trigger add_branch optimization for unaffected subtrees - let mut nibbles = vec![0x0, 0x0]; - nibbles.extend(vec![0x0; 62]); - let modify_path = AddressPath::new(Nibbles::from_nibbles(nibbles)); - - let updated_account = Account::new(10, U256::from(1000), EMPTY_ROOT_HASH, KECCAK_EMPTY); - let mut overlay_mut = OverlayStateMut::new(); - overlay_mut.insert( - modify_path.clone().into(), - Some(OverlayValue::Account(updated_account.clone())), - ); - let overlay = overlay_mut.freeze(); - - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify the computation is correct by comparing with direct modification - let mut verification_context = db.storage_engine.write_context(); - let mut verification_changes = Vec::new(); - for (i, (prefix, account)) in accounts_data.iter().enumerate() { - let mut nibbles = vec![prefix[0], prefix[1]]; - nibbles.extend(vec![0x0; 62]); - let path = AddressPath::new(Nibbles::from_nibbles(nibbles)); - - if i == 0 { - // Use updated account for first entry - verification_changes - .push((path.into(), Some(TrieValue::Account(updated_account.clone())))); - } else { - verification_changes.push((path.into(), Some(TrieValue::Account(account.clone())))); - } - } + compare_overlay_with_committed_root(&db, &mut context, &overlay4); - db.storage_engine.set_values(&mut verification_context, &mut verification_changes).unwrap(); - let committed_root = verification_context.root_node_hash; + // 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(); - assert_eq!( - overlay_root, committed_root, - "Performance test: Overlay root should match committed root" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay5); } #[test] @@ -1419,8 +1052,6 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - // Test: Add a NEW storage slot via overlay let mut overlay_mut = OverlayStateMut::new(); let storage_key2 = U256::from(20); // New storage key @@ -1432,20 +1063,51 @@ mod tests { overlay_mut.insert(storage_path2.full_path(), Some(OverlayValue::Storage(U256::from(222)))); let overlay = overlay_mut.freeze(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); + 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(); - // Verify by committing the addition - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine.set_values(&mut verification_context, &mut [ - (account_path.into(), Some(TrieValue::Account(account))), - (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), // Original - (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), // New - ]).unwrap(); - let committed_root = verification_context.root_node_hash; + // 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(); - assert_eq!(overlay_root, committed_root, "Adding storage slot via overlay should match"); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -1487,37 +1149,13 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - // 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(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing the changes - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context, - &mut [ - (account1_path.into(), Some(TrieValue::Account(account1))), - (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999)))), - (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888)))), - ], - ) - .unwrap(); - let committed_root = verification_context.root_node_hash; - - assert_eq!( - overlay_root, committed_root, - "Minimal multi-account storage overlay should match" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -1552,8 +1190,6 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - // Test: Apply MULTIPLE storage overlays to the same account let mut overlay_mut = OverlayStateMut::new(); @@ -1565,27 +1201,12 @@ mod tests { 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(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing all changes - db.storage_engine.set_values(&mut context, &mut [ - (account_path.into(), Some(TrieValue::Account(account))), - (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(1111)))), // Modified - (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), // Unchanged - (storage_path3.full_path(), Some(TrieValue::Storage(U256::from(444)))), // New - ]).unwrap(); - let committed_root = context.root_node_hash; - - assert_eq!( - overlay_root, committed_root, - "Multiple storage overlays same account should match" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -1620,30 +1241,12 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - // 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(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing: modify slot1, keep slot2 unchanged - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine.set_values(&mut verification_context, &mut [ - (account_path.into(), Some(TrieValue::Account(account))), - (storage_path1.full_path(), Some(TrieValue::Storage(U256::from(999)))), // Modified - (storage_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), // Unchanged - ]).unwrap(); - let committed_root = verification_context.root_node_hash; - - assert_eq!( - overlay_root, committed_root, - "Single account partial storage overlay should match" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -1691,8 +1294,6 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - // Test: Overlay changes to both accounts' storage let mut overlay_mut = OverlayStateMut::new(); @@ -1713,31 +1314,7 @@ mod tests { let overlay = overlay_mut.freeze(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing all changes - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context, - &mut [ - (account1_path.into(), Some(TrieValue::Account(account1))), - (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(1111)))), - (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(222)))), - (storage1_path3.full_path(), Some(TrieValue::Storage(U256::from(444)))), - (storage2_path1.full_path(), Some(TrieValue::Storage(U256::from(3333)))), - ], - ) - .unwrap(); - let committed_root = verification_context.root_node_hash; - - assert_eq!( - overlay_root, committed_root, - "Multi-account storage overlay root should match committed root" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -1779,31 +1356,12 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - // 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(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing the change - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine.set_values(&mut verification_context, &mut [ - (account1_path.into(), Some(TrieValue::Account(account1))), - (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999)))), // Modified - (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(222)))), // Unchanged - ]).unwrap(); - let committed_root = verification_context.root_node_hash; - - assert_eq!( - overlay_root, committed_root, - "Debug: Simple multi-account storage should match" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -1845,32 +1403,13 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - // 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(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing both changes - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine.set_values(&mut verification_context, &mut [ - (account1_path.into(), Some(TrieValue::Account(account1))), - (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path.full_path(), Some(TrieValue::Storage(U256::from(999)))), // Modified - (storage2_path.full_path(), Some(TrieValue::Storage(U256::from(888)))), // Modified - ]).unwrap(); - let committed_root = verification_context.root_node_hash; - - assert_eq!( - overlay_root, committed_root, - "Debug: Both accounts storage changes should match" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -1908,8 +1447,6 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - // Test: Add NEW storage to account1 let mut overlay_mut = OverlayStateMut::new(); let storage1_key2 = U256::from(40); // New storage key @@ -1920,21 +1457,7 @@ mod tests { .insert(storage1_path2.full_path(), Some(OverlayValue::Storage(U256::from(444)))); let overlay = overlay_mut.freeze(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing the addition - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine.set_values(&mut verification_context, &mut [ - (account1_path.into(), Some(TrieValue::Account(account1))), - (account2_path.into(), Some(TrieValue::Account(account2))), - (storage1_path1.full_path(), Some(TrieValue::Storage(U256::from(111)))), // Original - (storage1_path2.full_path(), Some(TrieValue::Storage(U256::from(444)))), // New - ]).unwrap(); - let committed_root = verification_context.root_node_hash; - - assert_eq!(overlay_root, committed_root, "Debug: Adding new storage should match"); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -1958,8 +1481,6 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - // Test: Add storage to account that had no storage before let mut overlay_mut = OverlayStateMut::new(); let storage_key = U256::from(42); @@ -1969,32 +1490,12 @@ mod tests { overlay_mut.insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value))); let overlay = overlay_mut.freeze(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing the storage addition - let mut verification_context = db.storage_engine.write_context(); - db.storage_engine - .set_values( - &mut verification_context, - &mut [ - (account_path.into(), Some(TrieValue::Account(account))), - (storage_path.full_path(), Some(TrieValue::Storage(storage_value))), - ], - ) - .unwrap(); - let committed_root = verification_context.root_node_hash; - - assert_eq!( - overlay_root, committed_root, - "Empty account storage overlay root should match committed root" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] fn test_1000_accounts_with_10_overlay() { - for _ in 0..100 { + 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(); @@ -2004,7 +1505,7 @@ mod tests { let mut changes: Vec<(Nibbles, Option)> = Vec::with_capacity(1000); - for i in 0..10000 { + for i in 0..1000 { let account_address = Address::random(); let account = Account::new(i, U256::from(rng.random::()), EMPTY_ROOT_HASH, KECCAK_EMPTY); @@ -2017,8 +1518,6 @@ mod tests { db.storage_engine.set_values(&mut context, &mut changes).unwrap(); - let initial_root = context.root_node_hash; - // Create overlay with modifications to every 100th account let mut overlay_mut = OverlayStateMut::new(); @@ -2036,55 +1535,9 @@ mod tests { } } } - - let mut overlay_mut_with_branches = overlay_mut.clone(); let overlay = overlay_mut.freeze(); - let (overlay_root, overlay_branches, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - for (path, branch) in overlay_branches.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(); - } - } - } - let overlay_with_branches = overlay_mut_with_branches.freeze(); - // println!("Overlay with branches: {:?}", overlay_with_branches); - let (overlay_root_with_branches, _, _) = db - .storage_engine - .compute_root_with_overlay(&context, &overlay_with_branches) - .unwrap(); - assert_eq!(overlay_root_with_branches, overlay_root); - - // Verify by committing the storage addition - 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(&mut context, &mut changes).unwrap(); - let committed_root = context.root_node_hash; - - // db.storage_engine - // .print_page(&context, std::io::stdout(), context.root_node_page_id) - // .unwrap(); - - assert_eq!(overlay_root, committed_root, "1000 accounts with 10 overlay should match"); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } } @@ -2107,36 +1560,42 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - 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.into(), Some(OverlayValue::Account(new_account.clone()))); + overlay_mut.insert( + new_account_path.clone().into(), + Some(OverlayValue::Account(new_account.clone())), + ); - let storage_key = U256::from(42); - let storage_path = - crate::path::StoragePath::for_address_and_slot(new_address, storage_key.into()); - let storage_value = U256::from(123); - overlay_mut.insert(storage_path.full_path(), Some(OverlayValue::Storage(storage_value))); + 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 overlay = overlay_mut.freeze(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); + 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))); - // Verify by committing the storage addition - 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(&mut context, &mut changes).unwrap(); - let committed_root = context.root_node_hash; + 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(); - assert_eq!(overlay_root, committed_root, "Overlay new account with storage should match"); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -2169,31 +1628,14 @@ mod tests { ) .unwrap(); - let initial_root = context.root_node_hash; - 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(); - let (overlay_root, _, _) = - db.storage_engine.compute_root_with_overlay(&context, &overlay).unwrap(); - assert_ne!(overlay_root, initial_root); - - // Verify by committing the storage addition - 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(&mut context, &mut changes).unwrap(); - let committed_root = context.root_node_hash; - assert_eq!( - overlay_root, committed_root, - "Overlay update account with storage should match" - ); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } #[test] @@ -2280,246 +1722,6 @@ mod tests { // 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 - let result = db.storage_engine.compute_root_with_overlay(&context, &overlay); - let overlay_root = result.unwrap().0; - - // commit the overlay, ensure that the root is the same as the initial 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(&mut context, &mut changes).unwrap(); - let committed_root = context.root_node_hash; - assert_eq!(committed_root, overlay_root); + compare_overlay_with_committed_root(&db, &mut context, &overlay); } - - #[test] - fn test_complex_case() {} - - // #[test] - // fn test_overlay_state_root_fuzz_rounds() { - // use crate::{path::AddressPath, path::StoragePath, Database}; - // use alloy_primitives::{Address, B256, U256}; - // use rand::{rngs::StdRng, Rng, SeedableRng}; - // use std::collections::{HashMap, HashSet}; - // use tempfile::TempDir; - - // // Deterministic RNG for reproducibility - // let mut rng = StdRng::seed_from_u64(0xDEADBEEFCAFEBABE); - - // // Create an empty DB - // let tmp_dir = TempDir::new().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(); - - // // In-memory expected state: AddressPath -> (Account, slot_hash -> value) - // let mut state: HashMap)> = HashMap::new(); - // let mut all_paths: Vec = Vec::new(); - // let mut used_paths: HashSet = HashSet::new(); - - // // Pre-seed 10,000 accounts - // let mut initial_changes: Vec<(Nibbles, Option)> = Vec::with_capacity(10_000); - // for i in 0..10_000u64 { - // let addr = Address::random(); - // let ap = AddressPath::for_address(addr); - // used_paths.insert(ap.clone()); - // all_paths.push(ap.clone()); - // let account = Account::new(i, U256::from(rng.random::()), EMPTY_ROOT_HASH, - // KECCAK_EMPTY); state.insert(ap.clone(), (account.clone(), HashMap::new())); - // initial_changes.push((ap.into(), Some(TrieValue::Account(account)))); - // } - - // // Of these, 1,000 should have storage: 100 with 10 slots, 900 with 1,000 slots - // let storage_accounts_small: Vec = - // all_paths.iter().cloned().take(100).collect(); let storage_accounts_large: - // Vec = all_paths.iter().cloned().skip(100).take(900).collect(); - - // // Helper to insert unique slots for an address - // let mut insert_slots = |ap: &AddressPath, num_slots: usize| { - // let (_, ref mut slots) = state.get_mut(ap).unwrap(); - // for _ in 0..num_slots { - // // Ensure unique slot per account - // let mut key: B256; - // loop { - // key = B256::from(rng.random::<[u8; 32]>()); - // if !slots.contains_key(&key) { - // break; - // } - // } - // let value = U256::from(rng.random::()); - // slots.insert(key, value); - // let storage_path = StoragePath::for_address_path_and_slot(ap.clone(), key); - // initial_changes.push((storage_path.full_path(), - // Some(TrieValue::Storage(value)))); } - // }; - - // for ap in storage_accounts_small.iter() { - // insert_slots(ap, 10); - // } - // for ap in storage_accounts_large.iter() { - // insert_slots(ap, 1_000); - // } - - // // Sort and commit initial state - // initial_changes.sort_by(|a, b| a.0.cmp(&b.0)); - // db.storage_engine.set_values(&mut context, &mut initial_changes).unwrap(); - - // // Accounts considered to have storage (non-empty map) - // let mut accounts_with_storage: Vec = state - // .iter() - // .filter_map(|(ap, (_acc, slots))| if !slots.is_empty() { Some(ap.clone()) } else { - // None }) .collect(); - - // // Helper to sample n unique elements from a candidate list without replacement - // let mut sample_n = |cands_len: usize, n: usize| -> Vec { - // let n = n.min(cands_len); - // let mut idxs: Vec = (0..cands_len).collect(); - // for i in 0..n { - // let rand_u = rng.random::() as usize; - // let j = i + (rand_u % (cands_len - i)); - // idxs.swap(i, j); - // } - // idxs.truncate(n); - // idxs - // }; - - // // 100 rounds of randomized operations - // for _round in 0..100 { - // let mut overlay_mut = OverlayStateMut::new(); - - // // Insert 50 new accounts - // let mut new_accounts: Vec = Vec::with_capacity(50); - // for i in 0..50u64 { - // let ap = loop { - // let a = Address::random(); - // let candidate = AddressPath::for_address(a); - // if !used_paths.contains(&candidate) { - // break candidate; - // } - // }; - // used_paths.insert(ap.clone()); - // new_accounts.push(ap.clone()); - // all_paths.push(ap.clone()); - // let account = Account::new(i, U256::from(rng.random::()), EMPTY_ROOT_HASH, - // KECCAK_EMPTY); state.insert(ap.clone(), (account.clone(), HashMap::new())); - // overlay_mut.insert(ap.into(), Some(OverlayValue::Account(account))); - // } - - // // Update 50 existing accounts - // let existing_paths: Vec = state.keys().cloned().collect(); - // let filtered: Vec = existing_paths - // .into_iter() - // .filter(|a| !new_accounts.contains(a)) - // .collect(); - // let idxs = sample_n(filtered.len(), 50); - // let update_targets: Vec = idxs.into_iter().map(|i| - // filtered[i].clone()).collect(); for ap in update_targets.iter().cloned() { - // if let Some((acc, _)) = state.get_mut(&ap) { - // acc.balance = U256::from(rng.random::()); - // acc.nonce = acc.nonce.wrapping_add(1); - // overlay_mut.insert(ap.into(), Some(OverlayValue::Account(acc.clone()))); - // } - // } - - // // Delete 10 accounts (avoid newly inserted and updated to keep sets disjoint) - // let delete_candidates: Vec = state - // .keys() - // .filter(|a| !new_accounts.contains(a) && !update_targets.contains(a)) - // .cloned() - // .collect(); - // let idxs = sample_n(delete_candidates.len(), 10); - // let delete_targets: Vec = idxs.into_iter().map(|i| - // delete_candidates[i].clone()).collect(); for ap in delete_targets.iter().cloned() - // { overlay_mut.insert(ap.clone().into(), None); - // state.remove(&ap); - // accounts_with_storage.retain(|a| *a != ap); - // } - - // // Storage operations - // // Helper: sample N accounts with at least K storage slots - // let mut sample_accounts_with_min_slots = |n: usize, k: usize| -> Vec { - // let candidates: Vec = accounts_with_storage - // .iter() - // .filter(|ap| state.get(*ap).map(|(_, m)| m.len()).unwrap_or(0) >= k) - // .cloned() - // .collect(); - // let idxs = sample_n(candidates.len(), n); - // idxs.into_iter().map(|i| candidates[i].clone()).collect() - // }; - - // // Insert 100 new slots into 10 random storage accounts - // let insert_slot_accounts = sample_accounts_with_min_slots(10, 0); - // for ap in insert_slot_accounts.iter().cloned() { - // if let Some((_, slots)) = state.get_mut(&ap) { - // for _ in 0..100 { - // let key = loop { - // let k = B256::from(rng.random::<[u8; 32]>()); - // if !slots.contains_key(&k) { - // break k; - // } - // }; - // let value = U256::from(rng.random::()); - // slots.insert(key, value); - // let storage_path = StoragePath::for_address_path_and_slot(ap.clone(), - // key); overlay_mut.insert(storage_path.full_path(), - // Some(OverlayValue::Storage(value))); } - // if !accounts_with_storage.contains(&ap) { - // accounts_with_storage.push(ap); - // } - // } - // } - - // // Update 10 slots in 10 random storage accounts (10 slots each) - // let update_slot_accounts = sample_accounts_with_min_slots(10, 10); - // for ap in update_slot_accounts.iter().cloned() { - // if let Some((_, map)) = state.get_mut(&ap) { - // let keys: Vec = map.keys().cloned().collect(); - // let idxs = sample_n(keys.len(), 10); - // for i in idxs { - // let key = keys[i]; - // let new_value = U256::from(rng.random::()); - // map.insert(key, new_value); - // let storage_path = StoragePath::for_address_path_and_slot(ap.clone(), - // key); overlay_mut - // .insert(storage_path.full_path(), - // Some(OverlayValue::Storage(new_value))); } - // } - // } - - // // Delete 10 slots in 10 random storage accounts (10 slots each) - // let delete_slot_accounts = sample_accounts_with_min_slots(10, 10); - // for ap in delete_slot_accounts.iter().cloned() { - // if let Some((_, map)) = state.get_mut(&ap) { - // let keys: Vec = map.keys().cloned().collect(); - // let idxs = sample_n(keys.len(), 10); - // for i in idxs { - // let key = keys[i]; - // map.remove(&key); - // let storage_path = StoragePath::for_address_path_and_slot(ap.clone(), - // key); overlay_mut.insert(storage_path.full_path(), None); - // } - // if map.is_empty() { - // accounts_with_storage.retain(|a| *a != ap); - // } - // } - // } - - // // Compute overlay root, apply, and compare - // let overlay = overlay_mut.freeze(); - // let (overlay_root, _, _) = db.storage_engine.compute_root_with_overlay(&context, - // &overlay).unwrap(); - - // let mut commit_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(&mut context, &mut commit_changes).unwrap(); - // let committed_root = context.root_node_hash; - - // assert_eq!(overlay_root, committed_root); - // } - // } } diff --git a/src/storage/overlay_root_iterative.rs b/src/storage/overlay_root_iterative.rs deleted file mode 100644 index d1031227..00000000 --- a/src/storage/overlay_root_iterative.rs +++ /dev/null @@ -1,1672 +0,0 @@ -use std::sync::Arc; - -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, -}; -use alloy_trie::{BranchNodeCompact, HashBuilder, Nibbles, EMPTY_ROOT_HASH}; -use arrayvec::ArrayVec; - -#[derive(Debug)] -enum TriePosition<'a> { - Node(Nibbles, Arc>, Node), - Pointer(Nibbles, Arc>, 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: Arc>, - overlay: OverlayState, - ) { - self.push(TriePosition::Node(path, page, node), overlay); - } - - fn push_pointer( - &mut self, - path: Nibbles, - pointer: Pointer, - page: Arc>, - can_add_by_hash: bool, - overlay: OverlayState, - ) { - self.push(TriePosition::Pointer(path, page, pointer, can_add_by_hash), overlay); - } - - fn push_none(&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() - } -} - -impl StorageEngine { - pub fn compute_state_root_with_overlay_iterative( - &self, - context: &TransactionContext, - overlay: OverlayState, - ) -> Result< - (B256, HashMap, B256Map>), - Error, - > { - if overlay.is_empty() { - return Ok((context.root_node_hash, HashMap::default(), B256Map::default())); - } - - let mut hash_builder = HashBuilder::default().with_updates(true); - let mut storage_branch_updates = B256Map::default(); - - 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_hash_builder( - &mut hash_builder, - &overlay, - &mut storage_branch_updates, - ); - let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); - return Ok((hash_builder.root(), updated_branch_nodes, B256Map::default())); - }; - - let root_node: Node = root_page.get_value(0)?; - let mut stack = TraversalStack::new(); - stack.push_node(root_node.prefix().clone(), root_node, Arc::new(root_page), overlay); - - self.compute_root_with_overlay_iterative( - context, - &mut stack, - &mut hash_builder, - &mut storage_branch_updates, - )?; - - let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); - Ok((hash_builder.root(), updated_branch_nodes, storage_branch_updates)) - } - - fn compute_root_with_overlay_iterative<'a>( - &'a self, - context: &TransactionContext, - stack: &mut TraversalStack<'a>, - hash_builder: &mut HashBuilder, - storage_branch_updates: &mut B256Map>, - ) -> 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_hash_builder( - hash_builder, - &overlay, - storage_branch_updates, - ); - } - 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 - hash_builder.add_branch(path, hash, true); - continue; - } - } - // println!("Adding pointer: {:?}", path); - // We have an overlay, need to process the child - self.process_overlayed_child( - context, - overlay, - hash_builder, - path, - &pointer, - page, - stack, - storage_branch_updates, - )?; - } - TriePosition::Node(path, page, node) => { - let (pre_overlay, overlay, post_overlay) = overlay.sub_slice_by_prefix(&path); - self.add_overlay_to_hash_builder( - hash_builder, - &pre_overlay, - storage_branch_updates, - ); - // Defer the post_overlay to be processed after the node is traversed - stack.push_none(post_overlay); - - if pre_overlay.contains_prefix_of(&path) { - // A prefix of the node has already been processed, so we can skip the rest - // of the db traversal - continue; - } - - match node.into_kind() { - NodeKind::Branch { children } => { - if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = - 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_hash_builder( - hash_builder, - &overlay, - storage_branch_updates, - ); - continue; - } - } - self.process_branch_node_with_overlay( - overlay, &path, children, page, stack, - )?; - } - NodeKind::AccountLeaf { - nonce_rlp, - balance_rlp, - code_hash, - storage_root, - } => { - self.process_account_leaf_with_overlay( - context, - &overlay, - hash_builder, - path, - page, - storage_branch_updates, - nonce_rlp, - balance_rlp, - code_hash, - storage_root, - )?; - } - NodeKind::StorageLeaf { value_rlp } => { - if let Some((overlay_path, _)) = 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_hash_builder( - hash_builder, - &overlay, - storage_branch_updates, - ); - continue; - } - } - // Leaf node, add it to the hash builder - // println!("Adding storage leaf: {:?}", path); - hash_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: Arc>, - 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 { - match child_overlay.first() { - Some((_, Some(_))) => { - // 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_none(child_overlay); - } - } - } - } - Ok(()) - } - - fn process_account_leaf_with_overlay<'a>( - &'a self, - context: &TransactionContext, - overlay: &OverlayState, - hash_builder: &mut HashBuilder, - path: Nibbles, - current_page: Arc>, - storage_branch_updates: &mut B256Map>, - 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 - // println!("Not adding removed account: {:?}", path); - 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 { - // println!("Adding account leaf with no storage overlays: {:?}", path); - 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_hash_builder( - hash_builder, - path, - &nonce_rlp, - &balance_rlp, - &code_hash, - &storage_root_hash, - ); - return Ok(()); - } - - let mut storage_hash_builder = HashBuilder::default().with_updates(true); - - // We have storage overlays, need to compute a new storage root - let storage_overlay = overlay.with_prefix_offset(64); - - match storage_root { - Some(pointer) => { - // println!("Processing overlayed storage root for: {:?}", path); - 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_iterative( - context, - &mut storage_stack, - &mut storage_hash_builder, - storage_branch_updates, - )? - } 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, - Arc::new(slotted_page), - storage_overlay, - ); - self.compute_root_with_overlay_iterative( - context, - &mut storage_stack, - &mut storage_hash_builder, - storage_branch_updates, - )?; - } - } - None => { - // No existing storage root, just add the overlay - self.add_overlay_to_hash_builder( - &mut storage_hash_builder, - &storage_overlay, - storage_branch_updates, - ); - } - }; - let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_hash_builder.split(); - let new_root = storage_hash_builder.root(); - // println!("New root: {:?}", new_root); - - storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); - - // println!("Adding overlayed account leaf: {:?}", path); - - self.add_account_leaf_to_hash_builder( - hash_builder, - path, - &nonce_rlp, - &balance_rlp, - &code_hash, - &new_root, - ); - - Ok(()) - } - - fn add_account_leaf_to_hash_builder( - &self, - hash_builder: &mut HashBuilder, - 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); - hash_builder.add_leaf(path, &buf[..account_rlp_length]); - } - - fn process_overlayed_child<'a>( - &'a self, - context: &TransactionContext, - overlay: OverlayState, - hash_builder: &mut HashBuilder, - mut child_path: Nibbles, - child: &Pointer, - current_page: Arc>, - stack: &mut TraversalStack<'a>, - storage_branch_updates: &mut B256Map>, - ) -> 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_hash_builder(hash_builder, &overlay, storage_branch_updates); - 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, Arc::new(child_slotted_page), overlay); - } - Ok(()) - } - - fn process_overlayed_account( - &self, - hash_builder: &mut HashBuilder, - path: Nibbles, - account: &Account, - storage_overlay: OverlayState, - storage_branch_updates: &mut B256Map>, - ) -> Result<(), Error> { - if storage_overlay.is_empty() { - let encoded = self.encode_account(account).unwrap(); - // println!("Adding overlayed account leaf with no storage overlays: {:?}", path); - hash_builder.add_leaf(path, &encoded); - return Ok(()); - } - - let mut storage_hash_builder = HashBuilder::default().with_updates(true); - self.add_overlay_to_hash_builder( - &mut storage_hash_builder, - &storage_overlay, - storage_branch_updates, - ); - - let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_hash_builder.split(); - let storage_root = storage_hash_builder.root(); - - // println!("Updated storage branch nodes: {:?}", updated_storage_branch_nodes); - - storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); - - let encoded = self.encode_account_with_root(account, storage_root).unwrap(); - // println!("Adding overlayed account leaf with storage overlays: {:?}", path); - hash_builder.add_leaf(path, &encoded); - Ok(()) - } - - fn add_overlay_to_hash_builder( - &self, - hash_builder: &mut HashBuilder, - overlay: &OverlayState, - storage_branch_updates: &mut B256Map>, - ) { - let mut last_processed_path: Option<&[u8]> = None; - for (path, value) in overlay.iter() { - if let Some(ref 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( - hash_builder, - Nibbles::from_nibbles(path), - &account, - storage_overlay, - storage_branch_updates, - ) - .unwrap(); - last_processed_path = Some(path); - } - Some(OverlayValue::Storage(storage_value)) => { - let encoded = self.encode_storage(&storage_value).unwrap(); - // println!("Adding overlayed storage leaf: {:?}", path); - hash_builder.add_leaf(Nibbles::from_nibbles(path), &encoded); - } - Some(OverlayValue::Hash(hash)) => { - // println!("Adding overlayed branch node: {:?}", path); - hash_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); - last_processed_path = Some(path); - } - None => { - // Tombstone - skip - // println!("Skipping tombstone: {:?}", path); - last_processed_path = Some(path); - } - } - } - } -} - -#[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 (overlay_root, account_branch_updates, storage_branch_updates) = db - .storage_engine - .compute_state_root_with_overlay_iterative(context, overlay.clone()) - .unwrap(); - assert_ne!(overlay_root, initial_root, "Overlay should not match initial root"); - - println!("Account branch updates: {:?}", account_branch_updates); - println!("Storage branch updates: {:?}", storage_branch_updates); - - 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 (overlay_root_with_branches, _, _) = db - .storage_engine - .compute_state_root_with_overlay_iterative(context, overlay_with_branches.clone()) - .unwrap(); - 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 (overlay_root_after_commit, _, _) = db - .storage_engine - .compute_state_root_with_overlay_iterative(context, overlay_with_branches) - .unwrap(); - 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 (root, _, _) = db - .storage_engine - .compute_state_root_with_overlay_iterative(&context, empty_overlay) - .unwrap(); - assert_eq!(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 (root, _, _) = - db.storage_engine.compute_state_root_with_overlay_iterative(&context, overlay).unwrap(); - - // The root should be different from the empty root (since we have changes) - assert_ne!(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()); - - // println!("Storage path 1: {:?}", storage_path1.full_path()); - - // 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()); - - // println!("Storage path 2: {:?}", storage_path2.full_path()); - - 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); - } -} diff --git a/src/transaction.rs b/src/transaction.rs index 28686578..82339101 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -8,13 +8,10 @@ use crate::{ node::TrieValue, overlay::OverlayState, path::{AddressPath, StoragePath}, - storage::proofs::AccountProof, + storage::{overlay_root::OverlayedRoot, proofs::AccountProof}, }; -use alloy_primitives::{ - map::{B256Map, HashMap}, - StorageValue, B256, -}; -use alloy_trie::{BranchNodeCompact, Nibbles}; +use alloy_primitives::{map::HashMap, StorageValue, B256}; +use alloy_trie::Nibbles; pub use error::TransactionError; pub use manager::TransactionManager; use sealed::sealed; @@ -91,10 +88,7 @@ impl, K: TransactionKind> Transaction { pub fn compute_root_with_overlay( &self, overlay_state: OverlayState, - ) -> Result< - (B256, HashMap, B256Map>), - TransactionError, - > { + ) -> Result { self.database .storage_engine .compute_state_root_with_overlay_iterative(&self.context, overlay_state) @@ -209,7 +203,6 @@ impl> Transaction { impl> Transaction { pub fn commit(mut self) -> Result<(), TransactionError> { - println!("transaction_metrics: {:?}", self.context.transaction_metrics); let mut transaction_manager = self.database.transaction_manager.lock(); transaction_manager.remove_tx(self.context.snapshot_id, false); From 41a9b585254b6fa65cf1d30ac26a5ec94af95e43 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Fri, 15 Aug 2025 16:47:45 -0700 Subject: [PATCH 13/17] Fix linting errors --- benches/crud_benchmarks.rs | 26 +------------------------- src/overlay.rs | 4 ++-- src/storage/overlay_root.rs | 10 +++++----- 3 files changed, 8 insertions(+), 32 deletions(-) diff --git a/benches/crud_benchmarks.rs b/benches/crud_benchmarks.rs index 4273ebc4..4641306d 100644 --- a/benches/crud_benchmarks.rs +++ b/benches/crud_benchmarks.rs @@ -20,7 +20,7 @@ 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, hint::black_box, io, path::Path, sync::Arc, thread, time::Duration}; +use std::{fs, io, path::Path, sync::Arc, thread, time::Duration}; use tempdir::TempDir; use triedb::{ account::Account, @@ -631,30 +631,6 @@ fn bench_state_root_with_overlay(c: &mut Criterion) { ); }); - // group.bench_function(BenchmarkId::new("state_root_with_storage_overlay", BATCH_SIZE), |b| { - // b.iter_with_setup( - // || { - // let dir = TempDir::new("triedb_bench_state_root_with_storage_overlay").unwrap(); - // copy_files(&base_dir, dir.path()).unwrap(); - // let db_path = dir.path().join(&file_name); - // Database::open(db_path).unwrap() - // }, - // |db| { - // black_box({ - // for _ in 0..100 { - // let tx = db.begin_ro().unwrap(); - - // // Compute the root hash with the overlay - // let _root_result = - // tx.compute_root_with_overlay(&storage_overlay).unwrap(); - - // tx.commit().unwrap(); - // } - // }) - // }, - // ); - // }); - group.finish(); } diff --git a/src/overlay.rs b/src/overlay.rs index 0c80f4f7..00ae9b59 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -1,7 +1,7 @@ use crate::{account::Account, node::TrieValue}; use alloy_primitives::{StorageValue, B256, U256}; use alloy_trie::Nibbles; -use std::sync::Arc; +use std::{cmp::min, sync::Arc}; #[derive(Debug, Clone)] pub enum OverlayValue { @@ -251,7 +251,7 @@ impl OverlayState { true } else { let adjusted_path = &stored_path[self.prefix_offset..]; - &adjusted_path[..prefix.len()] <= prefix + &adjusted_path[..min(adjusted_path.len(), prefix.len())] <= prefix } }) } diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 5e8133d8..877d60cd 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -54,7 +54,7 @@ impl<'a> TraversalStack<'a> { self.push(TriePosition::Pointer(path, page, pointer, can_add_by_hash), overlay); } - fn push_none(&mut self, overlay: OverlayState) { + fn push_overlay(&mut self, overlay: OverlayState) { self.push(TriePosition::None, overlay); } @@ -187,7 +187,7 @@ impl StorageEngine { storage_branch_updates, ); // Defer the post_overlay to be processed after the node is traversed - stack.push_none(post_overlay); + stack.push_overlay(post_overlay); if pre_overlay.contains_prefix_of(&path) { // A prefix of the node has already been processed, so we can skip the rest @@ -308,7 +308,7 @@ impl StorageEngine { // nothing here to add } else { // we have a nonconflicting overlay, add all of it to the hash builder - stack.push_none(child_overlay); + stack.push_overlay(child_overlay); } } } @@ -637,8 +637,8 @@ mod tests { (output.root, output.updated_branch_nodes, output.storage_branch_updates); assert_ne!(overlay_root, initial_root, "Overlay should not match initial root"); - println!("Account branch updates: {:?}", account_branch_updates); - println!("Storage branch updates: {:?}", storage_branch_updates); + // println!("Account branch updates: {:?}", account_branch_updates); + // println!("Storage branch updates: {:?}", storage_branch_updates); let mut overlay_mut_with_branches = OverlayStateMut::new(); From 12f227c942fad0fc7bf4b99e5b4495a20b9ab1c4 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Fri, 15 Aug 2025 17:25:59 -0700 Subject: [PATCH 14/17] Minor cleanup --- src/node.rs | 2 ++ src/transaction.rs | 2 +- src/transaction/error.rs | 14 ++------------ src/transaction/manager.rs | 2 +- 4 files changed, 6 insertions(+), 14 deletions(-) diff --git a/src/node.rs b/src/node.rs index 339d34a2..a0ccd590 100644 --- a/src/node.rs +++ b/src/node.rs @@ -633,6 +633,7 @@ impl Encodable for Node { } } +#[inline] pub fn encode_account_leaf( nonce_rlp: &ArrayVec, balance_rlp: &ArrayVec, @@ -653,6 +654,7 @@ pub fn encode_account_leaf( 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/transaction.rs b/src/transaction.rs index 82339101..86f33d10 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -92,7 +92,7 @@ impl, K: TransactionKind> Transaction { self.database .storage_engine .compute_state_root_with_overlay_iterative(&self.context, overlay_state) - .map_err(|_| TransactionError::Generic) + .map_err(|_| TransactionError) } pub fn get_account_with_proof( diff --git a/src/transaction/error.rs b/src/transaction/error.rs index de0bf916..eb3ea809 100644 --- a/src/transaction/error.rs +++ b/src/transaction/error.rs @@ -1,21 +1,11 @@ use std::{error::Error, fmt}; #[derive(Debug)] -pub enum TransactionError { - /// Generic transaction error (for backward compatibility) - Generic, - /// Overlay functionality is not enabled for this transaction - OverlayNotEnabled, -} +pub struct TransactionError; impl fmt::Display for TransactionError { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Generic => write!(f, "transaction error"), - Self::OverlayNotEnabled => { - write!(f, "overlay functionality is not enabled for this transaction") - } - } + write!(f, "transaction error") } } diff --git a/src/transaction/manager.rs b/src/transaction/manager.rs index 465ee9b0..3cb8efd5 100644 --- a/src/transaction/manager.rs +++ b/src/transaction/manager.rs @@ -25,7 +25,7 @@ impl TransactionManager { pub fn begin_rw(&mut self, snapshot_id: SnapshotId) -> Result { // only allow one writable transaction at a time if self.has_writer { - return Err(TransactionError::Generic); + return Err(TransactionError); } self.has_writer = true; self.open_txs.push(snapshot_id - 1); From 0fe3d76e6f1f2af765824f4c2f6360a2700401df Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Tue, 19 Aug 2025 14:57:41 -0700 Subject: [PATCH 15/17] Fix branch overlay bug --- src/storage/overlay_root.rs | 294 +++++++++++++++++++++--------------- src/transaction.rs | 2 +- 2 files changed, 176 insertions(+), 120 deletions(-) diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 877d60cd..023e7b35 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -91,9 +91,43 @@ impl OverlayedRoot { } } } +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_iterative( + pub fn compute_state_root_with_overlay( &self, context: &TransactionContext, overlay: OverlayState, @@ -102,47 +136,30 @@ impl StorageEngine { return Ok(OverlayedRoot::new_hash(context.root_node_hash)); } - let mut hash_builder = HashBuilder::default().with_updates(true); - let mut storage_branch_updates = B256Map::default(); + 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_hash_builder( - &mut hash_builder, - &overlay, - &mut storage_branch_updates, - ); - let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); - return Ok(OverlayedRoot::new( - hash_builder.root(), - updated_branch_nodes, - B256Map::default(), - )); + 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_iterative( - context, - &mut stack, - &mut hash_builder, - &mut storage_branch_updates, - )?; + self.compute_root_with_overlay(context, &mut stack, &mut root_builder)?; - let (mut hash_builder, updated_branch_nodes) = hash_builder.split(); - Ok(OverlayedRoot::new(hash_builder.root(), updated_branch_nodes, storage_branch_updates)) + Ok(root_builder.finalize()) } - fn compute_root_with_overlay_iterative<'a>( + fn compute_root_with_overlay<'a>( &'a self, context: &TransactionContext, stack: &mut TraversalStack<'a>, - hash_builder: &mut HashBuilder, - storage_branch_updates: &mut B256Map>, + 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. @@ -152,17 +169,14 @@ impl StorageEngine { match position { TriePosition::None => { // No trie position, process whatever is in the overlay - self.add_overlay_to_hash_builder( - hash_builder, - &overlay, - storage_branch_updates, - ); + 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 - hash_builder.add_branch(path, hash, true); + // println!("Adding pointer by hash: {:?}", path); + root_builder.add_branch(path, hash, true); continue; } } @@ -171,48 +185,52 @@ impl StorageEngine { self.process_overlayed_child( context, overlay, - hash_builder, + root_builder, path, &pointer, page, stack, - storage_branch_updates, )?; } TriePosition::Node(path, page, node) => { - let (pre_overlay, overlay, post_overlay) = overlay.sub_slice_by_prefix(&path); - self.add_overlay_to_hash_builder( - hash_builder, - &pre_overlay, - storage_branch_updates, - ); - // Defer the post_overlay to be processed after the node is traversed - stack.push_overlay(post_overlay); - + let (pre_overlay, matching_overlay, post_overlay) = + overlay.sub_slice_by_prefix(&path); if pre_overlay.contains_prefix_of(&path) { - // A prefix of the node has already been processed, so we can skip the rest - // of the db traversal + // 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. + // println!("Processing full overlay due to prefix: {:?}", path); + self.add_overlay_to_root_builder(root_builder, &overlay); continue; } + // println!("Adding pre_overlay before path: {:?}", path); + self.add_overlay_to_root_builder(root_builder, &pre_overlay); + // Defer the post_overlay to be processed after the node is traversed + // println!("Pushing post_overlay after path: {:?}", path); + stack.push_overlay(post_overlay); + match node.into_kind() { NodeKind::Branch { children } => { if let Some((overlay_path, Some(OverlayValue::Hash(_)))) = - overlay.first() + 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_hash_builder( - hash_builder, - &overlay, - storage_branch_updates, + self.add_overlay_to_root_builder( + root_builder, + &matching_overlay, ); continue; } } self.process_branch_node_with_overlay( - overlay, &path, children, page, stack, + matching_overlay, + &path, + children, + page, + stack, )?; } NodeKind::AccountLeaf { @@ -223,11 +241,10 @@ impl StorageEngine { } => { self.process_account_leaf_with_overlay( context, - &overlay, - hash_builder, + &matching_overlay, + root_builder, path, page, - storage_branch_updates, nonce_rlp, balance_rlp, code_hash, @@ -235,21 +252,20 @@ impl StorageEngine { )?; } NodeKind::StorageLeaf { value_rlp } => { - if let Some((overlay_path, _)) = overlay.first() { + 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_hash_builder( - hash_builder, - &overlay, - storage_branch_updates, + self.add_overlay_to_root_builder( + root_builder, + &matching_overlay, ); continue; } } // Leaf node, add it to the hash builder // println!("Adding storage leaf: {:?}", path); - hash_builder.add_leaf(path, &value_rlp); + root_builder.add_leaf(path, &value_rlp); } } } @@ -320,10 +336,9 @@ impl StorageEngine { &'a self, context: &TransactionContext, overlay: &OverlayState, - hash_builder: &mut HashBuilder, + root_builder: &mut RootBuilder, path: Nibbles, current_page: Rc>, - storage_branch_updates: &mut B256Map>, mut nonce_rlp: ArrayVec, mut balance_rlp: ArrayVec, mut code_hash: B256, @@ -354,8 +369,8 @@ impl StorageEngine { .as_ref() .map_or(EMPTY_ROOT_HASH, |p| p.rlp().as_hash().unwrap_or(EMPTY_ROOT_HASH)); - self.add_account_leaf_to_hash_builder( - hash_builder, + self.add_account_leaf_to_root_builder( + root_builder, path, &nonce_rlp, &balance_rlp, @@ -365,7 +380,7 @@ impl StorageEngine { return Ok(()); } - let mut storage_hash_builder = HashBuilder::default().with_updates(true); + 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); @@ -384,11 +399,10 @@ impl StorageEngine { current_page, storage_overlay, ); - self.compute_root_with_overlay_iterative( + self.compute_root_with_overlay( context, &mut storage_stack, - &mut storage_hash_builder, - storage_branch_updates, + &mut storage_root_builder, )? } else { let storage_page = @@ -401,33 +415,32 @@ impl StorageEngine { Rc::new(slotted_page), storage_overlay, ); - self.compute_root_with_overlay_iterative( + self.compute_root_with_overlay( context, &mut storage_stack, - &mut storage_hash_builder, - storage_branch_updates, + &mut storage_root_builder, )?; } } None => { // No existing storage root, just add the overlay - self.add_overlay_to_hash_builder( - &mut storage_hash_builder, - &storage_overlay, - storage_branch_updates, - ); + self.add_overlay_to_root_builder(&mut storage_root_builder, &storage_overlay); } }; - let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_hash_builder.split(); + let (mut storage_hash_builder, updated_storage_branch_nodes) = + storage_root_builder.hash_builder.split(); let new_root = storage_hash_builder.root(); // println!("New root: {:?}", new_root); - storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); + root_builder.add_storage_branch_updates( + B256::from_slice(&path.pack()), + updated_storage_branch_nodes, + ); // println!("Adding overlayed account leaf: {:?}", path); - self.add_account_leaf_to_hash_builder( - hash_builder, + self.add_account_leaf_to_root_builder( + root_builder, path, &nonce_rlp, &balance_rlp, @@ -438,9 +451,9 @@ impl StorageEngine { Ok(()) } - fn add_account_leaf_to_hash_builder( + fn add_account_leaf_to_root_builder( &self, - hash_builder: &mut HashBuilder, + root_builder: &mut RootBuilder, path: Nibbles, nonce_rlp: &ArrayVec, balance_rlp: &ArrayVec, @@ -452,19 +465,18 @@ impl StorageEngine { 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); - hash_builder.add_leaf(path, &buf[..account_rlp_length]); + root_builder.add_leaf(path, &buf[..account_rlp_length]); } fn process_overlayed_child<'a>( &'a self, context: &TransactionContext, overlay: OverlayState, - hash_builder: &mut HashBuilder, + root_builder: &mut RootBuilder, mut child_path: Nibbles, child: &Pointer, current_page: Rc>, stack: &mut TraversalStack<'a>, - storage_branch_updates: &mut B256Map>, ) -> 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 @@ -476,7 +488,7 @@ impl StorageEngine { !matches!(overlay_value, Some(OverlayValue::Account(_))) { // the child path is directly overlayed, so only use the overlay state - self.add_overlay_to_hash_builder(hash_builder, &overlay, storage_branch_updates); + self.add_overlay_to_root_builder(root_builder, &overlay); return Ok(()); } } @@ -498,45 +510,39 @@ impl StorageEngine { fn process_overlayed_account( &self, - hash_builder: &mut HashBuilder, + root_builder: &mut RootBuilder, path: Nibbles, account: &Account, storage_overlay: OverlayState, - storage_branch_updates: &mut B256Map>, ) -> Result<(), Error> { if storage_overlay.is_empty() { let encoded = self.encode_account(account); // println!("Adding overlayed account leaf with no storage overlays: {:?}", path); - hash_builder.add_leaf(path, &encoded); + root_builder.add_leaf(path, &encoded); return Ok(()); } - let mut storage_hash_builder = HashBuilder::default().with_updates(true); - self.add_overlay_to_hash_builder( - &mut storage_hash_builder, - &storage_overlay, - storage_branch_updates, - ); + 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_hash_builder.split(); + let (mut storage_hash_builder, updated_storage_branch_nodes) = + storage_root_builder.hash_builder.split(); let storage_root = storage_hash_builder.root(); // println!("Updated storage branch nodes: {:?}", updated_storage_branch_nodes); - storage_branch_updates.insert(B256::from_slice(&path.pack()), updated_storage_branch_nodes); + 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); // println!("Adding overlayed account leaf with storage overlays: {:?}", path); - hash_builder.add_leaf(path, &encoded); + root_builder.add_leaf(path, &encoded); Ok(()) } - fn add_overlay_to_hash_builder( - &self, - hash_builder: &mut HashBuilder, - overlay: &OverlayState, - storage_branch_updates: &mut B256Map>, - ) { + 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 { @@ -550,11 +556,10 @@ impl StorageEngine { Some(OverlayValue::Account(account)) => { let storage_overlay = overlay.sub_slice_for_prefix(path).with_prefix_offset(64); self.process_overlayed_account( - hash_builder, + root_builder, Nibbles::from_nibbles(path), account, storage_overlay, - storage_branch_updates, ) .unwrap(); last_processed_path = Some(path); @@ -562,11 +567,11 @@ impl StorageEngine { Some(OverlayValue::Storage(storage_value)) => { let encoded = self.encode_storage(storage_value); // println!("Adding overlayed storage leaf: {:?}", path); - hash_builder.add_leaf(Nibbles::from_nibbles(path), &encoded); + root_builder.add_leaf(Nibbles::from_nibbles(path), &encoded); } Some(OverlayValue::Hash(hash)) => { // println!("Adding overlayed branch node: {:?}", path); - hash_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); + root_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); last_processed_path = Some(path); } None => { @@ -629,10 +634,8 @@ mod tests { overlay: &OverlayState, ) -> B256 { let initial_root = context.root_node_hash; - let output = db - .storage_engine - .compute_state_root_with_overlay_iterative(context, overlay.clone()) - .unwrap(); + 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"); @@ -691,7 +694,7 @@ mod tests { let output = db .storage_engine - .compute_state_root_with_overlay_iterative(context, overlay_with_branches.clone()) + .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); @@ -710,7 +713,7 @@ mod tests { // committed root. let output = db .storage_engine - .compute_state_root_with_overlay_iterative(context, overlay_with_branches) + .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); @@ -728,10 +731,8 @@ mod tests { let context = db.storage_engine.read_context(); let empty_overlay = OverlayStateMut::new().freeze(); - let output = db - .storage_engine - .compute_state_root_with_overlay_iterative(&context, empty_overlay) - .unwrap(); + let output = + db.storage_engine.compute_state_root_with_overlay(&context, empty_overlay).unwrap(); assert_eq!(output.root, context.root_node_hash); } @@ -753,8 +754,7 @@ mod tests { let overlay = overlay_mut.freeze(); // Compute root with overlay - let output = - db.storage_engine.compute_state_root_with_overlay_iterative(&context, overlay).unwrap(); + 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); @@ -1724,4 +1724,60 @@ mod tests { // 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 86f33d10..f7600a41 100644 --- a/src/transaction.rs +++ b/src/transaction.rs @@ -91,7 +91,7 @@ impl, K: TransactionKind> Transaction { ) -> Result { self.database .storage_engine - .compute_state_root_with_overlay_iterative(&self.context, overlay_state) + .compute_state_root_with_overlay(&self.context, overlay_state) .map_err(|_| TransactionError) } From 12b689728fffd957141976086611091c4c6e30a8 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Wed, 20 Aug 2025 16:29:49 -0700 Subject: [PATCH 16/17] Address feedback --- benches/crud_benchmarks.rs | 29 ++++++++++++++++++++++------- src/overlay.rs | 8 -------- src/storage/overlay_root.rs | 26 -------------------------- 3 files changed, 22 insertions(+), 41 deletions(-) diff --git a/benches/crud_benchmarks.rs b/benches/crud_benchmarks.rs index 4641306d..9921136f 100644 --- a/benches/crud_benchmarks.rs +++ b/benches/crud_benchmarks.rs @@ -20,7 +20,13 @@ 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, sync::Arc, thread, time::Duration}; +use std::{ + fs, io, + path::Path, + sync::{Arc, Barrier}, + thread, + time::Duration, +}; use tempdir::TempDir; use triedb::{ account::Account, @@ -87,16 +93,15 @@ fn bench_account_reads(c: &mut Criterion) { b.iter_with_setup( || { let db_path = dir.path().join(&file_name); - Arc::new(Database::open(db_path.clone()).unwrap()) - }, - |db| { + let db = Arc::new(Database::open(db_path.clone()).unwrap()); + + // Spawn 4 reader threads let thread_count = 4; - // Divide addresses into 4 chunks - let addresses = addresses.clone(); + 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(); - // Spawn 4 threads for i in 0..thread_count { let start_idx = i * chunk_size; let end_idx = if i == thread_count { @@ -108,8 +113,12 @@ fn bench_account_reads(c: &mut Criterion) { 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(); @@ -126,6 +135,12 @@ fn bench_account_reads(c: &mut Criterion) { 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(); diff --git a/src/overlay.rs b/src/overlay.rs index 00ae9b59..de8305ed 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -118,14 +118,6 @@ pub struct OverlayState { prefix_offset: usize, } -// SAFETY: OverlayState is thread-safe because: -// - Arc<[...]> provides thread-safe reference counting -// - All fields are immutable after construction -// - No interior mutability is used -// - All operations are read-only or create new instances -unsafe impl Send for OverlayState {} -unsafe impl Sync for OverlayState {} - impl OverlayState { /// Creates an empty overlay state. pub fn empty() -> Self { diff --git a/src/storage/overlay_root.rs b/src/storage/overlay_root.rs index 023e7b35..482c04de 100644 --- a/src/storage/overlay_root.rs +++ b/src/storage/overlay_root.rs @@ -175,12 +175,10 @@ impl StorageEngine { if overlay.is_empty() && can_add_by_hash { if let Some(hash) = pointer.rlp().as_hash() { // No overlay, just add the pointer by hash - // println!("Adding pointer by hash: {:?}", path); root_builder.add_branch(path, hash, true); continue; } } - // println!("Adding pointer: {:?}", path); // We have an overlay, need to process the child self.process_overlayed_child( context, @@ -199,15 +197,12 @@ impl StorageEngine { // 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. - // println!("Processing full overlay due to prefix: {:?}", path); self.add_overlay_to_root_builder(root_builder, &overlay); continue; } - // println!("Adding pre_overlay before path: {:?}", path); self.add_overlay_to_root_builder(root_builder, &pre_overlay); // Defer the post_overlay to be processed after the node is traversed - // println!("Pushing post_overlay after path: {:?}", path); stack.push_overlay(post_overlay); match node.into_kind() { @@ -264,7 +259,6 @@ impl StorageEngine { } } // Leaf node, add it to the hash builder - // println!("Adding storage leaf: {:?}", path); root_builder.add_leaf(path, &value_rlp); } } @@ -348,7 +342,6 @@ impl StorageEngine { match overlayed_account { Some(None) => { // The account is removed in the overlay - // println!("Not adding removed account: {:?}", path); return Ok(()); } Some(Some(OverlayValue::Account(overlayed_account))) => { @@ -364,7 +357,6 @@ impl StorageEngine { let has_storage_overlays = overlay.iter().any(|(path, _)| path.len() > 64); if !has_storage_overlays { - // println!("Adding account leaf with no storage overlays: {:?}", path); let storage_root_hash = storage_root .as_ref() .map_or(EMPTY_ROOT_HASH, |p| p.rlp().as_hash().unwrap_or(EMPTY_ROOT_HASH)); @@ -387,7 +379,6 @@ impl StorageEngine { match storage_root { Some(pointer) => { - // println!("Processing overlayed storage root for: {:?}", path); let mut storage_stack = TraversalStack::new(); // load the root storage node @@ -430,15 +421,12 @@ impl StorageEngine { let (mut storage_hash_builder, updated_storage_branch_nodes) = storage_root_builder.hash_builder.split(); let new_root = storage_hash_builder.root(); - // println!("New root: {:?}", new_root); root_builder.add_storage_branch_updates( B256::from_slice(&path.pack()), updated_storage_branch_nodes, ); - // println!("Adding overlayed account leaf: {:?}", path); - self.add_account_leaf_to_root_builder( root_builder, path, @@ -517,7 +505,6 @@ impl StorageEngine { ) -> Result<(), Error> { if storage_overlay.is_empty() { let encoded = self.encode_account(account); - // println!("Adding overlayed account leaf with no storage overlays: {:?}", path); root_builder.add_leaf(path, &encoded); return Ok(()); } @@ -529,15 +516,12 @@ impl StorageEngine { storage_root_builder.hash_builder.split(); let storage_root = storage_hash_builder.root(); - // println!("Updated storage branch nodes: {:?}", updated_storage_branch_nodes); - 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); - // println!("Adding overlayed account leaf with storage overlays: {:?}", path); root_builder.add_leaf(path, &encoded); Ok(()) } @@ -566,17 +550,14 @@ impl StorageEngine { } Some(OverlayValue::Storage(storage_value)) => { let encoded = self.encode_storage(storage_value); - // println!("Adding overlayed storage leaf: {:?}", path); root_builder.add_leaf(Nibbles::from_nibbles(path), &encoded); } Some(OverlayValue::Hash(hash)) => { - // println!("Adding overlayed branch node: {:?}", path); root_builder.add_branch(Nibbles::from_nibbles(path), *hash, false); last_processed_path = Some(path); } None => { // Tombstone - skip - // println!("Skipping tombstone: {:?}", path); last_processed_path = Some(path); } } @@ -640,9 +621,6 @@ mod tests { (output.root, output.updated_branch_nodes, output.storage_branch_updates); assert_ne!(overlay_root, initial_root, "Overlay should not match initial root"); - // println!("Account branch updates: {:?}", account_branch_updates); - // println!("Storage branch updates: {:?}", storage_branch_updates); - let mut overlay_mut_with_branches = OverlayStateMut::new(); overlay.data().iter().for_each(|(path, value)| { @@ -1039,8 +1017,6 @@ mod tests { let storage_path1 = crate::path::StoragePath::for_address_and_slot(account_address, storage_key1.into()); - // println!("Storage path 1: {:?}", storage_path1.full_path()); - // Set up initial state with 1 storage slot db.storage_engine .set_values( @@ -1058,8 +1034,6 @@ mod tests { let storage_path2 = crate::path::StoragePath::for_address_and_slot(account_address, storage_key2.into()); - // println!("Storage path 2: {:?}", storage_path2.full_path()); - overlay_mut.insert(storage_path2.full_path(), Some(OverlayValue::Storage(U256::from(222)))); let overlay = overlay_mut.freeze(); From 0501a672b037f81048d7c9bd4e2313a326ecf628 Mon Sep 17 00:00:00 2001 From: Brian Bland Date: Thu, 21 Aug 2025 08:53:25 -0700 Subject: [PATCH 17/17] Remove test-log dep --- Cargo.toml | 1 - benches/crud_benchmarks.rs | 2 +- src/overlay.rs | 1 - 3 files changed, 1 insertion(+), 3 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 056c2024..ed0f05ca 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,7 +27,6 @@ rand = "0.9.1" walkdir = "2" serde_json = "1.0.140" tempfile = "3.19.1" -test-log = { version = "0.2.18", features = ["trace", "color"] } [lib] bench = false diff --git a/benches/crud_benchmarks.rs b/benches/crud_benchmarks.rs index 9921136f..d12d9024 100644 --- a/benches/crud_benchmarks.rs +++ b/benches/crud_benchmarks.rs @@ -124,7 +124,7 @@ fn bench_account_reads(c: &mut Criterion) { // Read its chunk of addresses for addr in thread_addresses { - let a = tx.get_account(addr.clone()).unwrap(); + let a = tx.get_account(&addr).unwrap(); assert!(a.is_some()); } diff --git a/src/overlay.rs b/src/overlay.rs index de8305ed..6ec6cc7f 100644 --- a/src/overlay.rs +++ b/src/overlay.rs @@ -267,7 +267,6 @@ impl OverlayState { /// /// This operation is O(1) and creates no allocations. The returned OverlayState /// shares the same underlying Arc<[...]> data and only tracks different bounds. - /// Perfect for thread pool scenarios where each thread processes a disjoint range. pub fn sub_slice(&self, start: usize, end: usize) -> OverlayState { let current_len = self.len(); let end = end.min(current_len);