Adding support for walking a prefix

Author: armon

radix.go

@@ -374,6 +374,38 @@ func (t *Tree) Walk(fn WalkFn) {
 	recursiveWalk(t.root, fn)
 }
 
+// WalkPrefix is used to walk the tree under a prefix
+func (t *Tree) WalkPrefix(prefix string, fn WalkFn) {
+	n := t.root
+	search := prefix
+	for {
+		// Check for key exhaution
+		if len(search) == 0 {
+			recursiveWalk(n, fn)
+			return
+		}
+
+		// Look for an edge
+		n = n.getEdge(search[0])
+		if n == nil {
+			break
+		}
+
+		// Consume the search prefix
+		if strings.HasPrefix(search, n.prefix) {
+			search = search[len(n.prefix):]
+
+		} else if strings.HasPrefix(n.prefix, search) {
+			// Child may be under our search prefix
+			recursiveWalk(n, fn)
+			return
+		} else {
+			break
+		}
+	}
+
+}
+
 // recursiveWalk is used to do a pre-order walk of a node
 // recursively. Returns true if the walk should be aborted
 func recursiveWalk(n *node, fn WalkFn) bool {

radix_test.go

@@ -3,6 +3,8 @@ package radix
 import (
 	crand "crypto/rand"
 	"fmt"
+	"reflect"
+	"sort"
 	"testing"
 )
 
@@ -112,6 +114,85 @@ func TestLongestPrefix(t *testing.T) {
 	}
 }
 
+func TestWalkPrefix(t *testing.T) {
+	r := New()
+
+	keys := []string{
+		"foobar",
+		"foo/bar/baz",
+		"foo/baz/bar",
+		"foo/zip/zap",
+		"zipzap",
+	}
+	for _, k := range keys {
+		r.Insert(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{
+		exp{
+			"f",
+			[]string{"foobar", "foo/bar/baz", "foo/baz/bar", "foo/zip/zap"},
+		},
+		exp{
+			"foo",
+			[]string{"foobar", "foo/bar/baz", "foo/baz/bar", "foo/zip/zap"},
+		},
+		exp{
+			"foob",
+			[]string{"foobar"},
+		},
+		exp{
+			"foo/",
+			[]string{"foo/bar/baz", "foo/baz/bar", "foo/zip/zap"},
+		},
+		exp{
+			"foo/b",
+			[]string{"foo/bar/baz", "foo/baz/bar"},
+		},
+		exp{
+			"foo/ba",
+			[]string{"foo/bar/baz", "foo/baz/bar"},
+		},
+		exp{
+			"foo/bar",
+			[]string{"foo/bar/baz"},
+		},
+		exp{
+			"foo/bar/baz",
+			[]string{"foo/bar/baz"},
+		},
+		exp{
+			"foo/bar/bazoo",
+			[]string{},
+		},
+		exp{
+			"z",
+			[]string{"zipzap"},
+		},
+	}
+
+	for _, test := range cases {
+		out := []string{}
+		fn := func(s string, v interface{}) bool {
+			out = append(out, s)
+			return false
+		}
+		r.WalkPrefix(test.inp, fn)
+		sort.Strings(out)
+		sort.Strings(test.out)
+		if !reflect.DeepEqual(out, test.out) {
+			t.Fatalf("mis-match: %v %v", out, test.out)
+		}
+	}
+}
+
 // generateUUID is used to generate a random UUID
 func generateUUID() string {
 	buf := make([]byte, 16)