chore(trie): return nibbles from TrieCursor::current

crates/trie/trie/src/trie.rs

@@ -299,7 +299,9 @@ where
                         };
 
                         trie_updates.extend(
-                            walker_deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)),
+                            walker_deleted_keys
+                                .into_iter()
+                                .map(|nibbles| (TrieKey::AccountNode(nibbles), TrieOp::Delete)),
                         );
                         trie_updates.extend_with_account_updates(hash_builder_updates);
 

crates/trie/trie/src/trie_cursor/database_cursors.rs

@@ -1,5 +1,5 @@
 use super::{TrieCursor, TrieCursorFactory};
-use crate::{updates::TrieKey, BranchNodeCompact, Nibbles, StoredNibbles, StoredNibblesSubKey};
+use crate::{BranchNodeCompact, Nibbles, StoredNibbles, StoredNibblesSubKey};
 use reth_db::{tables, DatabaseError};
 use reth_db_api::{
     cursor::{DbCursorRO, DbDupCursorRO},
@@ -60,8 +60,8 @@ where
     }
 
     /// Retrieves the current key in the cursor.
-    fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
-        Ok(self.0.current()?.map(|(k, _)| TrieKey::AccountNode(k.0)))
+    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
+        Ok(self.0.current()?.map(|(k, _)| k.0))
     }
 }
 
@@ -109,8 +109,8 @@ where
     }
 
     /// Retrieves the current value in the storage trie cursor.
-    fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
-        Ok(self.cursor.current()?.map(|(k, v)| TrieKey::StorageNode(k, v.nibbles.0)))
+    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
+        Ok(self.cursor.current()?.map(|(_, v)| v.nibbles.0))
     }
 }
 

crates/trie/trie/src/trie_cursor/in_memory.rs

@@ -42,7 +42,7 @@ impl<'a, CF: TrieCursorFactory> TrieCursorFactory for InMemoryTrieCursorFactory<
 pub struct InMemoryAccountTrieCursor<'a, C> {
     cursor: C,
     trie_updates: &'a TrieUpdatesSorted,
-    last_key: Option<TrieKey>,
+    last_key: Option<Nibbles>,
 }
 
 impl<'a, C> InMemoryAccountTrieCursor<'a, C> {
@@ -56,12 +56,12 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryAccountTrieCursor<'a, C> {
         &mut self,
         key: Nibbles,
     ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
-        if let Some((trie_key, trie_op)) = self.trie_updates.find_account_node(&key) {
-            self.last_key = Some(trie_key);
+        if let Some((nibbles, trie_op)) = self.trie_updates.find_account_node(&key) {
+            self.last_key = Some(nibbles);
             Ok(trie_op.into_update().map(|node| (key, node)))
         } else {
             let result = self.cursor.seek_exact(key)?;
-            self.last_key = result.as_ref().map(|(k, _)| TrieKey::AccountNode(k.clone()));
+            self.last_key = result.as_ref().map(|(key, _)| key.clone());
             Ok(result)
         }
     }
@@ -78,20 +78,22 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryAccountTrieCursor<'a, C> {
             .cloned();
 
         if let Some((trie_key, trie_op)) = trie_update_entry {
-            let nibbles = match &trie_key {
-                TrieKey::AccountNode(nibbles) => nibbles.clone(),
+            let nibbles = match trie_key {
+                TrieKey::AccountNode(nibbles) => {
+                    self.last_key = Some(nibbles.clone());
+                    nibbles
+                }
                 _ => panic!("Invalid trie key"),
             };
-            self.last_key = Some(trie_key);
             return Ok(trie_op.into_update().map(|node| (nibbles, node)))
         }
 
         let result = self.cursor.seek(key)?;
-        self.last_key = result.as_ref().map(|(k, _)| TrieKey::AccountNode(k.clone()));
+        self.last_key = result.as_ref().map(|(key, _)| key.clone());
         Ok(result)
     }
 
-    fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
+    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
         if self.last_key.is_some() {
             Ok(self.last_key.clone())
         } else {
@@ -108,7 +110,7 @@ pub struct InMemoryStorageTrieCursor<'a, C> {
     trie_update_index: usize,
     trie_updates: &'a TrieUpdatesSorted,
     hashed_address: B256,
-    last_key: Option<TrieKey>,
+    last_key: Option<Nibbles>,
 }
 
 impl<'a, C> InMemoryStorageTrieCursor<'a, C> {
@@ -129,8 +131,7 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryStorageTrieCursor<'a, C> {
             Ok(trie_op.into_update().map(|node| (key, node)))
         } else {
             let result = self.cursor.seek_exact(key)?;
-            self.last_key =
-                result.as_ref().map(|(k, _)| TrieKey::StorageNode(self.hashed_address, k.clone()));
+            self.last_key = result.as_ref().map(|(key, _)| key.clone());
             Ok(result)
         }
     }
@@ -151,20 +152,21 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryStorageTrieCursor<'a, C> {
             trie_update_entry.filter(|(k, _)| matches!(k, TrieKey::StorageNode(_, _)))
         {
             let nibbles = match trie_key {
-                TrieKey::StorageNode(_, nibbles) => nibbles.clone(),
+                TrieKey::StorageNode(_, nibbles) => {
+                    self.last_key = Some(nibbles.clone());
+                    nibbles.clone()
+                }
                 _ => panic!("this should not happen!"),
             };
-            self.last_key = Some(trie_key.clone());
             return Ok(trie_op.as_update().map(|node| (nibbles, node.clone())))
         }
 
         let result = self.cursor.seek(key)?;
-        self.last_key =
-            result.as_ref().map(|(k, _)| TrieKey::StorageNode(self.hashed_address, k.clone()));
+        self.last_key = result.as_ref().map(|(key, _)| key.clone());
         Ok(result)
     }
 
-    fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
+    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
         if self.last_key.is_some() {
             Ok(self.last_key.clone())
         } else {

crates/trie/trie/src/trie_cursor/mod.rs

@@ -1,4 +1,4 @@
-use crate::{updates::TrieKey, BranchNodeCompact, Nibbles};
+use crate::{BranchNodeCompact, Nibbles};
 use reth_db::DatabaseError;
 use reth_primitives::B256;
 
@@ -51,5 +51,5 @@ pub trait TrieCursor: Send + Sync {
         -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError>;
 
     /// Get the current entry.
-    fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError>;
+    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError>;
 }

crates/trie/trie/src/trie_cursor/noop.rs

@@ -1,5 +1,5 @@
 use super::{TrieCursor, TrieCursorFactory};
-use crate::{updates::TrieKey, BranchNodeCompact, Nibbles};
+use crate::{BranchNodeCompact, Nibbles};
 use reth_db::DatabaseError;
 use reth_primitives::B256;
 
@@ -49,7 +49,7 @@ impl TrieCursor for NoopAccountTrieCursor {
     }
 
     /// Retrieves the current cursor position within the account trie.
-    fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
+    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
         Ok(None)
     }
 }
@@ -77,7 +77,7 @@ impl TrieCursor for NoopStorageTrieCursor {
     }
 
     /// Retrieves the current cursor position within storage tries.
-    fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
+    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
         Ok(None)
     }
 }

crates/trie/trie/src/updates.rs

@@ -133,7 +133,9 @@ impl TrieUpdates {
     ) {
         // Add updates from trie walker.
         let (_, deleted_keys) = walker.split();
-        self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
+        self.extend(
+            deleted_keys.into_iter().map(|nibbles| (TrieKey::AccountNode(nibbles), TrieOp::Delete)),
+        );
 
         // Add account node updates from hash builder.
         let (_, hash_builder_updates) = hash_builder.split();
@@ -154,7 +156,11 @@ impl TrieUpdates {
     ) {
         // Add updates from trie walker.
         let (_, deleted_keys) = walker.split();
-        self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
+        self.extend(
+            deleted_keys
+                .into_iter()
+                .map(|nibbles| (TrieKey::StorageNode(hashed_address, nibbles), TrieOp::Delete)),
+        );
 
         // Add storage node updates from hash builder.
         let (_, hash_builder_updates) = hash_builder.split();
@@ -248,21 +254,24 @@ pub struct TrieUpdatesSorted {
 
 impl TrieUpdatesSorted {
     /// Find the account node with the given nibbles.
-    pub fn find_account_node(&self, key: &Nibbles) -> Option<(TrieKey, TrieOp)> {
-        self.trie_operations
-            .iter()
-            .find(|(k, _)| matches!(k, TrieKey::AccountNode(nibbles) if nibbles == key))
-            .cloned()
+    pub fn find_account_node(&self, key: &Nibbles) -> Option<(Nibbles, TrieOp)> {
+        self.trie_operations.iter().find_map(|(k, op)| {
+            k.as_account_node_key()
+                .filter(|nibbles| nibbles == &key)
+                .map(|nibbles| (nibbles.clone(), op.clone()))
+        })
     }
 
     /// Find the storage node with the given hashed address and key.
     pub fn find_storage_node(
         &self,
         hashed_address: &B256,
         key: &Nibbles,
-    ) -> Option<(TrieKey, TrieOp)> {
-        self.trie_operations.iter().find(|(k, _)| {
-            matches!(k, TrieKey::StorageNode(address, nibbles) if address == hashed_address && nibbles == key)
-        }).cloned()
+    ) -> Option<(Nibbles, TrieOp)> {
+        self.trie_operations.iter().find_map(|(k, op)| {
+            k.as_storage_node_key()
+                .filter(|(address, nibbles)| address == &hashed_address && nibbles == &key)
+                .map(|(_, nibbles)| (nibbles.clone(), op.clone()))
+        })
     }
 }

crates/trie/trie/src/walker.rs

@@ -1,7 +1,6 @@
 use crate::{
     prefix_set::PrefixSet,
     trie_cursor::{CursorSubNode, TrieCursor},
-    updates::TrieKey,
     BranchNodeCompact, Nibbles,
 };
 use reth_db::DatabaseError;
@@ -24,7 +23,7 @@ pub struct TrieWalker<C> {
     /// A `PrefixSet` representing the changes to be applied to the trie.
     pub changes: PrefixSet,
     /// The retained trie node keys that need to be deleted.
-    deleted_keys: Option<HashSet<TrieKey>>,
+    deleted_keys: Option<HashSet<Nibbles>>,
 }
 
 impl<C> TrieWalker<C> {
@@ -45,7 +44,7 @@ impl<C> TrieWalker<C> {
     }
 
     /// Split the walker into stack and trie updates.
-    pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<TrieKey>) {
+    pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<Nibbles>) {
         let keys = self.deleted_keys.take();
         (self.stack, keys.unwrap_or_default())
     }