reth_trie/
walker.rs

1use crate::{
2    prefix_set::PrefixSet,
3    trie_cursor::{subnode::SubNodePosition, CursorSubNode, TrieCursor},
4    BranchNodeCompact, Nibbles,
5};
6use alloy_primitives::{map::HashSet, B256};
7use alloy_trie::proof::AddedRemovedKeys;
8use reth_storage_errors::db::DatabaseError;
9use tracing::{instrument, trace};
10
11#[cfg(feature = "metrics")]
12use crate::metrics::WalkerMetrics;
13
14/// Traverses the trie in lexicographic order.
15///
16/// This iterator depends on the ordering guarantees of [`TrieCursor`].
17#[derive(Debug)]
18pub struct TrieWalker<C, K = AddedRemovedKeys> {
19    /// A mutable reference to a trie cursor instance used for navigating the trie.
20    pub cursor: C,
21    /// A vector containing the trie nodes that have been visited.
22    pub stack: Vec<CursorSubNode>,
23    /// A flag indicating whether the current node can be skipped when traversing the trie. This
24    /// is determined by whether the current key's prefix is included in the prefix set and if the
25    /// hash flag is set.
26    pub can_skip_current_node: bool,
27    /// A `PrefixSet` representing the changes to be applied to the trie.
28    pub changes: PrefixSet,
29    /// The retained trie node keys that need to be removed.
30    removed_keys: Option<HashSet<Nibbles>>,
31    /// Provided when it's necessary to not skip certain nodes during proof generation.
32    /// Specifically we don't skip certain branch nodes even when they are not in the `PrefixSet`,
33    /// when they might be required to support leaf removal.
34    added_removed_keys: Option<K>,
35    #[cfg(feature = "metrics")]
36    /// Walker metrics.
37    metrics: WalkerMetrics,
38}
39
40impl<C: TrieCursor, K: AsRef<AddedRemovedKeys>> TrieWalker<C, K> {
41    /// Constructs a new `TrieWalker` for the state trie from existing stack and a cursor.
42    pub fn state_trie_from_stack(cursor: C, stack: Vec<CursorSubNode>, changes: PrefixSet) -> Self {
43        Self::from_stack(
44            cursor,
45            stack,
46            changes,
47            #[cfg(feature = "metrics")]
48            crate::TrieType::State,
49        )
50    }
51
52    /// Constructs a new `TrieWalker` for the storage trie from existing stack and a cursor.
53    pub fn storage_trie_from_stack(
54        cursor: C,
55        stack: Vec<CursorSubNode>,
56        changes: PrefixSet,
57    ) -> Self {
58        Self::from_stack(
59            cursor,
60            stack,
61            changes,
62            #[cfg(feature = "metrics")]
63            crate::TrieType::Storage,
64        )
65    }
66
67    /// Constructs a new `TrieWalker` from existing stack and a cursor.
68    fn from_stack(
69        cursor: C,
70        stack: Vec<CursorSubNode>,
71        changes: PrefixSet,
72        #[cfg(feature = "metrics")] trie_type: crate::TrieType,
73    ) -> Self {
74        let mut this = Self {
75            cursor,
76            changes,
77            stack,
78            can_skip_current_node: false,
79            removed_keys: None,
80            added_removed_keys: None,
81            #[cfg(feature = "metrics")]
82            metrics: WalkerMetrics::new(trie_type),
83        };
84        this.update_skip_node();
85        this
86    }
87
88    /// Sets the flag whether the trie updates should be stored.
89    pub fn with_deletions_retained(mut self, retained: bool) -> Self {
90        if retained {
91            self.removed_keys = Some(HashSet::default());
92        }
93        self
94    }
95
96    /// Configures the walker to not skip certain branch nodes, even when they are not in the
97    /// `PrefixSet`, when they might be needed to support leaf removal.
98    pub fn with_added_removed_keys<K2>(self, added_removed_keys: Option<K2>) -> TrieWalker<C, K2> {
99        TrieWalker {
100            cursor: self.cursor,
101            stack: self.stack,
102            can_skip_current_node: self.can_skip_current_node,
103            changes: self.changes,
104            removed_keys: self.removed_keys,
105            added_removed_keys,
106            #[cfg(feature = "metrics")]
107            metrics: self.metrics,
108        }
109    }
110
111    /// Split the walker into stack and trie updates.
112    pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<Nibbles>) {
113        let keys = self.take_removed_keys();
114        (self.stack, keys)
115    }
116
117    /// Take removed keys from the walker.
118    pub fn take_removed_keys(&mut self) -> HashSet<Nibbles> {
119        self.removed_keys.take().unwrap_or_default()
120    }
121
122    /// Prints the current stack of trie nodes.
123    pub fn print_stack(&self) {
124        println!("====================== STACK ======================");
125        for node in &self.stack {
126            println!("{node:?}");
127        }
128        println!("====================== END STACK ======================\n");
129    }
130
131    /// The current length of the removed keys.
132    pub fn removed_keys_len(&self) -> usize {
133        self.removed_keys.as_ref().map_or(0, |u| u.len())
134    }
135
136    /// Returns the current key in the trie.
137    pub fn key(&self) -> Option<&Nibbles> {
138        self.stack.last().map(|n| n.full_key())
139    }
140
141    /// Returns the current hash in the trie, if any.
142    pub fn hash(&self) -> Option<B256> {
143        self.stack.last().and_then(|n| n.hash())
144    }
145
146    /// Returns the current hash in the trie, if any.
147    ///
148    /// Differs from [`Self::hash`] in that it returns `None` if the subnode is positioned at the
149    /// child without a hash mask bit set. [`Self::hash`] panics in that case.
150    pub fn maybe_hash(&self) -> Option<B256> {
151        self.stack.last().and_then(|n| n.maybe_hash())
152    }
153
154    /// Indicates whether the children of the current node are present in the trie.
155    pub fn children_are_in_trie(&self) -> bool {
156        self.stack.last().is_some_and(|n| n.tree_flag())
157    }
158
159    /// Returns the next unprocessed key in the trie along with its raw [`Nibbles`] representation.
160    #[instrument(level = "trace", skip(self), ret)]
161    pub fn next_unprocessed_key(&self) -> Option<(B256, Nibbles)> {
162        self.key()
163            .and_then(|key| if self.can_skip_current_node { key.increment() } else { Some(*key) })
164            .map(|key| {
165                let mut packed = key.pack();
166                packed.resize(32, 0);
167                (B256::from_slice(packed.as_slice()), key)
168            })
169    }
170
171    /// Updates the skip node flag based on the walker's current state.
172    fn update_skip_node(&mut self) {
173        let old = self.can_skip_current_node;
174        self.can_skip_current_node = self.stack.last().is_some_and(|node| {
175            // If the current key is not removed according to the [`AddedRemovedKeys`], and all of
176            // its siblings are removed, then we don't want to skip it. This allows the
177            // `ProofRetainer` to include this node in the returned proofs. Required to support
178            // leaf removal.
179            let key_is_only_nonremoved_child =
180                self.added_removed_keys.as_ref().is_some_and(|added_removed_keys| {
181                    node.full_key_is_only_nonremoved_child(added_removed_keys.as_ref())
182                });
183
184            trace!(
185                target: "trie::walker",
186                ?key_is_only_nonremoved_child,
187                full_key=?node.full_key(),
188                "Checked for only nonremoved child",
189            );
190
191            !self.changes.contains(node.full_key()) &&
192                node.hash_flag() &&
193                !key_is_only_nonremoved_child
194        });
195        trace!(
196            target: "trie::walker",
197            old,
198            new = self.can_skip_current_node,
199            last = ?self.stack.last(),
200            "updated skip node flag"
201        );
202    }
203
204    /// Constructs a new [`TrieWalker`] for the state trie.
205    pub fn state_trie(cursor: C, changes: PrefixSet) -> Self {
206        Self::new(
207            cursor,
208            changes,
209            #[cfg(feature = "metrics")]
210            crate::TrieType::State,
211        )
212    }
213
214    /// Constructs a new [`TrieWalker`] for the storage trie.
215    pub fn storage_trie(cursor: C, changes: PrefixSet) -> Self {
216        Self::new(
217            cursor,
218            changes,
219            #[cfg(feature = "metrics")]
220            crate::TrieType::Storage,
221        )
222    }
223
224    /// Constructs a new `TrieWalker`, setting up the initial state of the stack and cursor.
225    fn new(
226        cursor: C,
227        changes: PrefixSet,
228        #[cfg(feature = "metrics")] trie_type: crate::TrieType,
229    ) -> Self {
230        // Initialize the walker with a single empty stack element.
231        let mut this = Self {
232            cursor,
233            changes,
234            stack: vec![CursorSubNode::default()],
235            can_skip_current_node: false,
236            removed_keys: None,
237            added_removed_keys: Default::default(),
238            #[cfg(feature = "metrics")]
239            metrics: WalkerMetrics::new(trie_type),
240        };
241
242        // Set up the root node of the trie in the stack, if it exists.
243        if let Some((key, value)) = this.node(true).unwrap() {
244            this.stack[0] = CursorSubNode::new(key, Some(value));
245        }
246
247        // Update the skip state for the root node.
248        this.update_skip_node();
249        this
250    }
251
252    /// Advances the walker to the next trie node and updates the skip node flag.
253    /// The new key can then be obtained via `key()`.
254    ///
255    /// # Returns
256    ///
257    /// * `Result<(), Error>` - Unit on success or an error.
258    pub fn advance(&mut self) -> Result<(), DatabaseError> {
259        if let Some(last) = self.stack.last() {
260            if !self.can_skip_current_node && self.children_are_in_trie() {
261                trace!(
262                    target: "trie::walker",
263                    position = ?last.position(),
264                    "cannot skip current node and children are in the trie"
265                );
266                // If we can't skip the current node and the children are in the trie,
267                // either consume the next node or move to the next sibling.
268                match last.position() {
269                    SubNodePosition::ParentBranch => self.move_to_next_sibling(true)?,
270                    SubNodePosition::Child(_) => self.consume_node()?,
271                }
272            } else {
273                trace!(target: "trie::walker", "can skip current node");
274                // If we can skip the current node, move to the next sibling.
275                self.move_to_next_sibling(false)?;
276            }
277
278            // Update the skip node flag based on the new position in the trie.
279            self.update_skip_node();
280        }
281
282        Ok(())
283    }
284
285    /// Retrieves the current root node from the DB, seeking either the exact node or the next one.
286    fn node(&mut self, exact: bool) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError> {
287        let key = self.key().expect("key must exist");
288        let entry = if exact { self.cursor.seek_exact(*key)? } else { self.cursor.seek(*key)? };
289        #[cfg(feature = "metrics")]
290        self.metrics.inc_branch_nodes_seeked();
291
292        if let Some((_, node)) = &entry {
293            assert!(!node.state_mask.is_empty());
294        }
295
296        Ok(entry)
297    }
298
299    /// Consumes the next node in the trie, updating the stack.
300    #[instrument(level = "trace", skip(self), ret)]
301    fn consume_node(&mut self) -> Result<(), DatabaseError> {
302        let Some((key, node)) = self.node(false)? else {
303            // If no next node is found, clear the stack.
304            self.stack.clear();
305            return Ok(())
306        };
307
308        // Overwrite the root node's first nibble
309        // We need to sync the stack with the trie structure when consuming a new node. This is
310        // necessary for proper traversal and accurately representing the trie in the stack.
311        if !key.is_empty() && !self.stack.is_empty() {
312            self.stack[0].set_nibble(key.get_unchecked(0));
313        }
314
315        // The current tree mask might have been set incorrectly.
316        // Sanity check that the newly retrieved trie node key is the child of the last item
317        // on the stack. If not, advance to the next sibling instead of adding the node to the
318        // stack.
319        if let Some(subnode) = self.stack.last() {
320            if !key.starts_with(subnode.full_key()) {
321                #[cfg(feature = "metrics")]
322                self.metrics.inc_out_of_order_subnode(1);
323                self.move_to_next_sibling(false)?;
324                return Ok(())
325            }
326        }
327
328        // Create a new CursorSubNode and push it to the stack.
329        let subnode = CursorSubNode::new(key, Some(node));
330        let position = subnode.position();
331        self.stack.push(subnode);
332        self.update_skip_node();
333
334        // Delete the current node if it's included in the prefix set or it doesn't contain the root
335        // hash.
336        if !self.can_skip_current_node || position.is_child() {
337            if let Some((keys, key)) = self.removed_keys.as_mut().zip(self.cursor.current()?) {
338                keys.insert(key);
339            }
340        }
341
342        Ok(())
343    }
344
345    /// Moves to the next sibling node in the trie, updating the stack.
346    #[instrument(level = "trace", skip(self), ret)]
347    fn move_to_next_sibling(
348        &mut self,
349        allow_root_to_child_nibble: bool,
350    ) -> Result<(), DatabaseError> {
351        let Some(subnode) = self.stack.last_mut() else { return Ok(()) };
352
353        // Check if the walker needs to backtrack to the previous level in the trie during its
354        // traversal.
355        if subnode.position().is_last_child() ||
356            (subnode.position().is_parent() && !allow_root_to_child_nibble)
357        {
358            self.stack.pop();
359            self.move_to_next_sibling(false)?;
360            return Ok(())
361        }
362
363        subnode.inc_nibble();
364
365        if subnode.node.is_none() {
366            return self.consume_node()
367        }
368
369        // Find the next sibling with state.
370        loop {
371            let position = subnode.position();
372            if subnode.state_flag() {
373                trace!(target: "trie::walker", ?position, "found next sibling with state");
374                return Ok(())
375            }
376            if position.is_last_child() {
377                trace!(target: "trie::walker", ?position, "checked all siblings");
378                break
379            }
380            subnode.inc_nibble();
381        }
382
383        // Pop the current node and move to the next sibling.
384        self.stack.pop();
385        self.move_to_next_sibling(false)?;
386
387        Ok(())
388    }
389}