doc: improve documentation for trie crate

crates/trie/src/node_iter.rs

@@ -6,28 +6,39 @@ use crate::{
 };
 use reth_primitives::{trie::Nibbles, Account, StorageEntry, B256, U256};
 
+/// Represents a branch node in the trie.
 #[derive(Debug)]
 pub(crate) struct TrieBranchNode {
+    /// The key associated with the node.
     pub(crate) key: Nibbles,
+    /// The value associated with the node.
     pub(crate) value: B256,
+    /// Indicates whether children are in the trie.
     pub(crate) children_are_in_trie: bool,
 }
 
 impl TrieBranchNode {
+    /// Creates a new `TrieBranchNode`.
     pub(crate) fn new(key: Nibbles, value: B256, children_are_in_trie: bool) -> Self {
         Self { key, value, children_are_in_trie }
     }
 }
 
+/// Represents a variant of an account node.
 #[derive(Debug)]
 pub(crate) enum AccountNode {
+    /// Branch node.
     Branch(TrieBranchNode),
+    /// Leaf node.
     Leaf(B256, Account),
 }
 
+/// Represents a variant of a storage node.
 #[derive(Debug)]
 pub(crate) enum StorageNode {
+    /// Branch node.
     Branch(TrieBranchNode),
+    /// Leaf node.
     Leaf(B256, U256),
 }
 
@@ -49,6 +60,7 @@ pub(crate) struct AccountNodeIter<C, H> {
 }
 
 impl<C, H> AccountNodeIter<C, H> {
+    /// Creates a new `AccountNodeIter`.
     pub(crate) fn new(walker: TrieWalker<C>, hashed_account_cursor: H) -> Self {
         Self {
             walker,
@@ -59,6 +71,8 @@ impl<C, H> AccountNodeIter<C, H> {
         }
     }
 
+    /// Sets the last iterated account key and returns the modified `AccountNodeIter`.
+    /// This is used to resume iteration from the last checkpoint.
     pub(crate) fn with_last_account_key(mut self, previous_account_key: B256) -> Self {
         self.previous_account_key = Some(previous_account_key);
         self
@@ -83,9 +97,12 @@ where
     /// NOTE: The iteration will start from the key of the previous hashed entry if it was supplied.
     pub(crate) fn try_next(&mut self) -> Result<Option<AccountNode>, StateRootError> {
         loop {
+            // If the walker has a key...
             if let Some(key) = self.walker.key() {
+                // Check if the current walker key is unchecked and there's no previous account key
                 if !self.current_walker_key_checked && self.previous_account_key.is_none() {
                     self.current_walker_key_checked = true;
+                    // If it's possible to skip the current node in the walker, return a branch node
                     if self.walker.can_skip_current_node {
                         return Ok(Some(AccountNode::Branch(TrieBranchNode::new(
                             key.clone(),
@@ -96,22 +113,30 @@ where
                 }
             }
 
+            // If there's a hashed address and account...
             if let Some((hashed_address, account)) = self.current_hashed_entry.take() {
+                // If the walker's key is less than the unpacked hashed address, reset the checked
+                // status and continue
                 if self.walker.key().map_or(false, |key| key < &Nibbles::unpack(hashed_address)) {
                     self.current_walker_key_checked = false;
                     continue
                 }
 
+                // Set the next hashed entry as a leaf node and return
                 self.current_hashed_entry = self.hashed_account_cursor.next()?;
                 return Ok(Some(AccountNode::Leaf(hashed_address, account)))
             }
 
+            // Handle seeking and advancing based on the previous account key
             match self.previous_account_key.take() {
                 Some(account_key) => {
+                    // Seek to the previous account key and get the next hashed entry
                     self.hashed_account_cursor.seek(account_key)?;
                     self.current_hashed_entry = self.hashed_account_cursor.next()?;
                 }
                 None => {
+                    // Get the seek key and set the current hashed entry based on walker's next
+                    // unprocessed key
                     let seek_key = match self.walker.next_unprocessed_key() {
                         Some(key) => key,
                         None => break, // no more keys
@@ -142,6 +167,7 @@ pub(crate) struct StorageNodeIter<C, H> {
 }
 
 impl<C, H> StorageNodeIter<C, H> {
+    /// Creates a new instance of StorageNodeIter.
     pub(crate) fn new(
         walker: TrieWalker<C>,
         hashed_storage_cursor: H,
@@ -173,10 +199,14 @@ where
     /// 5. Repeat.
     pub(crate) fn try_next(&mut self) -> Result<Option<StorageNode>, StorageRootError> {
         loop {
+            // Check if there's a key in the walker.
             if let Some(key) = self.walker.key() {
+                // Check if the walker key hasn't been checked yet.
                 if !self.current_walker_key_checked {
                     self.current_walker_key_checked = true;
+                    // Check if the current node can be skipped in the walker.
                     if self.walker.can_skip_current_node {
+                        // Return a branch node based on the walker's properties.
                         return Ok(Some(StorageNode::Branch(TrieBranchNode::new(
                             key.clone(),
                             self.walker.hash().unwrap(),
@@ -186,23 +216,32 @@ where
                 }
             }
 
+            // Check for a current hashed storage entry.
             if let Some(StorageEntry { key: hashed_key, value }) = self.current_hashed_entry.take()
             {
+                // Compare keys and proceed accordingly.
                 if self.walker.key().map_or(false, |key| key < &Nibbles::unpack(hashed_key)) {
                     self.current_walker_key_checked = false;
                     continue
                 }
 
+                // Move to the next hashed storage entry and return the corresponding leaf node.
                 self.current_hashed_entry = self.hashed_storage_cursor.next()?;
                 return Ok(Some(StorageNode::Leaf(hashed_key, value)))
             }
 
-            let Some(seek_key) = self.walker.next_unprocessed_key() else { break };
-            self.current_hashed_entry =
-                self.hashed_storage_cursor.seek(self.hashed_address, seek_key)?;
-            self.walker.advance()?;
+            // Attempt to get the next unprocessed key from the walker.
+            if let Some(seek_key) = self.walker.next_unprocessed_key() {
+                // Seek and update the current hashed entry based on the new seek key.
+                self.current_hashed_entry =
+                    self.hashed_storage_cursor.seek(self.hashed_address, seek_key)?;
+                self.walker.advance()?;
+            } else {
+                // No more keys to process, break the loop.
+                break
+            }
         }
 
-        Ok(None)
+        Ok(None) // Return None if no more nodes are available.
     }
 }

crates/trie/src/trie_cursor/database_cursors.rs

@@ -45,22 +45,26 @@ impl<C> TrieCursor for DatabaseAccountTrieCursor<C>
 where
     C: DbCursorRO<tables::AccountsTrie>,
 {
+    /// The type of key used by this cursor.
     type Key = StoredNibbles;
 
+    /// Seeks an exact match for the provided key in the account trie.
     fn seek_exact(
         &mut self,
         key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(self.0.seek_exact(key)?.map(|value| (value.0 .0.to_vec(), value.1 .0)))
     }
 
+    /// Seeks a key in the account trie that matches or is greater than the provided key.
     fn seek(
         &mut self,
         key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(self.0.seek(key)?.map(|value| (value.0 .0.to_vec(), value.1 .0)))
     }
 
+    /// Retrieves the current key in the cursor.
     fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
         Ok(self.0.current()?.map(|(k, _)| TrieKey::AccountNode(k)))
     }
@@ -71,6 +75,7 @@ where
 pub struct DatabaseStorageTrieCursor<C> {
     /// The underlying cursor.
     pub cursor: C,
+    /// Hashed address used for cursor positioning.
     hashed_address: B256,
 }
 
@@ -85,8 +90,10 @@ impl<C> TrieCursor for DatabaseStorageTrieCursor<C>
 where
     C: DbDupCursorRO<tables::StoragesTrie> + DbCursorRO<tables::StoragesTrie>,
 {
+    /// Defines the type for keys used in the storage trie cursor.
     type Key = StoredNibblesSubKey;
 
+    /// Seeks an exact match for the given key in the storage trie.
     fn seek_exact(
         &mut self,
         key: Self::Key,
@@ -98,6 +105,7 @@ where
             .map(|value| (value.nibbles.to_vec(), value.node)))
     }
 
+    /// Seeks the given key in the storage trie.
     fn seek(
         &mut self,
         key: Self::Key,
@@ -108,6 +116,7 @@ where
             .map(|value| (value.nibbles.to_vec(), value.node)))
     }
 
+    /// Retrieves the current value in the storage trie cursor.
     fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
         Ok(self.cursor.current()?.map(|(k, v)| TrieKey::StorageNode(k, v.nibbles)))
     }

crates/trie/src/trie_cursor/noop.rs

@@ -1,22 +1,22 @@
+use super::{TrieCursor, TrieCursorFactory};
+use crate::updates::TrieKey;
 use reth_db::DatabaseError;
 use reth_primitives::trie::{BranchNodeCompact, StoredNibbles, StoredNibblesSubKey};
 
-use crate::updates::TrieKey;
-
-use super::{TrieCursor, TrieCursorFactory};
-
 /// Noop trie cursor factory.
 #[derive(Default, Debug)]
 #[non_exhaustive]
 pub struct NoopTrieCursorFactory;
 
 impl TrieCursorFactory for NoopTrieCursorFactory {
+    /// Generates a Noop account trie cursor.
     fn account_trie_cursor(
         &self,
     ) -> Result<Box<dyn TrieCursor<Key = StoredNibbles> + '_>, DatabaseError> {
         Ok(Box::<NoopAccountTrieCursor>::default())
     }
 
+    /// Generates a Noop storage trie cursor.
     fn storage_tries_cursor(
         &self,
         _hashed_address: reth_primitives::B256,
@@ -33,47 +33,53 @@ pub struct NoopAccountTrieCursor;
 impl TrieCursor for NoopAccountTrieCursor {
     type Key = StoredNibbles;
 
+    /// Seeks within the account trie.
     fn seek(
         &mut self,
         _key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(None)
     }
 
+    /// Seeks an exact match within the account trie.
     fn seek_exact(
         &mut self,
         _key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(None)
     }
 
+    /// Retrieves the current cursor position within the account trie.
     fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
         Ok(None)
     }
 }
 
-/// Noop account trie cursor.
+/// Noop storage trie cursor.
 #[derive(Default, Debug)]
 #[non_exhaustive]
 pub struct NoopStorageTrieCursor;
 
 impl TrieCursor for NoopStorageTrieCursor {
     type Key = StoredNibblesSubKey;
 
+    /// Seeks a key in storage tries.
     fn seek(
         &mut self,
         _key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(None)
     }
 
+    /// Seeks an exact match in storage tries.
     fn seek_exact(
         &mut self,
         _key: Self::Key,
     ) -> Result<Option<(Vec<u8>, BranchNodeCompact)>, DatabaseError> {
         Ok(None)
     }
 
+    /// Retrieves the current cursor position within storage tries.
     fn current(&mut self) -> Result<Option<TrieKey>, DatabaseError> {
         Ok(None)
     }

crates/trie/src/trie_cursor/subnode.rs

@@ -35,7 +35,12 @@ impl std::fmt::Debug for CursorSubNode {
     }
 }
 
+/// Implements conversion from `StoredSubNode` to `CursorSubNode`.
 impl From<StoredSubNode> 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 = match value.nibble {
             Some(n) => n as i8,
@@ -141,6 +146,7 @@ impl CursorSubNode {
     }
 }
 
+/// Constructs a full key from the given `Nibbles` and `nibble`.
 #[inline]
 fn full_key(mut key: Nibbles, nibble: i8) -> Nibbles {
     if nibble >= 0 {
@@ -149,6 +155,7 @@ fn full_key(mut key: Nibbles, nibble: i8) -> Nibbles {
     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 {

crates/trie/src/walker.rs

@@ -233,6 +233,7 @@ impl<C: TrieCursor> TrieWalker<C> {
             })
     }
 
+    /// Updates the skip node flag based on the walker's current state.
     fn update_skip_node(&mut self) {
         self.can_skip_current_node = if let Some(node) = self.stack.last() {
             !self.changes.contains(node.full_key()) && node.hash_flag()