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 index of the next child to visit.
10    nibble: i8,
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("nibble", &self.nibble)
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 nibble = value.nibble.map_or(-1, |n| n as i8);
44        let key = Nibbles::from_nibbles_unchecked(value.key);
45        let full_key = full_key(key.clone(), nibble);
46        Self { key, nibble, node: value.node, full_key }
47    }
48}
49
50impl From<CursorSubNode> for StoredSubNode {
51    fn from(value: CursorSubNode) -> Self {
52        let nibble = (value.nibble >= 0).then_some(value.nibble as u8);
53        Self { key: value.key.to_vec(), nibble, node: value.node }
54    }
55}
56
57impl CursorSubNode {
58    /// Creates a new `CursorSubNode` from a key and an optional node.
59    pub fn new(key: Nibbles, node: Option<BranchNodeCompact>) -> Self {
60        // Find the first nibble that is set in the state mask of the node.
61        let nibble = node.as_ref().filter(|n| n.root_hash.is_none()).map_or(-1, |n| {
62            CHILD_INDEX_RANGE.clone().find(|i| n.state_mask.is_bit_set(*i)).unwrap() as i8
63        });
64        let full_key = full_key(key.clone(), nibble);
65        Self { key, node, nibble, full_key }
66    }
67
68    /// Returns the full key of the current node.
69    #[inline]
70    pub const fn full_key(&self) -> &Nibbles {
71        &self.full_key
72    }
73
74    /// Returns `true` if the state flag is set for the current nibble.
75    #[inline]
76    pub fn state_flag(&self) -> bool {
77        self.node
78            .as_ref()
79            .is_none_or(|node| self.nibble < 0 || node.state_mask.is_bit_set(self.nibble as u8))
80    }
81
82    /// Returns `true` if the tree flag is set for the current nibble.
83    #[inline]
84    pub fn tree_flag(&self) -> bool {
85        self.node
86            .as_ref()
87            .is_none_or(|node| self.nibble < 0 || node.tree_mask.is_bit_set(self.nibble as u8))
88    }
89
90    /// Returns `true` if the current nibble has a root hash.
91    pub fn hash_flag(&self) -> bool {
92        self.node.as_ref().is_some_and(|node| match self.nibble {
93            // This guy has it
94            -1 => node.root_hash.is_some(),
95            // Or get it from the children
96            _ => node.hash_mask.is_bit_set(self.nibble as u8),
97        })
98    }
99
100    /// Returns the root hash of the current node, if it has one.
101    pub fn hash(&self) -> Option<B256> {
102        if self.hash_flag() {
103            let node = self.node.as_ref().unwrap();
104            match self.nibble {
105                -1 => node.root_hash,
106                _ => Some(node.hash_for_nibble(self.nibble as u8)),
107            }
108        } else {
109            None
110        }
111    }
112
113    /// Returns the next child index to visit.
114    #[inline]
115    pub const fn nibble(&self) -> i8 {
116        self.nibble
117    }
118
119    /// Increments the nibble index.
120    #[inline]
121    pub fn inc_nibble(&mut self) {
122        self.nibble += 1;
123        update_full_key(&mut self.full_key, self.nibble - 1, self.nibble);
124    }
125
126    /// Sets the nibble index.
127    #[inline]
128    pub fn set_nibble(&mut self, nibble: i8) {
129        let old_nibble = self.nibble;
130        self.nibble = nibble;
131        update_full_key(&mut self.full_key, old_nibble, self.nibble);
132    }
133}
134
135/// Constructs a full key from the given `Nibbles` and `nibble`.
136#[inline]
137fn full_key(mut key: Nibbles, nibble: i8) -> Nibbles {
138    if nibble >= 0 {
139        key.push(nibble as u8);
140    }
141    key
142}
143
144/// Updates the key by replacing or appending a nibble based on the old and new nibble values.
145#[inline]
146fn update_full_key(key: &mut Nibbles, old_nibble: i8, new_nibble: i8) {
147    if new_nibble >= 0 {
148        if old_nibble >= 0 {
149            let last_index = key.len() - 1;
150            key.set_at(last_index, new_nibble as u8);
151        } else {
152            key.push(new_nibble as u8);
153        }
154    } else if old_nibble >= 0 {
155        key.pop();
156    }
157}
reth_trie/trie_cursor/subnode.rs

reth_trie/trie_cursor/
subnode.rs
