fix(trie): exclude trie updates for root node (#10306)

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();

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<Nibbles, BranchNodeCompact>) -> Self {
-        Self { storage_nodes: updates, ..Default::default() }
+    pub fn new(updates: impl IntoIterator<Item = (Nibbles, BranchNodeCompact)>) -> 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<C>(&mut self, walker: TrieWalker<C>, 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>(I);
+
+impl<I: Iterator<Item = Nibbles>> ExcludeEmpty<I> {
+    fn from_iter<T: IntoIterator<Item = Nibbles, IntoIter = I>>(iter: T) -> Self {
+        Self(iter.into_iter())
+    }
+}
+
+impl<I: Iterator<Item = Nibbles>> Iterator for ExcludeEmpty<I> {
+    type Item = Nibbles;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        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>(I);
+
+impl<V, I: Iterator<Item = (Nibbles, V)>> ExcludeEmptyFromPair<I> {
+    fn from_iter<T: IntoIterator<Item = (Nibbles, V), IntoIter = I>>(iter: T) -> Self {
+        Self(iter.into_iter())
+    }
+}
+
+impl<V, I: Iterator<Item = (Nibbles, V)>> Iterator for ExcludeEmptyFromPair<I> {
+    type Item = (Nibbles, V);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        loop {
+            let next = self.0.next()?;
+            if !next.0.is_empty() {
+                return Some(next)
+            }
+        }
+    }
+}