reth_trie/trie_cursor/
subnode.rs

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.maybe_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 hash of the current node, if any.
158    ///
159    /// Differs from [`Self::hash`] in that it returns `None` if the subnode is positioned at the
160    /// child without a hash mask bit set. [`Self::hash`] panics in that case.
161    pub fn maybe_hash(&self) -> Option<B256> {
162        self.node.as_ref().and_then(|node| match self.position {
163            // Get the root hash for the parent branch node
164            SubNodePosition::ParentBranch => node.root_hash,
165            // Or get it from the children
166            SubNodePosition::Child(nibble) => {
167                node.hash_mask.is_bit_set(nibble).then(|| node.hash_for_nibble(nibble))
168            }
169        })
170    }
171
172    /// Returns the position to the current node.
173    #[inline]
174    pub const fn position(&self) -> SubNodePosition {
175        self.position
176    }
177
178    /// Increments the nibble index.
179    #[inline]
180    pub fn inc_nibble(&mut self) {
181        let old_position = self.position;
182        self.position.increment();
183        self.update_full_key(old_position);
184    }
185
186    /// Sets the nibble index.
187    #[inline]
188    pub fn set_nibble(&mut self, nibble: u8) {
189        let old_position = self.position;
190        self.position = SubNodePosition::Child(nibble);
191        self.update_full_key(old_position);
192    }
193}
194
195/// Represents a subnode position.
196#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
197pub enum SubNodePosition {
198    /// Positioned at the parent branch node.
199    ParentBranch,
200    /// Positioned at a child node at the given nibble.
201    Child(u8),
202}
203
204impl SubNodePosition {
205    /// Returns `true` if the position is set to the parent branch node.
206    pub const fn is_parent(&self) -> bool {
207        matches!(self, Self::ParentBranch)
208    }
209
210    /// Returns `true` if the position is set to a child node.
211    pub const fn is_child(&self) -> bool {
212        matches!(self, Self::Child(_))
213    }
214
215    /// Returns the nibble of the child node if the position is set to a child node.
216    pub const fn as_child(&self) -> Option<u8> {
217        match self {
218            Self::Child(nibble) => Some(*nibble),
219            _ => None,
220        }
221    }
222
223    /// Returns `true` if the position is set to a last child nibble (i.e. greater than or equal to
224    /// 0xf).
225    pub const fn is_last_child(&self) -> bool {
226        match self {
227            Self::ParentBranch => false,
228            Self::Child(nibble) => *nibble >= 0xf,
229        }
230    }
231
232    const fn increment(&mut self) {
233        match self {
234            Self::ParentBranch => *self = Self::Child(0),
235            Self::Child(nibble) => *nibble += 1,
236        }
237    }
238}
239
240#[cfg(test)]
241mod tests {
242    use super::*;
243
244    #[test]
245    fn subnode_position_ord() {
246        assert!([
247            SubNodePosition::ParentBranch,
248            SubNodePosition::Child(0),
249            SubNodePosition::Child(1),
250            SubNodePosition::Child(2),
251            SubNodePosition::Child(3),
252            SubNodePosition::Child(4),
253            SubNodePosition::Child(5),
254            SubNodePosition::Child(6),
255            SubNodePosition::Child(7),
256            SubNodePosition::Child(8),
257            SubNodePosition::Child(9),
258            SubNodePosition::Child(10),
259            SubNodePosition::Child(11),
260            SubNodePosition::Child(12),
261            SubNodePosition::Child(13),
262            SubNodePosition::Child(14),
263            SubNodePosition::Child(15),
264        ]
265        .is_sorted());
266    }
267
268    #[test]
269    fn subnode_position_is_last_child() {
270        assert!([
271            SubNodePosition::ParentBranch,
272            SubNodePosition::Child(0),
273            SubNodePosition::Child(1),
274            SubNodePosition::Child(2),
275            SubNodePosition::Child(3),
276            SubNodePosition::Child(4),
277            SubNodePosition::Child(5),
278            SubNodePosition::Child(6),
279            SubNodePosition::Child(7),
280            SubNodePosition::Child(8),
281            SubNodePosition::Child(9),
282            SubNodePosition::Child(10),
283            SubNodePosition::Child(11),
284            SubNodePosition::Child(12),
285            SubNodePosition::Child(13),
286            SubNodePosition::Child(14),
287        ]
288        .iter()
289        .all(|position| !position.is_last_child()));
290        assert!(SubNodePosition::Child(15).is_last_child());
291        assert!(SubNodePosition::Child(16).is_last_child());
292    }
293}
reth_trie/trie_cursor/subnode.rs

reth_trie/trie_cursor/
subnode.rs
