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#[derive(Debug)]
18pub struct TrieWalker<C, K = AddedRemovedKeys> {
19 pub cursor: C,
21 pub stack: Vec<CursorSubNode>,
23 pub can_skip_current_node: bool,
27 pub changes: PrefixSet,
29 removed_keys: Option<HashSet<Nibbles>>,
31 added_removed_keys: Option<K>,
35 #[cfg(feature = "metrics")]
36 metrics: WalkerMetrics,
38}
39
40impl<C: TrieCursor, K: AsRef<AddedRemovedKeys>> TrieWalker<C, K> {
41 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 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 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 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 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 pub fn split(mut self) -> (Vec<CursorSubNode>, HashSet<Nibbles>) {
113 let keys = self.take_removed_keys();
114 (self.stack, keys)
115 }
116
117 pub fn take_removed_keys(&mut self) -> HashSet<Nibbles> {
119 self.removed_keys.take().unwrap_or_default()
120 }
121
122 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 pub fn removed_keys_len(&self) -> usize {
133 self.removed_keys.as_ref().map_or(0, |u| u.len())
134 }
135
136 pub fn key(&self) -> Option<&Nibbles> {
138 self.stack.last().map(|n| n.full_key())
139 }
140
141 pub fn hash(&self) -> Option<B256> {
143 self.stack.last().and_then(|n| n.hash())
144 }
145
146 pub fn maybe_hash(&self) -> Option<B256> {
151 self.stack.last().and_then(|n| n.maybe_hash())
152 }
153
154 pub fn children_are_in_trie(&self) -> bool {
156 self.stack.last().is_some_and(|n| n.tree_flag())
157 }
158
159 #[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 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 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 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 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 fn new(
226 cursor: C,
227 changes: PrefixSet,
228 #[cfg(feature = "metrics")] trie_type: crate::TrieType,
229 ) -> Self {
230 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 if let Some((key, value)) = this.node(true).unwrap() {
244 this.stack[0] = CursorSubNode::new(key, Some(value));
245 }
246
247 this.update_skip_node();
249 this
250 }
251
252 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 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 self.move_to_next_sibling(false)?;
276 }
277
278 self.update_skip_node();
280 }
281
282 Ok(())
283 }
284
285 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 #[instrument(level = "trace", skip(self), ret)]
301 fn consume_node(&mut self) -> Result<(), DatabaseError> {
302 let Some((key, node)) = self.node(false)? else {
303 self.stack.clear();
305 return Ok(())
306 };
307
308 if !key.is_empty() && !self.stack.is_empty() {
312 self.stack[0].set_nibble(key.get_unchecked(0));
313 }
314
315 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 let subnode = CursorSubNode::new(key, Some(node));
330 let position = subnode.position();
331 self.stack.push(subnode);
332 self.update_skip_node();
333
334 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 #[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 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 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 self.stack.pop();
385 self.move_to_next_sibling(false)?;
386
387 Ok(())
388 }
389}