reth_trie/
node_iter.rs

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/// Represents a branch node in the trie.
10#[derive(Debug)]
11pub struct TrieBranchNode {
12    /// The key associated with the node.
13    pub key: Nibbles,
14    /// The value associated with the node.
15    pub value: B256,
16    /// Indicates whether children are in the trie.
17    pub children_are_in_trie: bool,
18}
19
20impl TrieBranchNode {
21    /// Creates a new `TrieBranchNode`.
22    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/// Represents variants of trie nodes returned by the iteration.
28#[derive(Debug)]
29pub enum TrieElement<Value> {
30    /// Branch node.
31    Branch(TrieBranchNode),
32    /// Leaf node.
33    Leaf(B256, Value),
34}
35
36/// Result of calling [`HashedCursor::seek`].
37#[derive(Debug)]
38struct SeekedHashedEntry<V> {
39    /// The key that was seeked.
40    seeked_key: B256,
41    /// The result of the seek.
42
43    /// If no entry was found for the provided key, this will be [`None`].
44    result: Option<(B256, V)>,
45}
46
47/// Iterates over trie nodes for hash building.
48///
49/// This iterator depends on the ordering guarantees of [`TrieCursor`],
50/// and additionally uses hashed cursor lookups when operating on storage tries.
51#[derive(Debug)]
52pub struct TrieNodeIter<C, H: HashedCursor, K> {
53    /// The walker over intermediate nodes.
54    pub walker: TrieWalker<C, K>,
55    /// The cursor for the hashed entries.
56    pub hashed_cursor: H,
57    /// The type of the trie.
58    trie_type: TrieType,
59    /// The previous hashed key. If the iteration was previously interrupted, this value can be
60    /// used to resume iterating from the last returned leaf node.
61    previous_hashed_key: Option<B256>,
62
63    /// Current hashed  entry.
64    current_hashed_entry: Option<(B256, H::Value)>,
65    /// Flag indicating whether we should check the current walker key.
66    should_check_walker_key: bool,
67
68    /// The last seeked hashed entry.
69    ///
70    /// We use it to not seek the same hashed entry twice, and instead reuse it.
71    last_seeked_hashed_entry: Option<SeekedHashedEntry<H::Value>>,
72
73    #[cfg(feature = "metrics")]
74    metrics: crate::metrics::TrieNodeIterMetrics,
75    /// Stores the result of the last successful [`Self::next_hashed_entry`], used to avoid a
76    /// redundant [`Self::seek_hashed_entry`] call if the walker points to the same key that
77    /// was just returned by `next()`.
78    last_next_result: Option<(B256, H::Value)>,
79}
80
81impl<C, H: HashedCursor, K> TrieNodeIter<C, H, K>
82where
83    H::Value: Clone,
84    K: AsRef<AddedRemovedKeys>,
85{
86    /// Creates a new [`TrieNodeIter`] for the state trie.
87    pub fn state_trie(walker: TrieWalker<C, K>, hashed_cursor: H) -> Self {
88        Self::new(walker, hashed_cursor, TrieType::State)
89    }
90
91    /// Creates a new [`TrieNodeIter`] for the storage trie.
92    pub fn storage_trie(walker: TrieWalker<C, K>, hashed_cursor: H) -> Self {
93        Self::new(walker, hashed_cursor, TrieType::Storage)
94    }
95
96    /// Creates a new [`TrieNodeIter`].
97    #[allow(clippy::missing_const_for_fn)]
98    fn new(walker: TrieWalker<C, K>, hashed_cursor: H, trie_type: TrieType) -> Self {
99        Self {
100            walker,
101            hashed_cursor,
102            trie_type,
103            previous_hashed_key: None,
104            current_hashed_entry: None,
105            should_check_walker_key: false,
106            last_seeked_hashed_entry: None,
107            #[cfg(feature = "metrics")]
108            metrics: crate::metrics::TrieNodeIterMetrics::new(trie_type),
109            last_next_result: None,
110        }
111    }
112
113    /// Sets the last iterated hashed key and returns the modified [`TrieNodeIter`].
114    /// This is used to resume iteration from the last checkpoint.
115    pub const fn with_last_hashed_key(mut self, previous_hashed_key: B256) -> Self {
116        self.previous_hashed_key = Some(previous_hashed_key);
117        self
118    }
119
120    /// Seeks the hashed cursor to the given key.
121    ///
122    /// If the key is the same as the last seeked key, the result of the last seek is returned.
123    ///
124    /// If `metrics` feature is enabled, it also updates the metrics.
125    fn seek_hashed_entry(&mut self, key: B256) -> Result<Option<(B256, H::Value)>, DatabaseError> {
126        if let Some((last_key, _)) = self.last_next_result.as_ref() &&
127            *last_key == key
128        {
129            trace!(target: "trie::node_iter", seek_key = ?key, "reusing result from last next() call instead of seeking");
130            let result = self.last_next_result.take();
131            self.last_seeked_hashed_entry =
132                Some(SeekedHashedEntry { seeked_key: key, result: result.clone() });
133
134            return Ok(result);
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.clone())
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 =
151            Some(SeekedHashedEntry { seeked_key: key, result: result.clone() });
152
153        #[cfg(feature = "metrics")]
154        {
155            self.metrics.inc_leaf_nodes_seeked();
156        }
157        Ok(result)
158    }
159
160    /// Advances the hashed cursor to the next entry.
161    ///
162    /// If `metrics` feature is enabled, it also updates the metrics.
163    fn next_hashed_entry(&mut self) -> Result<Option<(B256, H::Value)>, DatabaseError> {
164        let next = self.hashed_cursor.next()?;
165
166        self.last_next_result = next.clone();
167
168        #[cfg(feature = "metrics")]
169        {
170            self.metrics.inc_leaf_nodes_advanced();
171        }
172        Ok(next)
173    }
174}
175
176impl<C, H, K> TrieNodeIter<C, H, K>
177where
178    C: TrieCursor,
179    H: HashedCursor,
180    H::Value: Clone,
181    K: AsRef<AddedRemovedKeys>,
182{
183    /// Return the next trie node to be added to the hash builder.
184    ///
185    /// Returns the nodes using this algorithm:
186    /// 1. Return the current intermediate branch node if it hasn't been updated.
187    /// 2. Advance the trie walker to the next intermediate branch node and retrieve next
188    ///    unprocessed key.
189    /// 3. Reposition the hashed cursor on the next unprocessed key.
190    /// 4. Return every hashed entry up to the key of the current intermediate branch node.
191    /// 5. Repeat.
192    ///
193    /// NOTE: The iteration will start from the key of the previous hashed entry if it was supplied.
194    #[instrument(
195        level = "trace",
196        target = "trie::node_iter",
197        skip_all,
198        fields(trie_type = ?self.trie_type),
199        ret
200    )]
201    pub fn try_next(
202        &mut self,
203    ) -> Result<Option<TrieElement<<H as HashedCursor>::Value>>, DatabaseError> {
204        loop {
205            // If the walker has a key...
206            if let Some(key) = self.walker.key() {
207                // Ensure that the current walker key shouldn't be checked and there's no previous
208                // hashed key
209                if !self.should_check_walker_key && self.previous_hashed_key.is_none() {
210                    // Make sure we check the next walker key, because we only know we can skip the
211                    // current one.
212                    self.should_check_walker_key = true;
213                    // If it's possible to skip the current node in the walker, return a branch node
214                    if self.walker.can_skip_current_node {
215                        #[cfg(feature = "metrics")]
216                        self.metrics.inc_branch_nodes_returned();
217                        return Ok(Some(TrieElement::Branch(TrieBranchNode::new(
218                            *key,
219                            self.walker.hash().unwrap(),
220                            self.walker.children_are_in_trie(),
221                        ))))
222                    }
223                }
224            }
225
226            // If there's a hashed entry...
227            if let Some((hashed_key, value)) = self.current_hashed_entry.take() {
228                // Check if the walker's key is less than the key of the current hashed entry
229                if self.walker.key().is_some_and(|key| key < &Nibbles::unpack(hashed_key)) {
230                    self.should_check_walker_key = false;
231                    continue
232                }
233
234                // Set the next hashed entry as a leaf node and return
235                trace!(target: "trie::node_iter", ?hashed_key, "next hashed entry");
236                self.current_hashed_entry = self.next_hashed_entry()?;
237
238                #[cfg(feature = "metrics")]
239                self.metrics.inc_leaf_nodes_returned();
240                return Ok(Some(TrieElement::Leaf(hashed_key, value)))
241            }
242
243            // Handle seeking and advancing based on the previous hashed key
244            match self.previous_hashed_key.take() {
245                Some(hashed_key) => {
246                    trace!(target: "trie::node_iter", ?hashed_key, "seeking to the previous hashed entry");
247                    // Seek to the previous hashed key and get the next hashed entry
248                    self.seek_hashed_entry(hashed_key)?;
249                    self.current_hashed_entry = self.next_hashed_entry()?;
250                }
251                None => {
252                    // Get the seek key and set the current hashed entry based on walker's next
253                    // unprocessed key
254                    let (seek_key, seek_prefix) = match self.walker.next_unprocessed_key() {
255                        Some(key) => key,
256                        None => break, // no more keys
257                    };
258
259                    trace!(
260                        target: "trie::node_iter",
261                        ?seek_key,
262                        can_skip_current_node = self.walker.can_skip_current_node,
263                        last = ?self.walker.stack.last(),
264                        "seeking to the next unprocessed hashed entry"
265                    );
266                    let can_skip_node = self.walker.can_skip_current_node;
267                    self.walker.advance()?;
268                    trace!(
269                        target: "trie::node_iter",
270                        last = ?self.walker.stack.last(),
271                        "advanced walker"
272                    );
273
274                    // We should get the iterator to return a branch node if we can skip the
275                    // current node and the tree flag for the current node is set.
276                    //
277                    // `can_skip_node` is already set when the hash flag is set, so we don't need
278                    // to check for the hash flag explicitly.
279                    //
280                    // A deeper stored key may leave unvisited hashed entries between the seek
281                    // prefix and the walker, so only skip the seek at the exact prefix.
282                    if can_skip_node &&
283                        self.walker.key().is_some_and(|key| key == &seek_prefix) &&
284                        self.walker.children_are_in_trie()
285                    {
286                        trace!(
287                            target: "trie::node_iter",
288                            ?seek_key,
289                            walker_hash = ?self.walker.maybe_hash(),
290                            "skipping hashed seek"
291                        );
292
293                        self.should_check_walker_key = false;
294                        continue
295                    }
296
297                    self.current_hashed_entry = self.seek_hashed_entry(seek_key)?;
298                }
299            }
300        }
301
302        Ok(None)
303    }
304}
305
306#[cfg(test)]
307mod tests {
308    use super::*;
309    use crate::{
310        hashed_cursor::{
311            mock::MockHashedCursorFactory,
312            noop::{NoopHashedCursor, NoopHashedCursorFactory},
313            HashedCursorFactory, HashedPostStateCursor, HashedPostStateCursorFactory,
314        },
315        mock::{KeyVisit, KeyVisitType},
316        trie_cursor::{
317            mock::MockTrieCursorFactory,
318            noop::{NoopAccountTrieCursor, NoopTrieCursorFactory},
319            InMemoryTrieCursorFactory, TrieCursorFactory,
320        },
321        walker::TrieWalker,
322        StateRoot,
323    };
324    use alloy_primitives::{
325        b256,
326        map::{B256Map, HashMap},
327        B256, U256,
328    };
329    use alloy_trie::{
330        BranchNodeCompact, HashBuilder, Nibbles, TrieAccount, TrieMask, EMPTY_ROOT_HASH,
331    };
332    use itertools::Itertools;
333    use reth_primitives_traits::Account;
334    use reth_trie_common::{
335        prefix_set::PrefixSetMut, updates::TrieUpdates, BranchNode, HashedPostState, HashedStorage,
336        LeafNode, RlpNode,
337    };
338    use std::collections::BTreeMap;
339
340    /// Calculate the branch node stored in the database by feeding the provided state to the hash
341    /// builder and taking the trie updates.
342    fn get_hash_builder_branch_nodes(
343        state: impl IntoIterator<Item = (Nibbles, Account)> + Clone,
344    ) -> HashMap<Nibbles, BranchNodeCompact> {
345        let mut hash_builder = HashBuilder::default().with_updates(true);
346
347        let mut prefix_set = PrefixSetMut::default();
348        prefix_set.extend_keys(state.clone().into_iter().map(|(nibbles, _)| nibbles));
349        let walker =
350            TrieWalker::<_>::state_trie(NoopAccountTrieCursor, prefix_set.freeze()).unwrap();
351
352        let hashed_post_state = HashedPostState::default()
353            .with_accounts(state.into_iter().map(|(nibbles, account)| {
354                (nibbles.pack().into_inner().unwrap().into(), Some(account))
355            }))
356            .into_sorted();
357
358        let mut node_iter = TrieNodeIter::state_trie(
359            walker,
360            HashedPostStateCursor::new_account(
361                NoopHashedCursor::<Account>::default(),
362                &hashed_post_state,
363            ),
364        );
365
366        while let Some(node) = node_iter.try_next().unwrap() {
367            match node {
368                TrieElement::Branch(branch) => {
369                    hash_builder.add_branch(branch.key, branch.value, branch.children_are_in_trie);
370                }
371                TrieElement::Leaf(key, account) => {
372                    hash_builder.add_leaf(
373                        Nibbles::unpack(key),
374                        &alloy_rlp::encode(account.into_trie_account(EMPTY_ROOT_HASH)),
375                    );
376                }
377            }
378        }
379        hash_builder.root();
380
381        let mut trie_updates = TrieUpdates::default();
382        trie_updates.finalize(hash_builder, Default::default(), Default::default());
383
384        trie_updates.account_nodes
385    }
386
387    #[test]
388    fn inserted_leaf_before_deeper_stored_branch_is_included() {
389        fn slot(prefix: &[u8]) -> B256 {
390            let mut bytes = [0; 32];
391            bytes[..prefix.len()].copy_from_slice(prefix);
392            B256::from(bytes)
393        }
394
395        let address = B256::repeat_byte(0xaa);
396        let account = || Account { nonce: 1, ..Default::default() };
397        let old_slots = [
398            slot(&[0x00]),
399            slot(&[0x0f, 0x00]),
400            slot(&[0x0f, 0x10]),
401            slot(&[0x10, 0xa0, 0x00]),
402            slot(&[0x10, 0xa0, 0x01]),
403            slot(&[0x10, 0xa0, 0x10]),
404            slot(&[0x10, 0xa1]),
405        ];
406        let inserted = slot(&[0x10, 0x00]);
407        let old_state =
408            HashedPostState::default().with_accounts([(address, Some(account()))]).with_storages([
409                (address, HashedStorage::from_iter(old_slots.map(|slot| (slot, U256::from(1))))),
410            ]);
411        let old_sorted = old_state.clone().into_sorted();
412        let (old_root, old_updates) = StateRoot::new(
413            NoopTrieCursorFactory::default(),
414            HashedPostStateCursorFactory::new(NoopHashedCursorFactory::default(), &old_sorted),
415        )
416        .with_prefix_sets(old_state.construct_prefix_sets().freeze())
417        .root_with_updates()
418        .unwrap();
419
420        let new_state = HashedPostState::default()
421            .with_accounts([(address, Some(account()))])
422            .with_storages([(
423                address,
424                HashedStorage::from_iter(
425                    old_slots.into_iter().chain([inserted]).map(|slot| (slot, U256::from(1))),
426                ),
427            )]);
428        let delta = HashedPostState::default()
429            .with_accounts([(address, Some(account()))])
430            .with_storages([(address, HashedStorage::from_iter([(inserted, U256::from(1))]))]);
431        let new_sorted = new_state.clone().into_sorted();
432        let rebuilt_root = StateRoot::new(
433            NoopTrieCursorFactory::default(),
434            HashedPostStateCursorFactory::new(NoopHashedCursorFactory::default(), &new_sorted),
435        )
436        .with_prefix_sets(new_state.construct_prefix_sets().freeze())
437        .root()
438        .unwrap();
439        let old_trie = old_updates.into_sorted();
440        let incremental_root = StateRoot::new(
441            InMemoryTrieCursorFactory::new(NoopTrieCursorFactory::default(), &old_trie),
442            HashedPostStateCursorFactory::new(NoopHashedCursorFactory::default(), &new_sorted),
443        )
444        .with_prefix_sets(delta.construct_prefix_sets().freeze())
445        .root()
446        .unwrap();
447
448        assert_ne!(old_root, rebuilt_root);
449        assert_eq!(incremental_root, rebuilt_root);
450    }
451
452    #[test]
453    fn test_trie_node_iter() {
454        fn empty_leaf_rlp_for_key(key: Nibbles) -> RlpNode {
455            RlpNode::from_rlp(&alloy_rlp::encode(LeafNode::new(
456                key,
457                alloy_rlp::encode(TrieAccount::default()),
458            )))
459        }
460
461        reth_tracing::init_test_tracing();
462
463        // Extension (Key = 0x0000000000000000000000000000000000000000000000000000000000000)
464        // └── Branch (`branch_node_0`)
465        //     ├── 0 -> Branch (`branch_node_1`)
466        //     │      ├── 0 -> Leaf (`account_1`, Key = 0x0)
467        //     │      └── 1 -> Leaf (`account_2`, Key = 0x0)
468        //     ├── 1 -> Branch (`branch_node_2`)
469        //     │      ├── 0 -> Branch (`branch_node_3`)
470        //     │      │      ├── 0 -> Leaf (`account_3`, marked as changed)
471        //     │      │      └── 1 -> Leaf (`account_4`)
472        //     │      └── 1 -> Leaf (`account_5`, Key = 0x0)
473
474        let account_1 = B256::ZERO;
475        let account_2 = B256::with_last_byte(0x10);
476        let account_3 = b256!("0x0000000000000000000000000000000000000000000000000000000000000100");
477        let account_4 = b256!("0x0000000000000000000000000000000000000000000000000000000000000101");
478        let account_5 = b256!("0x0000000000000000000000000000000000000000000000000000000000000110");
479
480        let hash_builder_branch_nodes = get_hash_builder_branch_nodes(vec![
481            (Nibbles::unpack(account_1), Account::default()),
482            (Nibbles::unpack(account_2), Account::default()),
483            (Nibbles::unpack(account_3), Account::default()),
484            (Nibbles::unpack(account_4), Account::default()),
485            (Nibbles::unpack(account_5), Account::default()),
486        ]);
487
488        let branch_node_1_rlp = RlpNode::from_rlp(&alloy_rlp::encode(BranchNode::new(
489            vec![
490                empty_leaf_rlp_for_key(Nibbles::from_nibbles([0])),
491                empty_leaf_rlp_for_key(Nibbles::from_nibbles([0])),
492            ],
493            TrieMask::new(0b11),
494        )));
495
496        let branch_node_3_rlp = RlpNode::from_rlp(&alloy_rlp::encode(BranchNode::new(
497            vec![
498                empty_leaf_rlp_for_key(Nibbles::default()),
499                empty_leaf_rlp_for_key(Nibbles::default()),
500            ],
501            TrieMask::new(0b11),
502        )));
503
504        let branch_node_2 = (
505            Nibbles::from_nibbles([vec![0; 61], vec![1]].concat()),
506            BranchNodeCompact::new(
507                TrieMask::new(0b11),
508                TrieMask::new(0b00),
509                TrieMask::new(0b01),
510                vec![branch_node_3_rlp.as_hash().unwrap()],
511                None,
512            ),
513        );
514        let branch_node_2_rlp = RlpNode::from_rlp(&alloy_rlp::encode(BranchNode::new(
515            vec![branch_node_3_rlp, empty_leaf_rlp_for_key(Nibbles::from_nibbles([0]))],
516            TrieMask::new(0b11),
517        )));
518        let branch_node_0 = (
519            Nibbles::from_nibbles([0; 61]),
520            BranchNodeCompact::new(
521                TrieMask::new(0b11),
522                TrieMask::new(0b10),
523                TrieMask::new(0b11),
524                vec![branch_node_1_rlp.as_hash().unwrap(), branch_node_2_rlp.as_hash().unwrap()],
525                None,
526            ),
527        );
528
529        let mock_trie_nodes = vec![branch_node_0.clone(), branch_node_2.clone()];
530        pretty_assertions::assert_eq!(
531            hash_builder_branch_nodes.into_iter().sorted().collect::<Vec<_>>(),
532            mock_trie_nodes,
533        );
534
535        let trie_cursor_factory =
536            MockTrieCursorFactory::new(mock_trie_nodes.into_iter().collect(), B256Map::default());
537
538        // Mark the account 3 as changed.
539        let mut prefix_set = PrefixSetMut::default();
540        prefix_set.insert(Nibbles::unpack(account_3));
541        let prefix_set = prefix_set.freeze();
542
543        let walker = TrieWalker::<_>::state_trie(
544            trie_cursor_factory.account_trie_cursor().unwrap(),
545            prefix_set,
546        )
547        .unwrap();
548
549        let hashed_cursor_factory = MockHashedCursorFactory::new(
550            BTreeMap::from([
551                (account_1, Account::default()),
552                (account_2, Account::default()),
553                (account_3, Account::default()),
554                (account_4, Account::default()),
555                (account_5, Account::default()),
556            ]),
557            B256Map::default(),
558        );
559
560        let mut iter = TrieNodeIter::state_trie(
561            walker,
562            hashed_cursor_factory.hashed_account_cursor().unwrap(),
563        );
564
565        // Walk the iterator until it's exhausted.
566        while iter.try_next().unwrap().is_some() {}
567
568        pretty_assertions::assert_eq!(
569            *trie_cursor_factory.visited_account_keys(),
570            vec![
571                KeyVisit {
572                    visit_type: KeyVisitType::SeekExact(Nibbles::default()),
573                    visited_key: None
574                },
575                KeyVisit {
576                    visit_type: KeyVisitType::SeekNonExact(Nibbles::from_nibbles([0x0])),
577                    visited_key: Some(branch_node_0.0)
578                },
579                KeyVisit {
580                    visit_type: KeyVisitType::SeekNonExact(branch_node_2.0),
581                    visited_key: Some(branch_node_2.0)
582                },
583                KeyVisit {
584                    visit_type: KeyVisitType::SeekNonExact(Nibbles::from_nibbles([0x1])),
585                    visited_key: None
586                }
587            ]
588        );
589        pretty_assertions::assert_eq!(
590            *hashed_cursor_factory.visited_account_keys(),
591            vec![
592                // Why do we always seek this key first?
593                KeyVisit {
594                    visit_type: KeyVisitType::SeekNonExact(account_1),
595                    visited_key: Some(account_1)
596                },
597                // Seek to the modified account.
598                KeyVisit {
599                    visit_type: KeyVisitType::SeekNonExact(account_3),
600                    visited_key: Some(account_3)
601                },
602                // Collect the siblings of the modified account
603                KeyVisit { visit_type: KeyVisitType::Next, visited_key: Some(account_4) },
604                KeyVisit { visit_type: KeyVisitType::Next, visited_key: Some(account_5) },
605                KeyVisit { visit_type: KeyVisitType::Next, visited_key: None },
606            ],
607        );
608    }
609}
reth_trie/node_iter.rs

reth_trie/
node_iter.rs
