reth_trie/trie_cursor/
subnode.rs

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

reth_trie/trie_cursor/
subnode.rs
