Skip to main content

reth_trie/trie_cursor/
mod.rs

1use crate::{BranchNodeCompact, Nibbles};
2use alloy_primitives::B256;
3use reth_storage_errors::db::DatabaseError;
4
5/// In-memory implementations of trie cursors.
6mod in_memory;
7
8/// Cursor for iterating over a subtrie.
9pub mod subnode;
10
11/// Noop trie cursor implementations.
12pub mod noop;
13
14/// Masked trie cursor wrapper that skips nodes matching a prefix set.
15pub mod masked;
16
17/// Depth-first trie iterator.
18pub mod depth_first;
19
20/// Mock trie cursor implementations.
21#[cfg(any(test, feature = "test-utils"))]
22pub mod mock;
23
24/// Metrics tracking trie cursor implementations.
25pub mod metrics;
26#[cfg(feature = "metrics")]
27pub use metrics::TrieCursorMetrics;
28pub use metrics::{InstrumentedTrieCursor, TrieCursorMetricsCache};
29
30pub use self::{depth_first::DepthFirstTrieIterator, in_memory::*, subnode::CursorSubNode};
31
32/// Factory for creating trie cursors.
33#[auto_impl::auto_impl(&)]
34pub trait TrieCursorFactory {
35    /// The account trie cursor type.
36    type AccountTrieCursor<'a>: TrieCursor
37    where
38        Self: 'a;
39
40    /// The storage trie cursor type.
41    type StorageTrieCursor<'a>: TrieStorageCursor
42    where
43        Self: 'a;
44
45    /// Create an account trie cursor.
46    fn account_trie_cursor(&self) -> Result<Self::AccountTrieCursor<'_>, DatabaseError>;
47
48    /// Create a storage tries cursor.
49    fn storage_trie_cursor(
50        &self,
51        hashed_address: B256,
52    ) -> Result<Self::StorageTrieCursor<'_>, DatabaseError>;
53}
54
55/// A cursor for traversing stored trie nodes. The cursor must iterate over keys in
56/// lexicographical order.
57#[auto_impl::auto_impl(&mut, Box)]
58pub trait TrieCursor {
59    /// Move the cursor to the key and return if it is an exact match.
60    fn seek_exact(
61        &mut self,
62        key: Nibbles,
63    ) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError>;
64
65    /// Move the cursor to the key and return a value matching of greater than the key.
66    fn seek(&mut self, key: Nibbles)
67        -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError>;
68
69    /// Move the cursor to the next key.
70    fn next(&mut self) -> Result<Option<(Nibbles, BranchNodeCompact)>, DatabaseError>;
71
72    /// Get the current entry.
73    fn current(&mut self) -> Result<Option<Nibbles>, DatabaseError>;
74
75    /// Reset the cursor to the beginning.
76    ///
77    /// # Important
78    ///
79    /// After calling this method, the subsequent operation MUST be a [`TrieCursor::seek`] or
80    /// [`TrieCursor::seek_exact`] call.
81    fn reset(&mut self);
82}
83
84/// A cursor for traversing storage trie nodes.
85#[auto_impl::auto_impl(&mut, Box)]
86pub trait TrieStorageCursor: TrieCursor {
87    /// Set the hashed address for the storage trie cursor.
88    ///
89    /// # Important
90    ///
91    /// After calling this method, the subsequent operation MUST be a [`TrieCursor::seek`] or
92    /// [`TrieCursor::seek_exact`] call.
93    fn set_hashed_address(&mut self, hashed_address: B256);
94}
95
96/// Iterator wrapper for `TrieCursor` types
97#[derive(Debug)]
98pub struct TrieCursorIter<'a, C> {
99    cursor: &'a mut C,
100    /// The initial value from seek, if any
101    initial: Option<Result<(Nibbles, BranchNodeCompact), DatabaseError>>,
102}
103
104impl<'a, C> TrieCursorIter<'a, C> {
105    /// Create a new iterator from a mutable reference to a cursor. The Iterator will start from the
106    /// empty path.
107    pub fn new(cursor: &'a mut C) -> Self
108    where
109        C: TrieCursor,
110    {
111        let initial = cursor.seek(Nibbles::default()).transpose();
112        Self { cursor, initial }
113    }
114}
115
116impl<'a, C> From<&'a mut C> for TrieCursorIter<'a, C>
117where
118    C: TrieCursor,
119{
120    fn from(cursor: &'a mut C) -> Self {
121        Self::new(cursor)
122    }
123}
124
125impl<'a, C> Iterator for TrieCursorIter<'a, C>
126where
127    C: TrieCursor,
128{
129    type Item = Result<(Nibbles, BranchNodeCompact), DatabaseError>;
130
131    fn next(&mut self) -> Option<Self::Item> {
132        // If we have an initial value from seek, return it first
133        if let Some(initial) = self.initial.take() {
134            return Some(initial);
135        }
136
137        self.cursor.next().transpose()
138    }
139}