reth_trie/trie_cursor/
subnode.rs

Help
1use crate::{BranchNodeCompact, Nibbles, StoredSubNode, CHILD_INDEX_RANGE};
2use alloy_primitives::B256;
3
4/// Cursor for iterating over a subtrie.
5#[derive(Clone)]
6pub struct CursorSubNode {
7    /// The key of the current node.
8    pub key: Nibbles,
9    /// The position of the current subnode.
10    position: SubNodePosition,
11    /// The node itself.
12    pub node: Option<BranchNodeCompact>,
13    /// Full key
14    full_key: Nibbles,
15}
16
17impl Default for CursorSubNode {
18    fn default() -> Self {
19        Self::new(Nibbles::default(), None)
20    }
21}
22
23impl std::fmt::Debug for CursorSubNode {
24    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
25        f.debug_struct("CursorSubNode")
26            .field("key", &self.key)
27            .field("position", &self.position)
28            .field("state_flag", &self.state_flag())
29            .field("tree_flag", &self.tree_flag())
30            .field("hash_flag", &self.hash_flag())
31            .field("hash", &self.hash())
32            .finish()
33    }
34}
35
36/// Implements conversion from `StoredSubNode` to `CursorSubNode`.
37impl From<StoredSubNode> for CursorSubNode {
38    /// Converts a `StoredSubNode` into a `CursorSubNode`.
39    ///
40    /// Extracts necessary values from the `StoredSubNode` and constructs
41    /// a corresponding `CursorSubNode`.
42    fn from(value: StoredSubNode) -> Self {
43        let position = value.nibble.map_or(SubNodePosition::ParentBranch, SubNodePosition::Child);
44        let key = Nibbles::from_nibbles_unchecked(value.key);
45        Self::new_with_full_key(key, value.node, position)
46    }
47}
48
49impl From<CursorSubNode> for StoredSubNode {
50    fn from(value: CursorSubNode) -> Self {
51        Self { key: value.key.to_vec(), nibble: value.position.as_child(), node: value.node }
52    }
53}
54
55impl CursorSubNode {
56    /// Creates a new [`CursorSubNode`] from a key and an optional node.
57    pub fn new(key: Nibbles, node: Option<BranchNodeCompact>) -> Self {
58        // Find the first nibble that is set in the state mask of the node.
59        let position = node.as_ref().filter(|n| n.root_hash.is_none()).map_or(
60            SubNodePosition::ParentBranch,
61            |n| {
62                SubNodePosition::Child(
63                    CHILD_INDEX_RANGE.clone().find(|i| n.state_mask.is_bit_set(*i)).unwrap(),
64                )
65            },
66        );
67        Self::new_with_full_key(key, node, position)
68    }
69
70    /// Creates a new [`CursorSubNode`] and sets the full key according to the provided key and
71    /// position.
72    fn new_with_full_key(
73        key: Nibbles,
74        node: Option<BranchNodeCompact>,
75        position: SubNodePosition,
76    ) -> Self {
77        let mut full_key = key.clone();
78        if let Some(nibble) = position.as_child() {
79            full_key.push(nibble);
80        }
81        Self { key, node, position, full_key }
82    }
83
84    /// Returns the full key of the current node.
85    #[inline]
86    pub const fn full_key(&self) -> &Nibbles {
87        &self.full_key
88    }
89
90    /// Updates the full key by replacing or appending a child nibble based on the old subnode
91    /// position.
92    #[inline]
93    fn update_full_key(&mut self, old_position: SubNodePosition) {
94        if let Some(new_nibble) = self.position.as_child() {
95            if old_position.is_child() {
96                let last_index = self.full_key.len() - 1;
97                self.full_key.set_at(last_index, new_nibble);
98            } else {
99                self.full_key.push(new_nibble);
100            }
101        } else if old_position.is_child() {
102            self.full_key.pop();
103        }
104    }
105
106    /// Returns `true` if either of these:
107    /// - No current node is set.
108    /// - The current node is a parent branch node.
109    /// - The current node is a child with state mask bit set in the parent branch node.
110    #[inline]
111    pub fn state_flag(&self) -> bool {
112        self.node.as_ref().is_none_or(|node| {
113            self.position.as_child().is_none_or(|nibble| node.state_mask.is_bit_set(nibble))
114        })
115    }
116
117    /// Returns `true` if either of these:
118    /// - No current node is set.
119    /// - The current node is a parent branch node.
120    /// - The current node is a child with tree mask bit set in the parent branch node.
121    #[inline]
122    pub fn tree_flag(&self) -> bool {
123        self.node.as_ref().is_none_or(|node| {
124            self.position.as_child().is_none_or(|nibble| node.tree_mask.is_bit_set(nibble))
125        })
126    }
127
128    /// Returns `true` if the hash for the current node is set.
129    ///
130    /// It means either of two:
131    /// - Current node is a parent branch node, and it has a root hash.
132    /// - Current node is a child node, and it has a hash mask bit set in the parent branch node.
133    pub fn hash_flag(&self) -> bool {
134        self.node.as_ref().is_some_and(|node| match self.position {
135            // Check if the parent branch node has a root hash
136            SubNodePosition::ParentBranch => node.root_hash.is_some(),
137            // Or get it from the children
138            SubNodePosition::Child(nibble) => node.hash_mask.is_bit_set(nibble),
139        })
140    }
141
142    /// Returns the hash of the current node.
143    ///
144    /// It means either of two:
145    /// - Root hash of the parent branch node.
146    /// - Hash of the child node at the current nibble, if it has a hash mask bit set in the parent
147    ///   branch node.
148    pub fn hash(&self) -> Option<B256> {
149        self.node.as_ref().and_then(|node| match self.position {
150            // Get the root hash for the parent branch node
151            SubNodePosition::ParentBranch => node.root_hash,
152            // Or get it from the children
153            SubNodePosition::Child(nibble) => Some(node.hash_for_nibble(nibble)),
154        })
155    }
156
157    /// Returns the position to the current node.
158    #[inline]
159    pub const fn position(&self) -> SubNodePosition {
160        self.position
161    }
162
163    /// Increments the nibble index.
164    #[inline]
165    pub fn inc_nibble(&mut self) {
166        let old_position = self.position;
167        self.position.increment();
168        self.update_full_key(old_position);
169    }
170
171    /// Sets the nibble index.
172    #[inline]
173    pub fn set_nibble(&mut self, nibble: u8) {
174        let old_position = self.position;
175        self.position = SubNodePosition::Child(nibble);
176        self.update_full_key(old_position);
177    }
178}
179
180/// Represents a subnode position.
181#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
182pub enum SubNodePosition {
183    /// Positioned at the parent branch node.
184    ParentBranch,
185    /// Positioned at a child node at the given nibble.
186    Child(u8),
187}
188
189impl SubNodePosition {
190    /// Returns `true` if the position is set to the parent branch node.
191    pub const fn is_parent(&self) -> bool {
192        matches!(self, Self::ParentBranch)
193    }
194
195    /// Returns `true` if the position is set to a child node.
196    pub const fn is_child(&self) -> bool {
197        matches!(self, Self::Child(_))
198    }
199
200    /// Returns the nibble of the child node if the position is set to a child node.
201    pub const fn as_child(&self) -> Option<u8> {
202        match self {
203            Self::Child(nibble) => Some(*nibble),
204            _ => None,
205        }
206    }
207
208    /// Returns `true` if the position is set to a last child nibble (i.e. greater than or equal to
209    /// 0xf).
210    pub const fn is_last_child(&self) -> bool {
211        match self {
212            Self::ParentBranch => false,
213            Self::Child(nibble) => *nibble >= 0xf,
214        }
215    }
216
217    const fn increment(&mut self) {
218        match self {
219            Self::ParentBranch => *self = Self::Child(0),
220            Self::Child(nibble) => *nibble += 1,
221        }
222    }
223}
224
225#[cfg(test)]
226mod tests {
227    use super::*;
228
229    #[test]
230    fn subnode_position_ord() {
231        assert!([
232            SubNodePosition::ParentBranch,
233            SubNodePosition::Child(0),
234            SubNodePosition::Child(1),
235            SubNodePosition::Child(2),
236            SubNodePosition::Child(3),
237            SubNodePosition::Child(4),
238            SubNodePosition::Child(5),
239            SubNodePosition::Child(6),
240            SubNodePosition::Child(7),
241            SubNodePosition::Child(8),
242            SubNodePosition::Child(9),
243            SubNodePosition::Child(10),
244            SubNodePosition::Child(11),
245            SubNodePosition::Child(12),
246            SubNodePosition::Child(13),
247            SubNodePosition::Child(14),
248            SubNodePosition::Child(15),
249        ]
250        .is_sorted());
251    }
252
253    #[test]
254    fn subnode_position_is_last_child() {
255        assert!([
256            SubNodePosition::ParentBranch,
257            SubNodePosition::Child(0),
258            SubNodePosition::Child(1),
259            SubNodePosition::Child(2),
260            SubNodePosition::Child(3),
261            SubNodePosition::Child(4),
262            SubNodePosition::Child(5),
263            SubNodePosition::Child(6),
264            SubNodePosition::Child(7),
265            SubNodePosition::Child(8),
266            SubNodePosition::Child(9),
267            SubNodePosition::Child(10),
268            SubNodePosition::Child(11),
269            SubNodePosition::Child(12),
270            SubNodePosition::Child(13),
271            SubNodePosition::Child(14),
272        ]
273        .iter()
274        .all(|position| !position.is_last_child()));
275        assert!(SubNodePosition::Child(15).is_last_child());
276        assert!(SubNodePosition::Child(16).is_last_child());
277    }
278}
reth_trie/trie_cursor/subnode.rs

reth_trie/trie_cursor/
subnode.rs
