diff --git a/crates/trie/trie/src/trie_cursor/mod.rs b/crates/trie/trie/src/trie_cursor/mod.rs index 6cf788a781..69f7d1caec 100644 --- a/crates/trie/trie/src/trie_cursor/mod.rs +++ b/crates/trie/trie/src/trie_cursor/mod.rs @@ -6,7 +6,7 @@ use reth_storage_errors::db::DatabaseError; mod in_memory; /// Cursor for iterating over a subtrie. -mod subnode; +pub mod subnode; /// Noop trie cursor implementations. pub mod noop; diff --git a/crates/trie/trie/src/trie_cursor/subnode.rs b/crates/trie/trie/src/trie_cursor/subnode.rs index 457c1ba468..98dce055d5 100644 --- a/crates/trie/trie/src/trie_cursor/subnode.rs +++ b/crates/trie/trie/src/trie_cursor/subnode.rs @@ -6,8 +6,8 @@ use alloy_primitives::B256; pub struct CursorSubNode { /// The key of the current node. pub key: Nibbles, - /// The index of the next child to visit. - nibble: i8, + /// The position of the current subnode. + position: SubNodePosition, /// The node itself. pub node: Option, /// Full key @@ -24,7 +24,7 @@ impl std::fmt::Debug for CursorSubNode { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.debug_struct("CursorSubNode") .field("key", &self.key) - .field("nibble", &self.nibble) + .field("position", &self.position) .field("state_flag", &self.state_flag()) .field("tree_flag", &self.tree_flag()) .field("hash_flag", &self.hash_flag()) @@ -40,29 +40,45 @@ impl From for CursorSubNode { /// Extracts necessary values from the `StoredSubNode` and constructs /// a corresponding `CursorSubNode`. fn from(value: StoredSubNode) -> Self { - let nibble = value.nibble.map_or(-1, |n| n as i8); + let position = value.nibble.map_or(SubNodePosition::ParentBranch, SubNodePosition::Child); let key = Nibbles::from_nibbles_unchecked(value.key); - let full_key = full_key(key.clone(), nibble); - Self { key, nibble, node: value.node, full_key } + Self::new_with_full_key(key, value.node, position) } } impl From for StoredSubNode { fn from(value: CursorSubNode) -> Self { - let nibble = (value.nibble >= 0).then_some(value.nibble as u8); - Self { key: value.key.to_vec(), nibble, node: value.node } + Self { key: value.key.to_vec(), nibble: value.position.as_child(), node: value.node } } } impl CursorSubNode { - /// Creates a new `CursorSubNode` from a key and an optional node. + /// Creates a new [`CursorSubNode`] from a key and an optional node. pub fn new(key: Nibbles, node: Option) -> Self { // Find the first nibble that is set in the state mask of the node. - let nibble = node.as_ref().filter(|n| n.root_hash.is_none()).map_or(-1, |n| { - CHILD_INDEX_RANGE.clone().find(|i| n.state_mask.is_bit_set(*i)).unwrap() as i8 - }); - let full_key = full_key(key.clone(), nibble); - Self { key, node, nibble, full_key } + let position = node.as_ref().filter(|n| n.root_hash.is_none()).map_or( + SubNodePosition::ParentBranch, + |n| { + SubNodePosition::Child( + CHILD_INDEX_RANGE.clone().find(|i| n.state_mask.is_bit_set(*i)).unwrap(), + ) + }, + ); + Self::new_with_full_key(key, node, position) + } + + /// Creates a new [`CursorSubNode`] and sets the full key according to the provided key and + /// position. + fn new_with_full_key( + key: Nibbles, + node: Option, + position: SubNodePosition, + ) -> Self { + let mut full_key = key.clone(); + if let Some(nibble) = position.as_child() { + full_key.push(nibble); + } + Self { key, node, position, full_key } } /// Returns the full key of the current node. @@ -71,87 +87,192 @@ impl CursorSubNode { &self.full_key } - /// Returns `true` if the state flag is set for the current nibble. + /// Updates the full key by replacing or appending a child nibble based on the old subnode + /// position. #[inline] - pub fn state_flag(&self) -> bool { - self.node - .as_ref() - .is_none_or(|node| self.nibble < 0 || node.state_mask.is_bit_set(self.nibble as u8)) - } - - /// Returns `true` if the tree flag is set for the current nibble. - #[inline] - pub fn tree_flag(&self) -> bool { - self.node - .as_ref() - .is_none_or(|node| self.nibble < 0 || node.tree_mask.is_bit_set(self.nibble as u8)) - } - - /// Returns `true` if the current nibble has a root hash. - pub fn hash_flag(&self) -> bool { - self.node.as_ref().is_some_and(|node| match self.nibble { - // This guy has it - -1 => node.root_hash.is_some(), - // Or get it from the children - _ => node.hash_mask.is_bit_set(self.nibble as u8), - }) - } - - /// Returns the root hash of the current node, if it has one. - pub fn hash(&self) -> Option { - if self.hash_flag() { - let node = self.node.as_ref().unwrap(); - match self.nibble { - -1 => node.root_hash, - _ => Some(node.hash_for_nibble(self.nibble as u8)), + fn update_full_key(&mut self, old_position: SubNodePosition) { + if let Some(new_nibble) = self.position.as_child() { + if old_position.is_child() { + let last_index = self.full_key.len() - 1; + self.full_key.set_at(last_index, new_nibble); + } else { + self.full_key.push(new_nibble); } - } else { - None + } else if old_position.is_child() { + self.full_key.pop(); } } - /// Returns the next child index to visit. + /// Returns `true` if either of these: + /// - No current node is set. + /// - The current node is a parent branch node. + /// - The current node is a child with state mask bit set in the parent branch node. #[inline] - pub const fn nibble(&self) -> i8 { - self.nibble + pub fn state_flag(&self) -> bool { + self.node.as_ref().is_none_or(|node| { + self.position.as_child().is_none_or(|nibble| node.state_mask.is_bit_set(nibble)) + }) + } + + /// Returns `true` if either of these: + /// - No current node is set. + /// - The current node is a parent branch node. + /// - The current node is a child with tree mask bit set in the parent branch node. + #[inline] + pub fn tree_flag(&self) -> bool { + self.node.as_ref().is_none_or(|node| { + self.position.as_child().is_none_or(|nibble| node.tree_mask.is_bit_set(nibble)) + }) + } + + /// Returns `true` if the hash for the current node is set. + /// + /// It means either of two: + /// - Current node is a parent branch node, and it has a root hash. + /// - Current node is a child node, and it has a hash mask bit set in the parent branch node. + pub fn hash_flag(&self) -> bool { + self.node.as_ref().is_some_and(|node| match self.position { + // Check if the parent branch node has a root hash + SubNodePosition::ParentBranch => node.root_hash.is_some(), + // Or get it from the children + SubNodePosition::Child(nibble) => node.hash_mask.is_bit_set(nibble), + }) + } + + /// Returns the hash of the current node. + /// + /// It means either of two: + /// - Root hash of the parent branch node. + /// - Hash of the child node at the current nibble, if it has a hash mask bit set in the parent + /// branch node. + pub fn hash(&self) -> Option { + self.node.as_ref().and_then(|node| match self.position { + // Get the root hash for the parent branch node + SubNodePosition::ParentBranch => node.root_hash, + // Or get it from the children + SubNodePosition::Child(nibble) => Some(node.hash_for_nibble(nibble)), + }) + } + + /// Returns the position to the current node. + #[inline] + pub const fn position(&self) -> SubNodePosition { + self.position } /// Increments the nibble index. #[inline] pub fn inc_nibble(&mut self) { - self.nibble += 1; - update_full_key(&mut self.full_key, self.nibble - 1, self.nibble); + let old_position = self.position; + self.position.increment(); + self.update_full_key(old_position); } /// Sets the nibble index. #[inline] - pub fn set_nibble(&mut self, nibble: i8) { - let old_nibble = self.nibble; - self.nibble = nibble; - update_full_key(&mut self.full_key, old_nibble, self.nibble); + pub fn set_nibble(&mut self, nibble: u8) { + let old_position = self.position; + self.position = SubNodePosition::Child(nibble); + self.update_full_key(old_position); } } -/// Constructs a full key from the given `Nibbles` and `nibble`. -#[inline] -fn full_key(mut key: Nibbles, nibble: i8) -> Nibbles { - if nibble >= 0 { - key.push(nibble as u8); - } - key +/// Represents a subnode position. +#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)] +pub enum SubNodePosition { + /// Positioned at the parent branch node. + ParentBranch, + /// Positioned at a child node at the given nibble. + Child(u8), } -/// Updates the key by replacing or appending a nibble based on the old and new nibble values. -#[inline] -fn update_full_key(key: &mut Nibbles, old_nibble: i8, new_nibble: i8) { - if new_nibble >= 0 { - if old_nibble >= 0 { - let last_index = key.len() - 1; - key.set_at(last_index, new_nibble as u8); - } else { - key.push(new_nibble as u8); +impl SubNodePosition { + /// Returns `true` if the position is set to the parent branch node. + pub fn is_parent(&self) -> bool { + matches!(self, Self::ParentBranch) + } + + /// Returns `true` if the position is set to a child node. + pub fn is_child(&self) -> bool { + matches!(self, Self::Child(_)) + } + + /// Returns the nibble of the child node if the position is set to a child node. + pub fn as_child(&self) -> Option { + match self { + Self::Child(nibble) => Some(*nibble), + _ => None, + } + } + + /// Returns `true` if the position is set to a last child nibble (i.e. greater than or equal to + /// 0xf). + pub fn is_last_child(&self) -> bool { + match self { + Self::ParentBranch => false, + Self::Child(nibble) => *nibble >= 0xf, + } + } + + fn increment(&mut self) { + match self { + Self::ParentBranch => *self = Self::Child(0), + Self::Child(nibble) => *nibble += 1, } - } else if old_nibble >= 0 { - key.pop(); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn subnode_position_ord() { + assert!([ + SubNodePosition::ParentBranch, + SubNodePosition::Child(0), + SubNodePosition::Child(1), + SubNodePosition::Child(2), + SubNodePosition::Child(3), + SubNodePosition::Child(4), + SubNodePosition::Child(5), + SubNodePosition::Child(6), + SubNodePosition::Child(7), + SubNodePosition::Child(8), + SubNodePosition::Child(9), + SubNodePosition::Child(10), + SubNodePosition::Child(11), + SubNodePosition::Child(12), + SubNodePosition::Child(13), + SubNodePosition::Child(14), + SubNodePosition::Child(15), + ] + .is_sorted()); + } + + #[test] + fn subnode_position_is_last_child() { + assert!([ + SubNodePosition::ParentBranch, + SubNodePosition::Child(0), + SubNodePosition::Child(1), + SubNodePosition::Child(2), + SubNodePosition::Child(3), + SubNodePosition::Child(4), + SubNodePosition::Child(5), + SubNodePosition::Child(6), + SubNodePosition::Child(7), + SubNodePosition::Child(8), + SubNodePosition::Child(9), + SubNodePosition::Child(10), + SubNodePosition::Child(11), + SubNodePosition::Child(12), + SubNodePosition::Child(13), + SubNodePosition::Child(14), + ] + .iter() + .all(|position| !position.is_last_child())); + assert!(SubNodePosition::Child(15).is_last_child()); + assert!(SubNodePosition::Child(16).is_last_child()); } } diff --git a/crates/trie/trie/src/walker.rs b/crates/trie/trie/src/walker.rs index f1159f1fe6..16d1d990ea 100644 --- a/crates/trie/trie/src/walker.rs +++ b/crates/trie/trie/src/walker.rs @@ -1,6 +1,6 @@ use crate::{ prefix_set::PrefixSet, - trie_cursor::{CursorSubNode, TrieCursor}, + trie_cursor::{subnode::SubNodePosition, CursorSubNode, TrieCursor}, BranchNodeCompact, Nibbles, }; use alloy_primitives::{map::HashSet, B256}; @@ -162,14 +162,14 @@ impl TrieWalker { if !self.can_skip_current_node && self.children_are_in_trie() { trace!( target: "trie::walker", - nibble = ?last.nibble(), + position = ?last.position(), "cannot skip current node and children are in the trie" ); // If we can't skip the current node and the children are in the trie, // either consume the next node or move to the next sibling. - match last.nibble() { - -1 => self.move_to_next_sibling(true)?, - _ => self.consume_node()?, + match last.position() { + SubNodePosition::ParentBranch => self.move_to_next_sibling(true)?, + SubNodePosition::Child(_) => self.consume_node()?, } } else { trace!(target: "trie::walker", "can skip current node"); @@ -209,7 +209,7 @@ impl TrieWalker { // We need to sync the stack with the trie structure when consuming a new node. This is // necessary for proper traversal and accurately representing the trie in the stack. if !key.is_empty() && !self.stack.is_empty() { - self.stack[0].set_nibble(key[0] as i8); + self.stack[0].set_nibble(key[0]); } // The current tree mask might have been set incorrectly. @@ -227,13 +227,13 @@ impl TrieWalker { // Create a new CursorSubNode and push it to the stack. let subnode = CursorSubNode::new(key, Some(node)); - let nibble = subnode.nibble(); + let position = subnode.position(); self.stack.push(subnode); self.update_skip_node(); // Delete the current node if it's included in the prefix set or it doesn't contain the root // hash. - if !self.can_skip_current_node || nibble != -1 { + if !self.can_skip_current_node || position.is_child() { if let Some((keys, key)) = self.removed_keys.as_mut().zip(self.cursor.current()?) { keys.insert(key); } @@ -252,7 +252,9 @@ impl TrieWalker { // Check if the walker needs to backtrack to the previous level in the trie during its // traversal. - if subnode.nibble() >= 0xf || (subnode.nibble() < 0 && !allow_root_to_child_nibble) { + if subnode.position().is_last_child() || + (subnode.position().is_parent() && !allow_root_to_child_nibble) + { self.stack.pop(); self.move_to_next_sibling(false)?; return Ok(()) @@ -266,13 +268,13 @@ impl TrieWalker { // Find the next sibling with state. loop { - let nibble = subnode.nibble(); + let position = subnode.position(); if subnode.state_flag() { - trace!(target: "trie::walker", nibble, "found next sibling with state"); + trace!(target: "trie::walker", ?position, "found next sibling with state"); return Ok(()) } - if nibble == 0xf { - trace!(target: "trie::walker", nibble, "checked all siblings"); + if position.is_last_child() { + trace!(target: "trie::walker", ?position, "checked all siblings"); break } subnode.inc_nibble();