From 85cdf01d29cb4c47eb7b1917c53a4499eb3bc736 Mon Sep 17 00:00:00 2001 From: Roman Krasiuk Date: Sat, 17 Aug 2024 10:34:50 -0700 Subject: [PATCH] fix(trie): exclude trie updates for root node (#10306) --- crates/trie/db/tests/trie.rs | 18 +-------- crates/trie/trie/src/updates.rs | 67 ++++++++++++++++++++++++++++----- 2 files changed, 58 insertions(+), 27 deletions(-) diff --git a/crates/trie/db/tests/trie.rs b/crates/trie/db/tests/trie.rs index 8a9dbee9b2..007ce3069d 100644 --- a/crates/trie/db/tests/trie.rs +++ b/crates/trie/db/tests/trie.rs @@ -431,7 +431,7 @@ fn account_and_storage_trie() { assert_eq!(root, computed_expected_root); // Check account trie - let account_updates = trie_updates.clone().into_sorted(); + let account_updates = trie_updates.into_sorted(); let account_updates = account_updates.account_nodes_ref(); assert_eq!(account_updates.len(), 2); @@ -451,22 +451,6 @@ fn account_and_storage_trie() { assert_eq!(node2a.root_hash, None); assert_eq!(node2a.hashes.len(), 1); - // Check storage trie - let mut updated_storage_trie = - trie_updates.storage_tries_ref().iter().filter(|(_, u)| !u.storage_nodes_ref().is_empty()); - assert_eq!(updated_storage_trie.clone().count(), 1); - let (_, storage_trie_updates) = updated_storage_trie.next().unwrap(); - assert_eq!(storage_trie_updates.storage_nodes_ref().len(), 1); - - let (nibbles3, node3) = storage_trie_updates.storage_nodes_ref().iter().next().unwrap(); - assert!(nibbles3.is_empty()); - assert_eq!(node3.state_mask, TrieMask::new(0b1010)); - assert_eq!(node3.tree_mask, TrieMask::new(0b0000)); - assert_eq!(node3.hash_mask, TrieMask::new(0b0010)); - - assert_eq!(node3.hashes.len(), 1); - assert_eq!(node3.root_hash, Some(account3_storage_root)); - // Add an account // Some address whose hash starts with 0xB1 let address4b = Address::from_str("4f61f2d5ebd991b85aa1677db97307caf5215c91").unwrap(); diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs index 3cd802aa6f..b24eae2bd3 100644 --- a/crates/trie/trie/src/updates.rs +++ b/crates/trie/trie/src/updates.rs @@ -36,8 +36,8 @@ impl TrieUpdates { /// Extends the trie updates. pub fn extend(&mut self, other: Self) { - self.account_nodes.extend(other.account_nodes); - self.removed_nodes.extend(other.removed_nodes); + self.account_nodes.extend(ExcludeEmptyFromPair::from_iter(other.account_nodes)); + self.removed_nodes.extend(ExcludeEmpty::from_iter(other.removed_nodes)); for (hashed_address, storage_trie) in other.storage_tries { self.storage_tries.entry(hashed_address).or_default().extend(storage_trie); } @@ -62,11 +62,11 @@ impl TrieUpdates { ) { // Retrieve deleted keys from trie walker. let (_, removed_node_keys) = walker.split(); - self.removed_nodes.extend(removed_node_keys); + self.removed_nodes.extend(ExcludeEmpty::from_iter(removed_node_keys)); // Retrieve updated nodes from hash builder. let (_, updated_nodes) = hash_builder.split(); - self.account_nodes.extend(updated_nodes); + self.account_nodes.extend(ExcludeEmptyFromPair::from_iter(updated_nodes)); // Add deleted storage tries for destroyed accounts. for destroyed in destroyed_accounts { @@ -102,8 +102,11 @@ pub struct StorageTrieUpdates { #[cfg(feature = "test-utils")] impl StorageTrieUpdates { /// Creates a new storage trie updates that are not marked as deleted. - pub fn new(updates: HashMap) -> Self { - Self { storage_nodes: updates, ..Default::default() } + pub fn new(updates: impl IntoIterator) -> Self { + Self { + storage_nodes: ExcludeEmptyFromPair::from_iter(updates).collect(), + ..Default::default() + } } } @@ -150,19 +153,19 @@ impl StorageTrieUpdates { /// Extends storage trie updates. pub fn extend(&mut self, other: Self) { self.is_deleted |= other.is_deleted; - self.storage_nodes.extend(other.storage_nodes); - self.removed_nodes.extend(other.removed_nodes); + self.storage_nodes.extend(ExcludeEmptyFromPair::from_iter(other.storage_nodes)); + self.removed_nodes.extend(ExcludeEmpty::from_iter(other.removed_nodes)); } /// Finalize storage trie updates for by taking updates from walker and hash builder. pub fn finalize(&mut self, walker: TrieWalker, hash_builder: HashBuilder) { // Retrieve deleted keys from trie walker. let (_, removed_keys) = walker.split(); - self.removed_nodes.extend(removed_keys); + self.removed_nodes.extend(ExcludeEmpty::from_iter(removed_keys)); // Retrieve updated nodes from hash builder. let (_, updated_nodes) = hash_builder.split(); - self.storage_nodes.extend(updated_nodes); + self.storage_nodes.extend(ExcludeEmptyFromPair::from_iter(updated_nodes)); } /// Convert storage trie updates into [`StorageTrieUpdatesSorted`]. @@ -226,3 +229,47 @@ impl StorageTrieUpdatesSorted { &self.removed_nodes } } + +// A wrapper iterator to exclude empty nibbles. +struct ExcludeEmpty(I); + +impl> ExcludeEmpty { + fn from_iter>(iter: T) -> Self { + Self(iter.into_iter()) + } +} + +impl> Iterator for ExcludeEmpty { + type Item = Nibbles; + + fn next(&mut self) -> Option { + loop { + let next = self.0.next()?; + if !next.is_empty() { + return Some(next) + } + } + } +} + +// A wrapper iterator to exclude empty nibbles from pair where nibbles are the key. +struct ExcludeEmptyFromPair(I); + +impl> ExcludeEmptyFromPair { + fn from_iter>(iter: T) -> Self { + Self(iter.into_iter()) + } +} + +impl> Iterator for ExcludeEmptyFromPair { + type Item = (Nibbles, V); + + fn next(&mut self) -> Option { + loop { + let next = self.0.next()?; + if !next.0.is_empty() { + return Some(next) + } + } + } +}