diff --git a/crates/trie/trie/src/trie.rs b/crates/trie/trie/src/trie.rs index 4d23bfc9818..49e8d00a2b4 100644 --- a/crates/trie/trie/src/trie.rs +++ b/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); diff --git a/crates/trie/trie/src/trie_cursor/database_cursors.rs b/crates/trie/trie/src/trie_cursor/database_cursors.rs index 61e43c19b12..53a64a0b09f 100644 --- a/crates/trie/trie/src/trie_cursor/database_cursors.rs +++ b/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, DatabaseError> { - Ok(self.0.current()?.map(|(k, _)| TrieKey::AccountNode(k.0))) + fn current(&mut self) -> Result, 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, DatabaseError> { - Ok(self.cursor.current()?.map(|(k, v)| TrieKey::StorageNode(k, v.nibbles.0))) + fn current(&mut self) -> Result, DatabaseError> { + Ok(self.cursor.current()?.map(|(_, v)| v.nibbles.0)) } } diff --git a/crates/trie/trie/src/trie_cursor/in_memory.rs b/crates/trie/trie/src/trie_cursor/in_memory.rs index f7c521f69b6..3efd2fa5deb 100644 --- a/crates/trie/trie/src/trie_cursor/in_memory.rs +++ b/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, + last_key: Option, } impl<'a, C> InMemoryAccountTrieCursor<'a, C> { @@ -56,12 +56,12 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryAccountTrieCursor<'a, C> { &mut self, key: Nibbles, ) -> Result, 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, DatabaseError> { + fn current(&mut self) -> Result, 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, + last_key: Option, } 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, DatabaseError> { + fn current(&mut self) -> Result, DatabaseError> { if self.last_key.is_some() { Ok(self.last_key.clone()) } else { diff --git a/crates/trie/trie/src/trie_cursor/mod.rs b/crates/trie/trie/src/trie_cursor/mod.rs index a8e0a01cf09..e5160a5526a 100644 --- a/crates/trie/trie/src/trie_cursor/mod.rs +++ b/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, DatabaseError>; /// Get the current entry. - fn current(&mut self) -> Result, DatabaseError>; + fn current(&mut self) -> Result, DatabaseError>; } diff --git a/crates/trie/trie/src/trie_cursor/noop.rs b/crates/trie/trie/src/trie_cursor/noop.rs index c55bdb80f2c..e49c90613d3 100644 --- a/crates/trie/trie/src/trie_cursor/noop.rs +++ b/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, DatabaseError> { + fn current(&mut self) -> Result, DatabaseError> { Ok(None) } } @@ -77,7 +77,7 @@ impl TrieCursor for NoopStorageTrieCursor { } /// Retrieves the current cursor position within storage tries. - fn current(&mut self) -> Result, DatabaseError> { + fn current(&mut self) -> Result, DatabaseError> { Ok(None) } } diff --git a/crates/trie/trie/src/updates.rs b/crates/trie/trie/src/updates.rs index b982c945591..181f2b82d1d 100644 --- a/crates/trie/trie/src/updates.rs +++ b/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,11 +254,12 @@ 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. @@ -260,9 +267,11 @@ impl TrieUpdatesSorted { &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())) + }) } } diff --git a/crates/trie/trie/src/walker.rs b/crates/trie/trie/src/walker.rs index 990b7fb74d5..64a68ea17de 100644 --- a/crates/trie/trie/src/walker.rs +++ b/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 { /// 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>, + deleted_keys: Option>, } impl TrieWalker { @@ -45,7 +44,7 @@ impl TrieWalker { } /// Split the walker into stack and trie updates. - pub fn split(mut self) -> (Vec, HashSet) { + pub fn split(mut self) -> (Vec, HashSet) { let keys = self.deleted_keys.take(); (self.stack, keys.unwrap_or_default()) }