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

firewood/src/proof.rs

@@ -108,17 +108,32 @@ impl std::fmt::Debug for ProofNode {
 }
 
 impl Hashable for ProofNode {
-    fn parent_prefix_path(&self) -> impl IntoSplitPath + '_ {
+    type LeadingPath<'a>
+        = &'a [PathComponent]
+    where
+        Self: 'a;
+
+    type PartialPath<'a>
+        = &'a [PathComponent]
+    where
+        Self: 'a;
+
+    type FullPath<'a>
+        = &'a [PathComponent]
+    where
+        Self: 'a;
+
+    fn parent_prefix_path(&self) -> Self::LeadingPath<'_> {
         let (prefix, _) = self.key.split_at(self.partial_len);
         prefix
     }
 
-    fn partial_path(&self) -> impl IntoSplitPath + '_ {
+    fn partial_path(&self) -> Self::PartialPath<'_> {
         let (_, suffix) = self.key.split_at(self.partial_len);
         suffix
     }
 
-    fn full_path(&self) -> impl IntoSplitPath + '_ {
+    fn full_path(&self) -> Self::FullPath<'_> {
         &self.key
     }
 

storage/src/hashednode.rs

@@ -3,7 +3,6 @@
 
 use crate::{
     Children, HashType, HashableShunt, IntoSplitPath, Node, Path, PathComponent, SplitPath,
-    TriePath,
 };
 use smallvec::SmallVec;
 
@@ -153,22 +152,32 @@ impl<T: AsRef<[u8]>> AsRef<[u8]> for ValueDigest<T> {
 
 /// A node in the trie that can be hashed.
 pub trait Hashable: std::fmt::Debug {
+    /// The type of the leading path.
+    type LeadingPath<'a>: IntoSplitPath + 'a
+    where
+        Self: 'a;
+
+    /// The type of the partial path.
+    type PartialPath<'a>: IntoSplitPath + 'a
+    where
+        Self: 'a;
+
+    /// The type of the full path.
+    type FullPath<'a>: IntoSplitPath + 'a
+    where
+        Self: 'a;
+
     /// The full path of this node's parent where each byte is a nibble.
-    fn parent_prefix_path(&self) -> impl IntoSplitPath + '_;
+    fn parent_prefix_path(&self) -> Self::LeadingPath<'_>;
     /// The partial path of this node where each byte is a nibble.
-    fn partial_path(&self) -> impl IntoSplitPath + '_;
+    fn partial_path(&self) -> Self::PartialPath<'_>;
+    /// The full path of this node including the parent's prefix where each byte is a nibble.
+    fn full_path(&self) -> Self::FullPath<'_>;
     /// The node's value or hash.
     fn value_digest(&self) -> Option<ValueDigest<&[u8]>>;
     /// Each element is a child's index and hash.
     /// Yields 0 elements if the node is a leaf.
     fn children(&self) -> Children<Option<HashType>>;
-
-    /// The full path of this node including the parent's prefix where each byte is a nibble.
-    fn full_path(&self) -> impl IntoSplitPath + '_ {
-        self.parent_prefix_path()
-            .into_split_path()
-            .append(self.partial_path().into_split_path())
-    }
 }
 
 /// A preimage of a hash.

storage/src/hashedshunt.rs

@@ -1,7 +1,7 @@
 // Copyright (C) 2025, Ava Labs, Inc. All rights reserved.
 // See the file LICENSE.md for licensing terms.
 
-use crate::{Children, HashType, Hashable, IntoSplitPath, SplitPath, ValueDigest};
+use crate::{Children, HashType, Hashable, JoinedPath, SplitPath, ValueDigest};
 
 /// A shunt for a hasheable trie that we can use to compute the hash of a node
 /// using component parts.
@@ -51,14 +51,33 @@ impl<P1: SplitPath, P2: SplitPath> std::fmt::Debug for HashableShunt<'_, P1, P2>
 }
 
 impl<P1: SplitPath, P2: SplitPath> Hashable for HashableShunt<'_, P1, P2> {
-    fn parent_prefix_path(&self) -> impl IntoSplitPath + '_ {
+    type LeadingPath<'a>
+        = P1
+    where
+        Self: 'a;
+
+    type PartialPath<'a>
+        = P2
+    where
+        Self: 'a;
+
+    type FullPath<'a>
+        = JoinedPath<P1, P2>
+    where
+        Self: 'a;
+
+    fn parent_prefix_path(&self) -> Self::LeadingPath<'_> {
         self.parent_prefix
     }
 
-    fn partial_path(&self) -> impl IntoSplitPath + '_ {
+    fn partial_path(&self) -> Self::PartialPath<'_> {
         self.partial_path
     }
 
+    fn full_path(&self) -> Self::FullPath<'_> {
+        self.parent_prefix_path().append(self.partial_path())
+    }
+
     fn value_digest(&self) -> Option<ValueDigest<&[u8]>> {
         self.value.clone()
     }

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 {