use crate::{BranchNodeCompact, Nibbles, StoredSubNode, CHILD_INDEX_RANGE}; use reth_primitives::B256; /// Cursor for iterating over a subtrie. #[derive(Clone)] pub struct CursorSubNode { /// The key of the current node. pub key: Nibbles, /// The index of the next child to visit. nibble: i8, /// The node itself. pub node: Option, /// Full key full_key: Nibbles, } impl Default for CursorSubNode { fn default() -> Self { Self::new(Nibbles::default(), None) } } 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("state_flag", &self.state_flag()) .field("tree_flag", &self.tree_flag()) .field("hash_flag", &self.hash_flag()) .field("hash", &self.hash()) .finish() } } /// Implements conversion from `StoredSubNode` to `CursorSubNode`. impl From for CursorSubNode { /// Converts a `StoredSubNode` into a `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 key = Nibbles::from_nibbles_unchecked(value.key); let full_key = full_key(key.clone(), nibble); Self { key, nibble, node: value.node, full_key } } } impl From for StoredSubNode { fn from(value: CursorSubNode) -> Self { let nibble = if value.nibble >= 0 { Some(value.nibble as u8) } else { None }; Self { key: value.key.to_vec(), nibble, node: value.node } } } impl CursorSubNode { /// 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 } } /// Returns the full key of the current node. #[inline] pub const fn full_key(&self) -> &Nibbles { &self.full_key } /// Returns `true` if the state flag is set for the current nibble. #[inline] pub fn state_flag(&self) -> bool { self.node .as_ref() .map_or(true, |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() .map_or(true, |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().map_or(false, |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)), } } else { None } } /// Returns the next child index to visit. #[inline] pub const fn nibble(&self) -> i8 { self.nibble } /// 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); } /// 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); } } /// 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 } /// 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); } } else if old_nibble >= 0 { key.pop(); } }