Skip to content

Commit b045fc0

Browse files
committed
Adding support for walking a prefix
1 parent 519c5dd commit b045fc0

File tree

2 files changed

+113
-0
lines changed

2 files changed

+113
-0
lines changed

‎radix.go‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -374,6 +374,38 @@ func (t *Tree) Walk(fn WalkFn) {
374374
recursiveWalk(t.root, fn)
375375
}
376376

377+
// WalkPrefix is used to walk the tree under a prefix
378+
func (t *Tree) WalkPrefix(prefix string, fn WalkFn) {
379+
n := t.root
380+
search := prefix
381+
for {
382+
// Check for key exhaution
383+
if len(search) == 0 {
384+
recursiveWalk(n, fn)
385+
return
386+
}
387+
388+
// Look for an edge
389+
n = n.getEdge(search[0])
390+
if n == nil {
391+
break
392+
}
393+
394+
// Consume the search prefix
395+
if strings.HasPrefix(search, n.prefix) {
396+
search = search[len(n.prefix):]
397+
398+
} else if strings.HasPrefix(n.prefix, search) {
399+
// Child may be under our search prefix
400+
recursiveWalk(n, fn)
401+
return
402+
} else {
403+
break
404+
}
405+
}
406+
407+
}
408+
377409
// recursiveWalk is used to do a pre-order walk of a node
378410
// recursively. Returns true if the walk should be aborted
379411
func recursiveWalk(n *node, fn WalkFn) bool {

‎radix_test.go‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@ package radix
33
import (
44
crand "crypto/rand"
55
"fmt"
6+
"reflect"
7+
"sort"
68
"testing"
79
)
810

@@ -112,6 +114,85 @@ func TestLongestPrefix(t *testing.T) {
112114
}
113115
}
114116

117+
func TestWalkPrefix(t *testing.T) {
118+
r := New()
119+
120+
keys := []string{
121+
"foobar",
122+
"foo/bar/baz",
123+
"foo/baz/bar",
124+
"foo/zip/zap",
125+
"zipzap",
126+
}
127+
for _, k := range keys {
128+
r.Insert(k, nil)
129+
}
130+
if r.Len() != len(keys) {
131+
t.Fatalf("bad len: %v %v", r.Len(), len(keys))
132+
}
133+
134+
type exp struct {
135+
inp string
136+
out []string
137+
}
138+
cases := []exp{
139+
exp{
140+
"f",
141+
[]string{"foobar", "foo/bar/baz", "foo/baz/bar", "foo/zip/zap"},
142+
},
143+
exp{
144+
"foo",
145+
[]string{"foobar", "foo/bar/baz", "foo/baz/bar", "foo/zip/zap"},
146+
},
147+
exp{
148+
"foob",
149+
[]string{"foobar"},
150+
},
151+
exp{
152+
"foo/",
153+
[]string{"foo/bar/baz", "foo/baz/bar", "foo/zip/zap"},
154+
},
155+
exp{
156+
"foo/b",
157+
[]string{"foo/bar/baz", "foo/baz/bar"},
158+
},
159+
exp{
160+
"foo/ba",
161+
[]string{"foo/bar/baz", "foo/baz/bar"},
162+
},
163+
exp{
164+
"foo/bar",
165+
[]string{"foo/bar/baz"},
166+
},
167+
exp{
168+
"foo/bar/baz",
169+
[]string{"foo/bar/baz"},
170+
},
171+
exp{
172+
"foo/bar/bazoo",
173+
[]string{},
174+
},
175+
exp{
176+
"z",
177+
[]string{"zipzap"},
178+
},
179+
}
180+
181+
for _, test := range cases {
182+
out := []string{}
183+
fn := func(s string, v interface{}) bool {
184+
out = append(out, s)
185+
return false
186+
}
187+
r.WalkPrefix(test.inp, fn)
188+
sort.Strings(out)
189+
sort.Strings(test.out)
190+
if !reflect.DeepEqual(out, test.out) {
191+
t.Fatalf("mis-match: %v %v", out, test.out)
192+
}
193+
}
194+
}
195+
115196
// generateUUID is used to generate a random UUID
116197
func generateUUID() string {
117198
buf := make([]byte, 16)

0 commit comments

Comments
 (0)