1use crate::{
2 hashed_cursor::HashedCursor, trie_cursor::TrieCursor, walker::TrieWalker, Nibbles, TrieType,
3};
4use alloy_primitives::B256;
5use alloy_trie::proof::AddedRemovedKeys;
6use reth_storage_errors::db::DatabaseError;
7use tracing::{instrument, trace};
8
9#[derive(Debug)]
11pub struct TrieBranchNode {
12 pub key: Nibbles,
14 pub value: B256,
16 pub children_are_in_trie: bool,
18}
19
20impl TrieBranchNode {
21 pub const fn new(key: Nibbles, value: B256, children_are_in_trie: bool) -> Self {
23 Self { key, value, children_are_in_trie }
24 }
25}
26
27#[derive(Debug)]
29pub enum TrieElement<Value> {
30 Branch(TrieBranchNode),
32 Leaf(B256, Value),
34}
35
36#[derive(Debug)]
38struct SeekedHashedEntry<V> {
39 seeked_key: B256,
41 result: Option<(B256, V)>,
45}
46
47#[derive(Debug)]
52pub struct TrieNodeIter<C, H: HashedCursor, K> {
53 pub walker: TrieWalker<C, K>,
55 pub hashed_cursor: H,
57 trie_type: TrieType,
59 previous_hashed_key: Option<B256>,
62
63 current_hashed_entry: Option<(B256, H::Value)>,
65 should_check_walker_key: bool,
67
68 last_seeked_hashed_entry: Option<SeekedHashedEntry<H::Value>>,
72
73 #[cfg(feature = "metrics")]
74 metrics: crate::metrics::TrieNodeIterMetrics,
75 last_next_result: Option<(B256, H::Value)>,
79}
80
81impl<C, H: HashedCursor, K> TrieNodeIter<C, H, K>
82where
83 H::Value: Copy,
84 K: AsRef<AddedRemovedKeys>,
85{
86 pub fn state_trie(walker: TrieWalker<C, K>, hashed_cursor: H) -> Self {
88 Self::new(walker, hashed_cursor, TrieType::State)
89 }
90
91 pub fn storage_trie(walker: TrieWalker<C, K>, hashed_cursor: H) -> Self {
93 Self::new(walker, hashed_cursor, TrieType::Storage)
94 }
95
96 fn new(walker: TrieWalker<C, K>, hashed_cursor: H, trie_type: TrieType) -> Self {
98 Self {
99 walker,
100 hashed_cursor,
101 trie_type,
102 previous_hashed_key: None,
103 current_hashed_entry: None,
104 should_check_walker_key: false,
105 last_seeked_hashed_entry: None,
106 #[cfg(feature = "metrics")]
107 metrics: crate::metrics::TrieNodeIterMetrics::new(trie_type),
108 last_next_result: None,
109 }
110 }
111
112 pub const fn with_last_hashed_key(mut self, previous_hashed_key: B256) -> Self {
115 self.previous_hashed_key = Some(previous_hashed_key);
116 self
117 }
118
119 fn seek_hashed_entry(&mut self, key: B256) -> Result<Option<(B256, H::Value)>, DatabaseError> {
125 if let Some((last_key, last_value)) = self.last_next_result {
126 if last_key == key {
127 trace!(target: "trie::node_iter", seek_key = ?key, "reusing result from last next() call instead of seeking");
128 self.last_next_result = None; let result = Some((last_key, last_value));
131 self.last_seeked_hashed_entry = Some(SeekedHashedEntry { seeked_key: key, result });
132
133 return Ok(result);
134 }
135 }
136
137 if let Some(entry) = self
138 .last_seeked_hashed_entry
139 .as_ref()
140 .filter(|entry| entry.seeked_key == key)
141 .map(|entry| entry.result)
142 {
143 #[cfg(feature = "metrics")]
144 self.metrics.inc_leaf_nodes_same_seeked();
145 return Ok(entry);
146 }
147
148 trace!(target: "trie::node_iter", ?key, "performing hashed cursor seek");
149 let result = self.hashed_cursor.seek(key)?;
150 self.last_seeked_hashed_entry = Some(SeekedHashedEntry { seeked_key: key, result });
151
152 #[cfg(feature = "metrics")]
153 {
154 self.metrics.inc_leaf_nodes_seeked();
155 }
156 Ok(result)
157 }
158
159 fn next_hashed_entry(&mut self) -> Result<Option<(B256, H::Value)>, DatabaseError> {
163 let result = self.hashed_cursor.next();
164
165 self.last_next_result = result.clone()?;
166
167 #[cfg(feature = "metrics")]
168 {
169 self.metrics.inc_leaf_nodes_advanced();
170 }
171 result
172 }
173}
174
175impl<C, H, K> TrieNodeIter<C, H, K>
176where
177 C: TrieCursor,
178 H: HashedCursor,
179 H::Value: Copy,
180 K: AsRef<AddedRemovedKeys>,
181{
182 #[instrument(
194 level = "trace",
195 target = "trie::node_iter",
196 skip_all,
197 fields(trie_type = ?self.trie_type),
198 ret
199 )]
200 pub fn try_next(
201 &mut self,
202 ) -> Result<Option<TrieElement<<H as HashedCursor>::Value>>, DatabaseError> {
203 loop {
204 if let Some(key) = self.walker.key() {
206 if !self.should_check_walker_key && self.previous_hashed_key.is_none() {
209 self.should_check_walker_key = true;
212 if self.walker.can_skip_current_node {
214 #[cfg(feature = "metrics")]
215 self.metrics.inc_branch_nodes_returned();
216 return Ok(Some(TrieElement::Branch(TrieBranchNode::new(
217 *key,
218 self.walker.hash().unwrap(),
219 self.walker.children_are_in_trie(),
220 ))))
221 }
222 }
223 }
224
225 if let Some((hashed_key, value)) = self.current_hashed_entry.take() {
227 if self.walker.key().is_some_and(|key| key < &Nibbles::unpack(hashed_key)) {
229 self.should_check_walker_key = false;
230 continue
231 }
232
233 trace!(target: "trie::node_iter", ?hashed_key, "next hashed entry");
235 self.current_hashed_entry = self.next_hashed_entry()?;
236
237 #[cfg(feature = "metrics")]
238 self.metrics.inc_leaf_nodes_returned();
239 return Ok(Some(TrieElement::Leaf(hashed_key, value)))
240 }
241
242 match self.previous_hashed_key.take() {
244 Some(hashed_key) => {
245 trace!(target: "trie::node_iter", ?hashed_key, "seeking to the previous hashed entry");
246 self.seek_hashed_entry(hashed_key)?;
248 self.current_hashed_entry = self.next_hashed_entry()?;
249 }
250 None => {
251 let (seek_key, seek_prefix) = match self.walker.next_unprocessed_key() {
254 Some(key) => key,
255 None => break, };
257
258 trace!(
259 target: "trie::node_iter",
260 ?seek_key,
261 can_skip_current_node = self.walker.can_skip_current_node,
262 last = ?self.walker.stack.last(),
263 "seeking to the next unprocessed hashed entry"
264 );
265 let can_skip_node = self.walker.can_skip_current_node;
266 self.walker.advance()?;
267 trace!(
268 target: "trie::node_iter",
269 last = ?self.walker.stack.last(),
270 "advanced walker"
271 );
272
273 if can_skip_node &&
284 self.walker.key().is_some_and(|key| key.starts_with(&seek_prefix)) &&
285 self.walker.children_are_in_trie()
286 {
287 trace!(
288 target: "trie::node_iter",
289 ?seek_key,
290 walker_hash = ?self.walker.maybe_hash(),
291 "skipping hashed seek"
292 );
293
294 self.should_check_walker_key = false;
295 continue
296 }
297
298 self.current_hashed_entry = self.seek_hashed_entry(seek_key)?;
299 }
300 }
301 }
302
303 Ok(None)
304 }
305}
306
307#[cfg(test)]
308mod tests {
309 use crate::{
310 hashed_cursor::{
311 mock::MockHashedCursorFactory, noop::NoopHashedAccountCursor, HashedCursorFactory,
312 HashedPostStateAccountCursor,
313 },
314 mock::{KeyVisit, KeyVisitType},
315 trie_cursor::{
316 mock::MockTrieCursorFactory, noop::NoopAccountTrieCursor, TrieCursorFactory,
317 },
318 walker::TrieWalker,
319 };
320 use alloy_primitives::{
321 b256,
322 map::{B256Map, HashMap},
323 };
324 use alloy_trie::{
325 BranchNodeCompact, HashBuilder, Nibbles, TrieAccount, TrieMask, EMPTY_ROOT_HASH,
326 };
327 use itertools::Itertools;
328 use reth_primitives_traits::Account;
329 use reth_trie_common::{
330 prefix_set::PrefixSetMut, updates::TrieUpdates, BranchNode, HashedPostState, LeafNode,
331 RlpNode,
332 };
333 use std::collections::BTreeMap;
334
335 use super::{TrieElement, TrieNodeIter};
336
337 fn get_hash_builder_branch_nodes(
340 state: impl IntoIterator<Item = (Nibbles, Account)> + Clone,
341 ) -> HashMap<Nibbles, BranchNodeCompact> {
342 let mut hash_builder = HashBuilder::default().with_updates(true);
343
344 let mut prefix_set = PrefixSetMut::default();
345 prefix_set.extend_keys(state.clone().into_iter().map(|(nibbles, _)| nibbles));
346 let walker = TrieWalker::<_>::state_trie(NoopAccountTrieCursor, prefix_set.freeze());
347
348 let hashed_post_state = HashedPostState::default()
349 .with_accounts(state.into_iter().map(|(nibbles, account)| {
350 (nibbles.pack().into_inner().unwrap().into(), Some(account))
351 }))
352 .into_sorted();
353
354 let mut node_iter = TrieNodeIter::state_trie(
355 walker,
356 HashedPostStateAccountCursor::new(
357 NoopHashedAccountCursor::default(),
358 hashed_post_state.accounts(),
359 ),
360 );
361
362 while let Some(node) = node_iter.try_next().unwrap() {
363 match node {
364 TrieElement::Branch(branch) => {
365 hash_builder.add_branch(branch.key, branch.value, branch.children_are_in_trie);
366 }
367 TrieElement::Leaf(key, account) => {
368 hash_builder.add_leaf(
369 Nibbles::unpack(key),
370 &alloy_rlp::encode(account.into_trie_account(EMPTY_ROOT_HASH)),
371 );
372 }
373 }
374 }
375 hash_builder.root();
376
377 let mut trie_updates = TrieUpdates::default();
378 trie_updates.finalize(hash_builder, Default::default(), Default::default());
379
380 trie_updates.account_nodes
381 }
382
383 #[test]
384 fn test_trie_node_iter() {
385 fn empty_leaf_rlp_for_key(key: Nibbles) -> RlpNode {
386 RlpNode::from_rlp(&alloy_rlp::encode(LeafNode::new(
387 key,
388 alloy_rlp::encode(TrieAccount::default()),
389 )))
390 }
391
392 reth_tracing::init_test_tracing();
393
394 let account_1 = b256!("0x0000000000000000000000000000000000000000000000000000000000000000");
406 let account_2 = b256!("0x0000000000000000000000000000000000000000000000000000000000000010");
407 let account_3 = b256!("0x0000000000000000000000000000000000000000000000000000000000000100");
408 let account_4 = b256!("0x0000000000000000000000000000000000000000000000000000000000000101");
409 let account_5 = b256!("0x0000000000000000000000000000000000000000000000000000000000000110");
410 let empty_account = Account::default();
411
412 let hash_builder_branch_nodes = get_hash_builder_branch_nodes(vec![
413 (Nibbles::unpack(account_1), empty_account),
414 (Nibbles::unpack(account_2), empty_account),
415 (Nibbles::unpack(account_3), empty_account),
416 (Nibbles::unpack(account_4), empty_account),
417 (Nibbles::unpack(account_5), empty_account),
418 ]);
419
420 let branch_node_1_rlp = RlpNode::from_rlp(&alloy_rlp::encode(BranchNode::new(
421 vec![
422 empty_leaf_rlp_for_key(Nibbles::from_nibbles([0])),
423 empty_leaf_rlp_for_key(Nibbles::from_nibbles([0])),
424 ],
425 TrieMask::new(0b11),
426 )));
427
428 let branch_node_3_rlp = RlpNode::from_rlp(&alloy_rlp::encode(BranchNode::new(
429 vec![
430 empty_leaf_rlp_for_key(Nibbles::default()),
431 empty_leaf_rlp_for_key(Nibbles::default()),
432 ],
433 TrieMask::new(0b11),
434 )));
435
436 let branch_node_2 = (
437 Nibbles::from_nibbles([vec![0; 61], vec![1]].concat()),
438 BranchNodeCompact::new(
439 TrieMask::new(0b11),
440 TrieMask::new(0b00),
441 TrieMask::new(0b01),
442 vec![branch_node_3_rlp.as_hash().unwrap()],
443 None,
444 ),
445 );
446 let branch_node_2_rlp = RlpNode::from_rlp(&alloy_rlp::encode(BranchNode::new(
447 vec![branch_node_3_rlp, empty_leaf_rlp_for_key(Nibbles::from_nibbles([0]))],
448 TrieMask::new(0b11),
449 )));
450 let branch_node_0 = (
451 Nibbles::from_nibbles([0; 61]),
452 BranchNodeCompact::new(
453 TrieMask::new(0b11),
454 TrieMask::new(0b10),
455 TrieMask::new(0b11),
456 vec![branch_node_1_rlp.as_hash().unwrap(), branch_node_2_rlp.as_hash().unwrap()],
457 None,
458 ),
459 );
460
461 let mock_trie_nodes = vec![branch_node_0.clone(), branch_node_2.clone()];
462 pretty_assertions::assert_eq!(
463 hash_builder_branch_nodes.into_iter().sorted().collect::<Vec<_>>(),
464 mock_trie_nodes,
465 );
466
467 let trie_cursor_factory =
468 MockTrieCursorFactory::new(mock_trie_nodes.into_iter().collect(), B256Map::default());
469
470 let mut prefix_set = PrefixSetMut::default();
472 prefix_set.insert(Nibbles::unpack(account_3));
473 let prefix_set = prefix_set.freeze();
474
475 let walker = TrieWalker::<_>::state_trie(
476 trie_cursor_factory.account_trie_cursor().unwrap(),
477 prefix_set,
478 );
479
480 let hashed_cursor_factory = MockHashedCursorFactory::new(
481 BTreeMap::from([
482 (account_1, empty_account),
483 (account_2, empty_account),
484 (account_3, empty_account),
485 (account_4, empty_account),
486 (account_5, empty_account),
487 ]),
488 B256Map::default(),
489 );
490
491 let mut iter = TrieNodeIter::state_trie(
492 walker,
493 hashed_cursor_factory.hashed_account_cursor().unwrap(),
494 );
495
496 while iter.try_next().unwrap().is_some() {}
498
499 pretty_assertions::assert_eq!(
500 *trie_cursor_factory.visited_account_keys(),
501 vec![
502 KeyVisit {
503 visit_type: KeyVisitType::SeekExact(Nibbles::default()),
504 visited_key: None
505 },
506 KeyVisit {
507 visit_type: KeyVisitType::SeekNonExact(Nibbles::from_nibbles([0x0])),
508 visited_key: Some(branch_node_0.0)
509 },
510 KeyVisit {
511 visit_type: KeyVisitType::SeekNonExact(branch_node_2.0),
512 visited_key: Some(branch_node_2.0)
513 },
514 KeyVisit {
515 visit_type: KeyVisitType::SeekNonExact(Nibbles::from_nibbles([0x1])),
516 visited_key: None
517 }
518 ]
519 );
520 pretty_assertions::assert_eq!(
521 *hashed_cursor_factory.visited_account_keys(),
522 vec![
523 KeyVisit {
525 visit_type: KeyVisitType::SeekNonExact(account_1),
526 visited_key: Some(account_1)
527 },
528 KeyVisit {
530 visit_type: KeyVisitType::SeekNonExact(account_3),
531 visited_key: Some(account_3)
532 },
533 KeyVisit { visit_type: KeyVisitType::Next, visited_key: Some(account_4) },
535 KeyVisit { visit_type: KeyVisitType::Next, visited_key: Some(account_5) },
536 KeyVisit { visit_type: KeyVisitType::Next, visited_key: None },
537 ],
538 );
539 }
540}