Rework the signature of WalkFn to allow setting new values

Before this commit, you could change the nodes during walk, assuming they were pointers.
But it was not possible to replace a node with a new one (unless they were stored as *interface{}, but that doesn's sound practical).

You could do a new `Insert`, but that be data racy if you e.g. partitioned the tree and walked and modified it concurrently.

Author: bep

go.mod

@@ -2,8 +2,9 @@ module github.com/gohugoio/go-radix
 
 go 1.24.0
 
+require github.com/frankban/quicktest v1.14.6
+
 require (
-	github.com/frankban/quicktest v1.14.6 // indirect
 	github.com/google/go-cmp v0.5.9 // indirect
 	github.com/kr/pretty v0.3.1 // indirect
 	github.com/kr/text v0.2.0 // indirect

radix.go

@@ -5,10 +5,29 @@ import (
 	"strings"
 )
 
+// WalkFlag is a bitmask return value for Walk functions.
+type WalkFlag uint32
+
+const (
+	WalkContinue WalkFlag = 0
+	WalkStop     WalkFlag = 1 << 0
+	WalkSet      WalkFlag = 1 << 1
+)
+
+// shouldStop returns true if the walk should terminate.
+func (w WalkFlag) shouldStop() bool {
+	return w&WalkStop != 0
+}
+
+// shouldSet returns true if the walk function wants to set a new value.
+func (w WalkFlag) shouldSet() bool {
+	return w&WalkSet != 0
+}
+
 // WalkFn is used when walking the tree. Takes a
-// key and value, returning if iteration should
-// be terminated.
-type WalkFn[T any] func(s string, v T) bool
+// key and value, returning a WalkFlag to indicate
+// how to proceed and possibly a new value to set.
+type WalkFn[T any] func(s string, v T) (WalkFlag, T)
 
 // leafNode is used to represent a value
 type leafNode[T any] struct {
@@ -313,10 +332,10 @@ func (t *Tree[T]) deletePrefix(parent, n *node[T], prefix string) int {
 	if len(prefix) == 0 {
 		// Remove the leaf node
 		subTreeSize := 0
-		//recursively walk from all edges of the node to be deleted
-		recursiveWalk(n, func(s string, v T) bool {
+		// recursively walk from all edges of the node to be deleted
+		recursiveWalk(n, func(s string, v T) (WalkFlag, T) {
 			subTreeSize++
-			return false
+			return WalkContinue, v
 		})
 		if n.isLeaf() {
 			n.leaf = nil
@@ -501,8 +520,14 @@ func (t *Tree[T]) WalkPath(path string, fn WalkFn[T]) {
 	search := path
 	for {
 		// Visit the leaf values if any
-		if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
-			return
+		if n.leaf != nil {
+			f, n2 := fn(n.leaf.key, n.leaf.val)
+			if f.shouldSet() {
+				n.leaf.val = n2
+			}
+			if f.shouldStop() {
+				return
+			}
 		}
 
 		// Check for key exhaution
@@ -529,8 +554,14 @@ func (t *Tree[T]) WalkPath(path string, fn WalkFn[T]) {
 // recursively. Returns true if the walk should be aborted
 func recursiveWalk[T any](n *node[T], fn WalkFn[T]) bool {
 	// Visit the leaf values if any
-	if n.leaf != nil && fn(n.leaf.key, n.leaf.val) {
-		return true
+	if n.leaf != nil {
+		f, n2 := fn(n.leaf.key, n.leaf.val)
+		if f.shouldSet() {
+			n.leaf.val = n2
+		}
+		if f.shouldStop() {
+			return true
+		}
 	}
 
 	// Recurse on the children
@@ -562,9 +593,10 @@ func recursiveWalk[T any](n *node[T], fn WalkFn[T]) bool {
 // ToMap is used to walk the tree and convert it into a map
 func (t *Tree[T]) ToMap() map[string]T {
 	out := make(map[string]T, t.size)
-	t.Walk(func(k string, v T) bool {
+	var zero T
+	t.Walk(func(k string, v T) (WalkFlag, T) {
 		out[k] = v
-		return false
+		return WalkContinue, zero
 	})
 	return out
 }

radix_test.go

@@ -28,8 +28,8 @@ func TestRadix(t *testing.T) {
 	r := NewFromMap(inp)
 	c.Assert(r.Len(), qt.Equals, len(inp))
 
-	r.Walk(func(k string, v int) bool {
-		return false
+	r.Walk(func(k string, v int) (WalkFlag, int) {
+		return WalkContinue, v
 	})
 
 	for k, v := range inp {
@@ -67,6 +67,45 @@ func TestRoot(t *testing.T) {
 	c.Assert(val, qt.IsTrue)
 }
 
+func TestWalkSet(t *testing.T) {
+	c := qt.New(t)
+	r := New[int]()
+
+	for i := range 10 {
+		r.Insert(fmt.Sprintf("key%d", i), i)
+	}
+
+	r.Walk(func(s string, v int) (WalkFlag, int) {
+		k := fmt.Sprintf("key%d", v)
+		c.Assert(s, qt.Equals, k)
+		v2 := v
+		if v%2 == 0 {
+			v2 = v * 10
+		}
+		return WalkSet, v2
+	})
+
+	var ints []int
+	r.Walk(func(s string, v int) (WalkFlag, int) {
+		ints = append(ints, v)
+		return WalkContinue, 0
+	})
+
+	c.Assert(ints, qt.DeepEquals, []int{0, 1, 20, 3, 40, 5, 60, 7, 80, 9})
+}
+
+func TestWalkFlag(t *testing.T) {
+	c := qt.New(t)
+
+	f := WalkSet
+	c.Assert(f.shouldSet(), qt.IsTrue)
+	c.Assert(f.shouldStop(), qt.IsFalse)
+
+	f = WalkSet | WalkStop
+	c.Assert(f.shouldSet(), qt.IsTrue)
+	c.Assert(f.shouldStop(), qt.IsTrue)
+}
+
 func TestDelete(t *testing.T) {
 	c := qt.New(t)
 	r := New[bool]()
@@ -110,9 +149,9 @@ func TestDeletePrefix(t *testing.T) {
 		c.Assert(deleted, qt.Equals, test.numDeleted)
 
 		out := []string{}
-		fn := func(s string, v bool) bool {
+		fn := func(s string, v bool) (WalkFlag, bool) {
 			out = append(out, s)
-			return false
+			return WalkContinue, false
 		}
 		r.Walk(fn)
 
@@ -133,7 +172,7 @@ func TestLongestPrefix(t *testing.T) {
 		"foozip",
 	}
 	for _, k := range keys {
-		r.Insert(k, nil)
+		r.Insert(k, "")
 	}
 	c.Assert(r.Len(), qt.Equals, len(keys))
 
@@ -167,7 +206,7 @@ func TestLongestPrefix(t *testing.T) {
 
 func TestWalkPrefix(t *testing.T) {
 	c := qt.New(t)
-	r := New[any]()
+	r := New[string]()
 
 	keys := []string{
 		"foobar",
@@ -177,7 +216,7 @@ func TestWalkPrefix(t *testing.T) {
 		"zipzap",
 	}
 	for _, k := range keys {
-		r.Insert(k, nil)
+		r.Insert(k, "")
 	}
 	c.Assert(r.Len(), qt.Equals, len(keys))
 
@@ -230,9 +269,9 @@ func TestWalkPrefix(t *testing.T) {
 
 	for _, test := range cases {
 		out := []string{}
-		fn := func(s string, v any) bool {
+		fn := func(s string, v string) (WalkFlag, string) {
 			out = append(out, s)
-			return false
+			return WalkContinue, ""
 		}
 		r.WalkPrefix(test.inp, fn)
 		sort.Strings(out)
@@ -243,7 +282,7 @@ func TestWalkPrefix(t *testing.T) {
 
 func TestWalkPath(t *testing.T) {
 	c := qt.New(t)
-	r := New[any]()
+	r := New[string]()
 
 	keys := []string{
 		"foo",
@@ -254,7 +293,7 @@ func TestWalkPath(t *testing.T) {
 		"zipzap",
 	}
 	for _, k := range keys {
-		r.Insert(k, nil)
+		r.Insert(k, "")
 	}
 	c.Assert(r.Len(), qt.Equals, len(keys))
 
@@ -299,9 +338,9 @@ func TestWalkPath(t *testing.T) {
 
 	for _, test := range cases {
 		out := []string{}
-		fn := func(s string, v any) bool {
+		fn := func(s string, v string) (WalkFlag, string) {
 			out = append(out, s)
-			return false
+			return WalkContinue, ""
 		}
 		r.WalkPath(test.inp, fn)
 		sort.Strings(out)
@@ -312,20 +351,20 @@ func TestWalkPath(t *testing.T) {
 
 func TestWalkDelete(t *testing.T) {
 	c := qt.New(t)
-	r := New[any]()
-	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 any) bool {
+	r := New[string]()
+	r.Insert("init0/0", "")
+	r.Insert("init0/1", "")
+	r.Insert("init0/2", "")
+	r.Insert("init0/3", "")
+	r.Insert("init1/0", "")
+	r.Insert("init1/1", "")
+	r.Insert("init1/2", "")
+	r.Insert("init1/3", "")
+	r.Insert("init2", "")
+
+	deleteFn := func(s string, v string) (WalkFlag, string) {
 		r.Delete(s)
-		return false
+		return WalkContinue, ""
 	}
 
 	r.WalkPrefix("init1", deleteFn)
@@ -387,24 +426,24 @@ func BenchmarkRadix(b *testing.B) {
 
 	b.Run("Walk", func(b *testing.B) {
 		for b.Loop() {
-			r.Walk(func(s string, v *v) bool {
-				return false
+			r.Walk(func(s string, v *v) (WalkFlag, *v) {
+				return WalkContinue, nil
 			})
 		}
 	})
 
 	b.Run("WalkPrefix", func(b *testing.B) {
 		for b.Loop() {
-			r.WalkPrefix("init50", func(s string, v *v) bool {
-				return false
+			r.WalkPrefix("init50", func(s string, v *v) (WalkFlag, *v) {
+				return WalkContinue, nil
 			})
 		}
 	})
 
 	b.Run("WalkPath", func(b *testing.B) {
 		for b.Loop() {
-			r.WalkPath("init50/50", func(s string, v *v) bool {
-				return false
+			r.WalkPath("init50/50", func(s string, v *v) (WalkFlag, *v) {
+				return WalkContinue, nil
 			})
 		}
 	})