Merge pull request #22 from tiborvass/fix_walk_delete

Allow calling Delete() inside Walk() and WalkPrefix()

Author: armon

radix.go

@@ -458,7 +458,7 @@ func (t *Tree) WalkPrefix(prefix string, fn WalkFn) {
 	n := t.root
 	search := prefix
 	for {
-		// Check for key exhaution
+		// Check for key exhaustion
 		if len(search) == 0 {
 			recursiveWalk(n, fn)
 			return
@@ -467,22 +467,20 @@ func (t *Tree) WalkPrefix(prefix string, fn WalkFn) {
 		// Look for an edge
 		n = n.getEdge(search[0])
 		if n == nil {
-			break
+			return
 		}
 
 		// Consume the search prefix
 		if strings.HasPrefix(search, n.prefix) {
 			search = search[len(n.prefix):]
-
-		} else if strings.HasPrefix(n.prefix, search) {
+			continue
+		}
+		if strings.HasPrefix(n.prefix, search) {
 			// Child may be under our search prefix
 			recursiveWalk(n, fn)
-			return
-		} else {
-			break
 		}
+		return
 	}
-
 }
 
 // WalkPath is used to walk the tree, but only visiting nodes
@@ -527,10 +525,27 @@ func recursiveWalk(n *node, fn WalkFn) bool {
 	}
 
 	// Recurse on the children
-	for _, e := range n.edges {
+	i := 0
+	k := len(n.edges) // keeps track of number of edges in previous iteration
+	for i < k {
+		e := n.edges[i]
 		if recursiveWalk(e.node, fn) {
 			return true
 		}
+		// It is a possibility that the WalkFn modified the node we are
+		// iterating on. If there are no more edges, mergeChild happened,
+		// so the last edge became the current node n, on which we'll
+		// iterate one last time.
+		if len(n.edges) == 0 {
+			return recursiveWalk(n, fn)
+		}
+		// If there are now less edges than in the previous iteration,
+		// then do not increment the loop index, since the current index
+		// points to a new edge. Otherwise, get to the next index.
+		if len(n.edges) >= k {
+			i++
+		}
+		k = len(n.edges)
 	}
 	return false
 }

radix_test.go

@@ -344,6 +344,40 @@ func TestWalkPath(t *testing.T) {
 	}
 }
 
+func TestWalkDelete(t *testing.T) {
+	r := New()
+	r.Insert("init0/0", nil)
+	r.Insert("init0/1", nil)
+	r.Insert("init0/2", nil)
+	r.Insert("init0/3", nil)
+	r.Insert("init1/0", nil)
+	r.Insert("init1/1", nil)
+	r.Insert("init1/2", nil)
+	r.Insert("init1/3", nil)
+	r.Insert("init2", nil)
+
+	deleteFn := func(s string, v interface{}) bool {
+		r.Delete(s)
+		return false
+	}
+
+	r.WalkPrefix("init1", deleteFn)
+
+	for _, s := range []string{"init0/0", "init0/1", "init0/2", "init0/3", "init2"} {
+		if _, ok := r.Get(s); !ok {
+			t.Fatalf("expecting to still find %q", s)
+		}
+	}
+	if n := r.Len(); n != 5 {
+		t.Fatalf("expected to find exactly 5 nodes, instead found %d: %v", n, r.ToMap())
+	}
+
+	r.Walk(deleteFn)
+	if n := r.Len(); n != 0 {
+		t.Fatalf("expected to find exactly 0 nodes, instead found %d: %v", n, r.ToMap())
+	}
+}
+
 // generateUUID is used to generate a random UUID
 func generateUUID() string {
 	buf := make([]byte, 16)