reth_trie/trie_cursor/
subnode.rs

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

reth_trie/trie_cursor/
subnode.rs
