reth_trie/
node_iter.rs

Help
1use crate::{hashed_cursor::HashedCursor, trie_cursor::TrieCursor, walker::TrieWalker, Nibbles};
2use alloy_primitives::B256;
3use reth_storage_errors::db::DatabaseError;
4
5/// Represents a branch node in the trie.
6#[derive(Debug)]
7pub struct TrieBranchNode {
8    /// The key associated with the node.
9    pub key: Nibbles,
10    /// The value associated with the node.
11    pub value: B256,
12    /// Indicates whether children are in the trie.
13    pub children_are_in_trie: bool,
14}
15
16impl TrieBranchNode {
17    /// Creates a new `TrieBranchNode`.
18    pub const fn new(key: Nibbles, value: B256, children_are_in_trie: bool) -> Self {
19        Self { key, value, children_are_in_trie }
20    }
21}
22
23/// Represents variants of trie nodes returned by the iteration.
24#[derive(Debug)]
25pub enum TrieElement<Value> {
26    /// Branch node.
27    Branch(TrieBranchNode),
28    /// Leaf node.
29    Leaf(B256, Value),
30}
31
32/// An iterator over existing intermediate branch nodes and updated leaf nodes.
33#[derive(Debug)]
34pub struct TrieNodeIter<C, H: HashedCursor> {
35    /// The walker over intermediate nodes.
36    pub walker: TrieWalker<C>,
37    /// The cursor for the hashed entries.
38    pub hashed_cursor: H,
39    /// The previous hashed key. If the iteration was previously interrupted, this value can be
40    /// used to resume iterating from the last returned leaf node.
41    previous_hashed_key: Option<B256>,
42
43    /// Current hashed  entry.
44    current_hashed_entry: Option<(B256, <H as HashedCursor>::Value)>,
45    /// Flag indicating whether we should check the current walker key.
46    current_walker_key_checked: bool,
47}
48
49impl<C, H: HashedCursor> TrieNodeIter<C, H> {
50    /// Creates a new [`TrieNodeIter`].
51    pub const fn new(walker: TrieWalker<C>, hashed_cursor: H) -> Self {
52        Self {
53            walker,
54            hashed_cursor,
55            previous_hashed_key: None,
56            current_hashed_entry: None,
57            current_walker_key_checked: false,
58        }
59    }
60
61    /// Sets the last iterated hashed key and returns the modified [`TrieNodeIter`].
62    /// This is used to resume iteration from the last checkpoint.
63    pub const fn with_last_hashed_key(mut self, previous_hashed_key: B256) -> Self {
64        self.previous_hashed_key = Some(previous_hashed_key);
65        self
66    }
67}
68
69impl<C, H> TrieNodeIter<C, H>
70where
71    C: TrieCursor,
72    H: HashedCursor,
73{
74    /// Return the next trie node to be added to the hash builder.
75    ///
76    /// Returns the nodes using this algorithm:
77    /// 1. Return the current intermediate branch node if it hasn't been updated.
78    /// 2. Advance the trie walker to the next intermediate branch node and retrieve next
79    ///    unprocessed key.
80    /// 3. Reposition the hashed cursor on the next unprocessed key.
81    /// 4. Return every hashed entry up to the key of the current intermediate branch node.
82    /// 5. Repeat.
83    ///
84    /// NOTE: The iteration will start from the key of the previous hashed entry if it was supplied.
85    pub fn try_next(
86        &mut self,
87    ) -> Result<Option<TrieElement<<H as HashedCursor>::Value>>, DatabaseError> {
88        loop {
89            // If the walker has a key...
90            if let Some(key) = self.walker.key() {
91                // Check if the current walker key is unchecked and there's no previous hashed key
92                if !self.current_walker_key_checked && self.previous_hashed_key.is_none() {
93                    self.current_walker_key_checked = true;
94                    // If it's possible to skip the current node in the walker, return a branch node
95                    if self.walker.can_skip_current_node {
96                        return Ok(Some(TrieElement::Branch(TrieBranchNode::new(
97                            key.clone(),
98                            self.walker.hash().unwrap(),
99                            self.walker.children_are_in_trie(),
100                        ))))
101                    }
102                }
103            }
104
105            // If there's a hashed entry...
106            if let Some((hashed_key, value)) = self.current_hashed_entry.take() {
107                // If the walker's key is less than the unpacked hashed key,
108                // reset the checked status and continue
109                if self.walker.key().is_some_and(|key| key < &Nibbles::unpack(hashed_key)) {
110                    self.current_walker_key_checked = false;
111                    continue
112                }
113
114                // Set the next hashed entry as a leaf node and return
115                self.current_hashed_entry = self.hashed_cursor.next()?;
116                return Ok(Some(TrieElement::Leaf(hashed_key, value)))
117            }
118
119            // Handle seeking and advancing based on the previous hashed key
120            match self.previous_hashed_key.take() {
121                Some(hashed_key) => {
122                    // Seek to the previous hashed key and get the next hashed entry
123                    self.hashed_cursor.seek(hashed_key)?;
124                    self.current_hashed_entry = self.hashed_cursor.next()?;
125                }
126                None => {
127                    // Get the seek key and set the current hashed entry based on walker's next
128                    // unprocessed key
129                    let seek_key = match self.walker.next_unprocessed_key() {
130                        Some(key) => key,
131                        None => break, // no more keys
132                    };
133                    self.current_hashed_entry = self.hashed_cursor.seek(seek_key)?;
134                    self.walker.advance()?;
135                }
136            }
137        }
138
139        Ok(None)
140    }
141}
reth_trie/node_iter.rs

reth_trie/
node_iter.rs
