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: Copy,
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    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    /// Sets the last iterated hashed key and returns the modified [`TrieNodeIter`].
113    /// This is used to resume iteration from the last checkpoint.
114    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    /// Seeks the hashed cursor to the given key.
120    ///
121    /// If the key is the same as the last seeked key, the result of the last seek is returned.
122    ///
123    /// If `metrics` feature is enabled, also updates the metrics.
124    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; // Consume the cached value
129
130                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    /// Advances the hashed cursor to the next entry.
160    ///
161    /// If `metrics` feature is enabled, also updates the metrics.
162    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    /// Return the next trie node to be added to the hash builder.
183    ///
184    /// Returns the nodes using this algorithm:
185    /// 1. Return the current intermediate branch node if it hasn't been updated.
186    /// 2. Advance the trie walker to the next intermediate branch node and retrieve next
187    ///    unprocessed key.
188    /// 3. Reposition the hashed cursor on the next unprocessed key.
189    /// 4. Return every hashed entry up to the key of the current intermediate branch node.
190    /// 5. Repeat.
191    ///
192    /// NOTE: The iteration will start from the key of the previous hashed entry if it was supplied.
193    #[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 the walker has a key...
205            if let Some(key) = self.walker.key() {
206                // Ensure that the current walker key shouldn't be checked and there's no previous
207                // hashed key
208                if !self.should_check_walker_key && self.previous_hashed_key.is_none() {
209                    // Make sure we check the next walker key, because we only know we can skip the
210                    // current one.
211                    self.should_check_walker_key = true;
212                    // If it's possible to skip the current node in the walker, return a branch node
213                    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 there's a hashed entry...
226            if let Some((hashed_key, value)) = self.current_hashed_entry.take() {
227                // Check if the walker's key is less than the key of the current hashed entry
228                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                // Set the next hashed entry as a leaf node and return
234                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            // Handle seeking and advancing based on the previous hashed key
243            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                    // Seek to the previous hashed key and get the next hashed entry
247                    self.seek_hashed_entry(hashed_key)?;
248                    self.current_hashed_entry = self.next_hashed_entry()?;
249                }
250                None => {
251                    // Get the seek key and set the current hashed entry based on walker's next
252                    // unprocessed key
253                    let (seek_key, seek_prefix) = match self.walker.next_unprocessed_key() {
254                        Some(key) => key,
255                        None => break, // no more keys
256                    };
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                    // We should get the iterator to return a branch node if we can skip the
274                    // current node and the tree flag for the current node is set.
275                    //
276                    // `can_skip_node` is already set when the hash flag is set, so we don't need
277                    // to check for the hash flag explicitly.
278                    //
279                    // It is possible that the branch node at the key `seek_key` is not stored in
280                    // the database, so the walker will advance to the branch node after it. Because
281                    // of this, we need to check that the current walker key has a prefix of the key
282                    // that we seeked to.
283                    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    /// Calculate the branch node stored in the database by feeding the provided state to the hash
338    /// builder and taking the trie updates.
339    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        // Extension (Key = 0x0000000000000000000000000000000000000000000000000000000000000)
395        // └── Branch (`branch_node_0`)
396        //     ├── 0 -> Branch (`branch_node_1`)
397        //     │      ├── 0 -> Leaf (`account_1`, Key = 0x0)
398        //     │      └── 1 -> Leaf (`account_2`, Key = 0x0)
399        //     ├── 1 -> Branch (`branch_node_2`)
400        //     │      ├── 0 -> Branch (`branch_node_3`)
401        //     │      │      ├── 0 -> Leaf (`account_3`, marked as changed)
402        //     │      │      └── 1 -> Leaf (`account_4`)
403        //     │      └── 1 -> Leaf (`account_5`, Key = 0x0)
404
405        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        // Mark the account 3 as changed.
471        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        // Walk the iterator until it's exhausted.
497        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                // Why do we always seek this key first?
524                KeyVisit {
525                    visit_type: KeyVisitType::SeekNonExact(account_1),
526                    visited_key: Some(account_1)
527                },
528                // Seek to the modified account.
529                KeyVisit {
530                    visit_type: KeyVisitType::SeekNonExact(account_3),
531                    visited_key: Some(account_3)
532                },
533                // Collect the siblings of the modified account
534                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}