Add Path Iterator to return results like WalkPath would visit

node.go

@@ -221,6 +221,12 @@ func (n *Node[T]) ReverseIterator() *ReverseIterator[T] {
 	return NewReverseIterator(n)
 }
 
+// Iterator is used to return an iterator at
+// the given node to walk the tree
+func (n *Node[T]) PathIterator(path []byte) *PathIterator[T] {
+	return &PathIterator[T]{node: n, path: path}
+}
+
 // rawIterator is used to return a raw iterator at the given node to walk the
 // tree.
 func (n *Node[T]) rawIterator() *rawIterator[T] {
@@ -274,30 +280,12 @@ func (n *Node[T]) WalkPrefix(prefix []byte, fn WalkFn[T]) {
 // all the entries *under* the given prefix, this walks the
 // entries *above* the given prefix.
 func (n *Node[T]) WalkPath(path []byte, fn WalkFn[T]) {
-	search := path
-	for {
-		// Visit the leaf values if any
-		if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
-			return
-		}
-
-		// Check for key exhaustion
-		if len(search) == 0 {
-			return
-		}
+	i := n.PathIterator(path)
 
-		// Look for an edge
-		_, n = n.getEdge(search[0])
-		if n == nil {
+	for path, val, ok := i.Next(); ok; path, val, ok = i.Next() {
+		if fn(path, val) {
 			return
 		}
-
-		// Consume the search prefix
-		if bytes.HasPrefix(search, n.prefix) {
-			search = search[len(n.prefix):]
-		} else {
-			break
-		}
 	}
 }
 

path_iter.go

@@ -0,0 +1,59 @@
+package iradix
+
+import "bytes"
+
+// PathIterator is used to iterate over a set of nodes from the root
+// down to a specified path. This will iterate overthe same values that
+// the Node.WalkPath method will.
+type PathIterator[T any] struct {
+	node *Node[T]
+	path []byte
+	done bool
+}
+
+// Next returns the next node in order
+func (i *PathIterator[T]) Next() ([]byte, T, bool) {
+	// This is mostly just an asyncrhonous implementation of the WalkPath
+	// method on the node.
+	var zero T
+	var leaf *leafNode[T]
+
+	for leaf == nil && i.node != nil {
+		// visit the leaf values if any
+		if i.node.leaf != nil {
+			leaf = i.node.leaf
+		}
+
+		i.iterate()
+	}
+
+	if leaf != nil {
+		return leaf.key, leaf.val, true
+	}
+
+	return nil, zero, false
+}
+
+func (i *PathIterator[T]) iterate() {
+	// Check for key exhaustion
+	if len(i.path) == 0 {
+		i.node = nil
+		return
+	}
+
+	// Look for an edge
+	_, i.node = i.node.getEdge(i.path[0])
+	if i.node == nil {
+		return
+	}
+
+	// Consume the search prefix
+	if bytes.HasPrefix(i.path, i.node.prefix) {
+		i.path = i.path[len(i.node.prefix):]
+	} else {
+		// there are no more nodes to iterate through so
+		// nil out the node to prevent returning results
+		// for subsequent calls to Next()
+		i.node = nil
+	}
+}

path_iter_test.go

@@ -0,0 +1,104 @@
+package iradix
+
+import (
+	"reflect"
+	"sort"
+	"testing"
+)
+
+func TestPathIterator(t *testing.T) {
+	r := New[any]()
+
+	keys := []string{
+		"foo",
+		"foo/bar",
+		"foo/bar/baz",
+		"foo/baz/bar",
+		"foo/zip/zap",
+		"zipzap",
+	}
+	for _, k := range keys {
+		r, _, _ = r.Insert([]byte(k), nil)
+	}
+	if r.Len() != len(keys) {
+		t.Fatalf("bad len: %v %v", r.Len(), len(keys))
+	}
+
+	type exp struct {
+		inp string
+		out []string
+	}
+	cases := []exp{
+		{
+			"f",
+			[]string{},
+		},
+		{
+			"foo",
+			[]string{"foo"},
+		},
+		{
+			"foo/",
+			[]string{"foo"},
+		},
+		{
+			"foo/ba",
+			[]string{"foo"},
+		},
+		{
+			"foo/bar",
+			[]string{"foo", "foo/bar"},
+		},
+		{
+			"foo/bar/baz",
+			[]string{"foo", "foo/bar", "foo/bar/baz"},
+		},
+		{
+			"foo/bar/bazoo",
+			[]string{"foo", "foo/bar", "foo/bar/baz"},
+		},
+		{
+			"z",
+			[]string{},
+		},
+	}
+
+	root := r.Root()
+
+	for _, test := range cases {
+		iter := root.PathIterator([]byte(test.inp))
+
+		// Radix tree iteration on our string indices will return values with keys in
+		// ascending order. So before we check the iteration ordering we must sort
+		// the expected outputs.
+		sort.Strings(test.out)
+
+		// verify that all the expected values come out in ascending order.
+		for idx, expected := range test.out {
+			actual, _, found := iter.Next()
+			// ensure we found a value (i.e. iteration is finding the correct number of values)
+			if !found {
+				t.Fatalf("iteration returned fewer values than expected: %d, actual: %d", len(test.out), idx)
+			}
+
+			// ensure that the values are coming out in the sorted order.
+			if !reflect.DeepEqual([]byte(expected), actual) {
+				t.Errorf("expected: %#v", []byte(expected))
+				t.Errorf("actual:   %#v", actual)
+				t.Fatalf("value returned during iteration doesn't match our expectation (iteration num: %d)", idx)
+			}
+		}
+
+		// Now ensure there are no trailing values that the tree will output.
+		_, _, found := iter.Next()
+		if found {
+			t.Fatalf("iteration returned more values than expected: %d, actual: %d+", len(test.out), len(test.out)+1)
+		}
+
+		// Verify that continued calls to next on a completed iterator do not panic or return values.
+		_, _, found = iter.Next()
+		if found {
+			t.Fatalf("iteration returned a value after previously indicating iteration was complete")
+		}
+	}
+}