feat: add TriePathIter for iterating over the nodes along a path in a trie

Author: demosdemon

storage/src/lib.rs

@@ -66,7 +66,7 @@ pub use path::{
 pub use path::{PackedBytes, PackedPathComponents, PackedPathRef};
 pub use tries::{
     DuplicateKeyError, HashedKeyValueTrieRoot, HashedTrieNode, IterAscending, IterDescending,
-    KeyValueTrieRoot, TrieEdgeIter, TrieEdgeState, TrieNode, TrieValueIter,
+    KeyValueTrieRoot, TrieEdgeIter, TrieEdgeState, TrieNode, TriePathIter, TrieValueIter,
 };
 pub use u4::{TryFromIntError, U4};
 

storage/src/tries/iter.rs

@@ -1,7 +1,9 @@
 // Copyright (C) 2025, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE.md for licensing terms.
 
-use crate::{HashType, PathBuf, PathComponent, TrieNode, TriePath, tries::TrieEdgeState};
+use crate::{
+    HashType, IntoSplitPath, PathBuf, PathComponent, SplitPath, TrieEdgeState, TrieNode, TriePath,
+};
 
 /// A marker type for [`TrieEdgeIter`] that indicates that the iterator traverses
 /// the trie in ascending order.
@@ -33,6 +35,18 @@ pub struct TrieValueIter<'a, N: ?Sized, V: ?Sized, D> {
     edges: TrieEdgeIter<'a, N, V, D>,
 }
 
+/// An iterator over the edges along a specified path in a key-value trie
+/// terminating at the node corresponding to the path, if it exists; otherwise,
+/// terminating at the deepest existing edge along the path.
+#[derive(Debug)]
+#[must_use = "iterators are lazy and do nothing unless consumed"]
+pub struct TriePathIter<'a, P, N: ?Sized, V: ?Sized> {
+    needle: P,
+    current_path: PathBuf,
+    current_edge: Option<TrieEdgeState<'a, N>>,
+    marker: std::marker::PhantomData<V>,
+}
+
 #[derive(Debug)]
 struct Frame<'a, N: ?Sized, V: ?Sized> {
     node: &'a N,
@@ -83,6 +97,26 @@ where
     }
 }
 
+impl<'a, P, N, V> TriePathIter<'a, P, N, V>
+where
+    P: SplitPath,
+    N: TrieNode<V> + ?Sized,
+    V: AsRef<[u8]> + ?Sized,
+{
+    /// Creates a new iterator over the edges along the given path in the
+    /// specified key-value trie.
+    pub const fn new(root: &'a N, root_hash: Option<&'a HashType>, path: P) -> Self {
+        let mut this = Self {
+            needle: path,
+            current_path: PathBuf::new_const(),
+            current_edge: None,
+            marker: std::marker::PhantomData,
+        };
+        this.current_edge = Some(TrieEdgeState::from_node(root, root_hash));
+        this
+    }
+}
+
 /// Both iterators share this logic to descend into a node's children.
 ///
 /// The passed in `children_iter` should be an iterator over the indices into
@@ -218,6 +252,45 @@ where
     }
 }
 
+impl<'a, P, N, V> Iterator for TriePathIter<'a, P, N, V>
+where
+    P: SplitPath,
+    N: TrieNode<V> + ?Sized,
+    V: AsRef<[u8]> + ?Sized,
+{
+    type Item = (PathBuf, TrieEdgeState<'a, N>);
+
+    fn next(&mut self) -> Option<Self::Item> {
+        // qualified path to `Option::take` because rust-analyzer thinks
+        // `self.current_edge` is an `Iterator` and displays an error for
+        // `self.current_edge.take()` where `rustc` does not
+        let edge = Option::take(&mut self.current_edge)?;
+        let path = self.current_path.clone();
+
+        let Some(node) = edge.node() else {
+            // the current edge is remote, so we cannot descend further
+            return Some((path, edge));
+        };
+
+        self.current_path.extend(node.partial_path().components());
+
+        if let (None, Some((pc, needle_suffix)), _) = node
+            .partial_path()
+            .into_split_path()
+            .longest_common_prefix(self.needle)
+            .split_first_parts()
+        {
+            // the target path continues beyond the current node, descend
+            // into it if there's a child along the path
+            self.current_path.push(pc);
+            self.current_edge = node.child_state(pc);
+            self.needle = needle_suffix;
+        }
+
+        Some((path, edge))
+    }
+}
+
 // auto-derived implementations would require N: Clone, V: Clone which is too much
 
 impl<N: ?Sized, V: ?Sized, D> Clone for TrieEdgeIter<'_, N, V, D> {
@@ -247,6 +320,23 @@ impl<N: ?Sized, V: ?Sized, D> Clone for TrieValueIter<'_, N, V, D> {
     }
 }
 
+impl<P: SplitPath, N: ?Sized, V: ?Sized> Clone for TriePathIter<'_, P, N, V> {
+    fn clone(&self) -> Self {
+        Self {
+            needle: self.needle,
+            current_path: self.current_path.clone(),
+            current_edge: self.current_edge,
+            marker: std::marker::PhantomData,
+        }
+    }
+
+    fn clone_from(&mut self, source: &Self) {
+        self.needle = source.needle;
+        self.current_path.clone_from(&source.current_path);
+        self.current_edge = source.current_edge;
+    }
+}
+
 impl<N: ?Sized, V: ?Sized> Clone for Frame<'_, N, V> {
     fn clone(&self) -> Self {
         Self {