Skip to content

Commit ad8ec33

Browse files
rkrasiukrkrasiuk
authored
chore(trie): return nibbles from TrieCursor::current (#9227)
1 parent 52068cc commit ad8ec33

File tree

7 files changed

+54
-42
lines changed

7 files changed

+54
-42
lines changed

‎crates/trie/trie/src/trie.rs

+3-1
Original file line numberDiff line numberDiff line change
@@ -299,7 +299,9 @@ where
299299
};
300300

301301
trie_updates.extend(
302-
walker_deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)),
302+
walker_deleted_keys
303+
.into_iter()
304+
.map(|nibbles| (TrieKey::AccountNode(nibbles), TrieOp::Delete)),
303305
);
304306
trie_updates.extend_with_account_updates(hash_builder_updates);
305307

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

+5-5
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use super::{TrieCursor, TrieCursorFactory};
2-
use crate::{updates::TrieKey, BranchNodeCompact, Nibbles, StoredNibbles, StoredNibblesSubKey};
2+
use crate::{BranchNodeCompact, Nibbles, StoredNibbles, StoredNibblesSubKey};
33
use reth_db::{tables, DatabaseError};
44
use reth_db_api::{
55
cursor::{DbCursorRO, DbDupCursorRO},
@@ -60,8 +60,8 @@ where
6060
}
6161

6262
/// Retrieves the current key in the cursor.
63-
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
64-
Ok(self.0.current()?.map(|(k, _)| TrieKey::AccountNode(k.0)))
63+
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
64+
Ok(self.0.current()?.map(|(k, _)| k.0))
6565
}
6666
}
6767

@@ -109,8 +109,8 @@ where
109109
}
110110

111111
/// Retrieves the current value in the storage trie cursor.
112-
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
113-
Ok(self.cursor.current()?.map(|(k, v)| TrieKey::StorageNode(k, v.nibbles.0)))
112+
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
113+
Ok(self.cursor.current()?.map(|(_, v)| v.nibbles.0))
114114
}
115115
}
116116

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

+19-17
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,7 @@ impl<'a, CF: TrieCursorFactory> TrieCursorFactory for InMemoryTrieCursorFactory<
4242
pub struct InMemoryAccountTrieCursor<'a, C> {
4343
cursor: C,
4444
trie_updates: &'a TrieUpdatesSorted,
45-
last_key: Option<TrieKey>,
45+
last_key: Option<Nibbles>,
4646
}
4747

4848
impl<'a, C> InMemoryAccountTrieCursor<'a, C> {
@@ -56,12 +56,12 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryAccountTrieCursor<'a, C> {
5656
&mut self,
5757
key: Nibbles,
5858
) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
59-
if let Some((trie_key, trie_op)) = self.trie_updates.find_account_node(&key) {
60-
self.last_key = Some(trie_key);
59+
if let Some((nibbles, trie_op)) = self.trie_updates.find_account_node(&key) {
60+
self.last_key = Some(nibbles);
6161
Ok(trie_op.into_update().map(|node| (key, node)))
6262
} else {
6363
let result = self.cursor.seek_exact(key)?;
64-
self.last_key = result.as_ref().map(|(k, _)| TrieKey::AccountNode(k.clone()));
64+
self.last_key = result.as_ref().map(|(key, _)| key.clone());
6565
Ok(result)
6666
}
6767
}
@@ -78,20 +78,22 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryAccountTrieCursor<'a, C> {
7878
.cloned();
7979

8080
if let Some((trie_key, trie_op)) = trie_update_entry {
81-
let nibbles = match &trie_key {
82-
TrieKey::AccountNode(nibbles) => nibbles.clone(),
81+
let nibbles = match trie_key {
82+
TrieKey::AccountNode(nibbles) => {
83+
self.last_key = Some(nibbles.clone());
84+
nibbles
85+
}
8386
_ => panic!("Invalid trie key"),
8487
};
85-
self.last_key = Some(trie_key);
8688
return Ok(trie_op.into_update().map(|node| (nibbles, node)))
8789
}
8890

8991
let result = self.cursor.seek(key)?;
90-
self.last_key = result.as_ref().map(|(k, _)| TrieKey::AccountNode(k.clone()));
92+
self.last_key = result.as_ref().map(|(key, _)| key.clone());
9193
Ok(result)
9294
}
9395

94-
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
96+
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
9597
if self.last_key.is_some() {
9698
Ok(self.last_key.clone())
9799
} else {
@@ -108,7 +110,7 @@ pub struct InMemoryStorageTrieCursor<'a, C> {
108110
trie_update_index: usize,
109111
trie_updates: &'a TrieUpdatesSorted,
110112
hashed_address: B256,
111-
last_key: Option<TrieKey>,
113+
last_key: Option<Nibbles>,
112114
}
113115

114116
impl<'a, C> InMemoryStorageTrieCursor<'a, C> {
@@ -129,8 +131,7 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryStorageTrieCursor<'a, C> {
129131
Ok(trie_op.into_update().map(|node| (key, node)))
130132
} else {
131133
let result = self.cursor.seek_exact(key)?;
132-
self.last_key =
133-
result.as_ref().map(|(k, _)| TrieKey::StorageNode(self.hashed_address, k.clone()));
134+
self.last_key = result.as_ref().map(|(key, _)| key.clone());
134135
Ok(result)
135136
}
136137
}
@@ -151,20 +152,21 @@ impl<'a, C: TrieCursor> TrieCursor for InMemoryStorageTrieCursor<'a, C> {
151152
trie_update_entry.filter(|(k, _)| matches!(k, TrieKey::StorageNode(_, _)))
152153
{
153154
let nibbles = match trie_key {
154-
TrieKey::StorageNode(_, nibbles) => nibbles.clone(),
155+
TrieKey::StorageNode(_, nibbles) => {
156+
self.last_key = Some(nibbles.clone());
157+
nibbles.clone()
158+
}
155159
_ => panic!("this should not happen!"),
156160
};
157-
self.last_key = Some(trie_key.clone());
158161
return Ok(trie_op.as_update().map(|node| (nibbles, node.clone())))
159162
}
160163

161164
let result = self.cursor.seek(key)?;
162-
self.last_key =
163-
result.as_ref().map(|(k, _)| TrieKey::StorageNode(self.hashed_address, k.clone()));
165+
self.last_key = result.as_ref().map(|(key, _)| key.clone());
164166
Ok(result)
165167
}
166168

167-
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
169+
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
168170
if self.last_key.is_some() {
169171
Ok(self.last_key.clone())
170172
} else {

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

+2-2
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
use crate::{updates::TrieKey, BranchNodeCompact, Nibbles};
1+
use crate::{BranchNodeCompact, Nibbles};
22
use reth_db::DatabaseError;
33
use reth_primitives::B256;
44

@@ -51,5 +51,5 @@ pub trait TrieCursor: Send + Sync {
5151
-> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError>;
5252

5353
/// Get the current entry.
54-
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError>;
54+
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError>;
5555
}

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

+3-3
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
use super::{TrieCursor, TrieCursorFactory};
2-
use crate::{updates::TrieKey, BranchNodeCompact, Nibbles};
2+
use crate::{BranchNodeCompact, Nibbles};
33
use reth_db::DatabaseError;
44
use reth_primitives::B256;
55

@@ -49,7 +49,7 @@ impl TrieCursor for NoopAccountTrieCursor {
4949
}
5050

5151
/// Retrieves the current cursor position within the account trie.
52-
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
52+
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
5353
Ok(None)
5454
}
5555
}
@@ -77,7 +77,7 @@ impl TrieCursor for NoopStorageTrieCursor {
7777
}
7878

7979
/// Retrieves the current cursor position within storage tries.
80-
fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
80+
fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError> {
8181
Ok(None)
8282
}
8383
}

‎crates/trie/trie/src/updates.rs

+20-11
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,9 @@ impl TrieUpdates {
133133
) {
134134
// Add updates from trie walker.
135135
let (_, deleted_keys) = walker.split();
136-
self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
136+
self.extend(
137+
deleted_keys.into_iter().map(|nibbles| (TrieKey::AccountNode(nibbles), TrieOp::Delete)),
138+
);
137139

138140
// Add account node updates from hash builder.
139141
let (_, hash_builder_updates) = hash_builder.split();
@@ -154,7 +156,11 @@ impl TrieUpdates {
154156
) {
155157
// Add updates from trie walker.
156158
let (_, deleted_keys) = walker.split();
157-
self.extend(deleted_keys.into_iter().map(|key| (key, TrieOp::Delete)));
159+
self.extend(
160+
deleted_keys
161+
.into_iter()
162+
.map(|nibbles| (TrieKey::StorageNode(hashed_address, nibbles), TrieOp::Delete)),
163+
);
158164

159165
// Add storage node updates from hash builder.
160166
let (_, hash_builder_updates) = hash_builder.split();
@@ -248,21 +254,24 @@ pub struct TrieUpdatesSorted {
248254

249255
impl TrieUpdatesSorted {
250256
/// Find the account node with the given nibbles.
251-
pub fn find_account_node(&self, key: &Nibbles) -> Option<(TrieKey, TrieOp)> {
252-
self.trie_operations
253-
.iter()
254-
.find(|(k, _)| matches!(k, TrieKey::AccountNode(nibbles) if nibbles == key))
255-
.cloned()
257+
pub fn find_account_node(&self, key: &Nibbles) -> Option<(Nibbles, TrieOp)> {
258+
self.trie_operations.iter().find_map(|(k, op)| {
259+
k.as_account_node_key()
260+
.filter(|nibbles| nibbles == &key)
261+
.map(|nibbles| (nibbles.clone(), op.clone()))
262+
})
256263
}
257264

258265
/// Find the storage node with the given hashed address and key.
259266
pub fn find_storage_node(
260267
&self,
261268
hashed_address: &B256,
262269
key: &Nibbles,
263-
) -> Option<(TrieKey, TrieOp)> {
264-
self.trie_operations.iter().find(|(k, _)| {
265-
matches!(k, TrieKey::StorageNode(address, nibbles) if address == hashed_address && nibbles == key)
266-
}).cloned()
270+
) -> Option<(Nibbles, TrieOp)> {
271+
self.trie_operations.iter().find_map(|(k, op)| {
272+
k.as_storage_node_key()
273+
.filter(|(address, nibbles)| address == &hashed_address && nibbles == &key)
274+
.map(|(_, nibbles)| (nibbles.clone(), op.clone()))
275+
})
267276
}
268277
}

‎crates/trie/trie/src/walker.rs

+2-3
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
use crate::{
22
prefix_set::PrefixSet,
33
trie_cursor::{CursorSubNode, TrieCursor},
4-
updates::TrieKey,
54
BranchNodeCompact, Nibbles,
65
};
76
use reth_db::DatabaseError;
@@ -24,7 +23,7 @@ pub struct TrieWalker<C> {
2423
/// A `PrefixSet` representing the changes to be applied to the trie.
2524
pub changes: PrefixSet,
2625
/// The retained trie node keys that need to be deleted.
27-
deleted_keys: Option<HashSet<TrieKey>>,
26+
deleted_keys: Option<HashSet<Nibbles>>,
2827
}
2928

3029
impl<C> TrieWalker<C> {
@@ -45,7 +44,7 @@ impl<C> TrieWalker<C> {
4544
}
4645

4746
/// Split the walker into stack and trie updates.
48-
pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<TrieKey>) {
47+
pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<Nibbles>) {
4948
let keys = self.deleted_keys.take();
5049
(self.stack, keys.unwrap_or_default())
5150
}

0 commit comments

Comments
 (0)