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

tiborvass · tiborvass · commit f30034d78803 · 2022-11-17T13:08:30.000-08:00
In recursiveWalk called by Walk and WalkPrefix, we were iterating over
the edges of a node. The edges as well as the current node object itself
(when mergeChild is called) can be modified by calls to Delete.

This commit modifies the loop to handle modifications correctly and
prevent a nil pointer dereference panic in the case detailed in TestWalkDelete.

Signed-off-by: Tibor Vass &lt;teabee89@gmail.com&gt;
diff --git a/radix.go b/radix.go
@@ -525,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
 }
diff --git a/radix_test.go b/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)
