Add some iterator utilities to go with the v1.23.0 upgrade. (#4238)

* Add some iterator utilities to go with the v1.23.0 upgrade.

* Ditch the unused var.

* Update from PR and add index and delete functions to sliceutil.

Author: scr-oath

util/iterutil/doc.go

@@ -0,0 +1,36 @@
+/*
+Package iterators provides a set of iterators for various data structures.
+
+# Elements of slices as pointers
+
+The main motivation is making it easier to work with slices of large structures, where people often choose an iterator
+pattern, which makes a shallow copy of each struct in the slice:
+
+	var myStructs []MyStruct
+	for _, myStruct := range myStructs {
+		// myStruct is a shallow *COPY* of each element - any changes in the loop have no effect on the slice element.
+	}
+
+In order to improve that, and also not have to use the throw-away `_` variable, [SlicePointerValues] may be used to
+iterate across pointers to each element in the slice:
+
+	var myStructs []MyStruct
+	for myStructPtr := range iterators.SlicePointerValues(myStructs) {
+		// myStructPtr is a pointer to each element - any changes in the loop will be reflected in the slice element.
+	}
+
+If you still need the index, you can use [SlicePointers]:
+
+	var myStructs []MyStruct
+	for i, myStructPtr := range iterators.SlicePointers(myStructs) {
+		// myStructPtr is a pointer to each element - any changes in the loop will be reflected in the slice element.
+		// i is the index of the element.
+	}
+
+# Walking json
+
+The https://github.com/tidwall/gjson library is already included as a dependency. If you need to walk a json document,
+the [WalkGjsonLeaves] iterator is available, which will yield the json paths to all leaves in the document along with
+the [gjson.Result] for each leaf.
+*/
+package iterators

util/iterutil/gjson.go

@@ -0,0 +1,40 @@
+package iterators
+
+import (
+	"iter"
+
+	"github.com/tidwall/gjson"
+)
+
+type (
+	GjsonElem struct {
+		path   string
+		result gjson.Result
+	}
+)
+
+// WalkGjsonLeaves returns an iterator of (path, result) from a gjson result; only leaves are yielded.
+func WalkGjsonLeaves(result gjson.Result) iter.Seq2[string, gjson.Result] {
+	return func(yield func(string, gjson.Result) bool) {
+		stack := []GjsonElem{{path: "", result: result}}
+		for stackLen := len(stack); stackLen > 0; stackLen = len(stack) {
+			elem := stack[stackLen-1]
+			stack = stack[:stackLen-1]
+			if elem.result.Type == gjson.JSON {
+				if elem.path != "" {
+					elem.result.ForEach(func(key, value gjson.Result) bool {
+						stack = append(stack, GjsonElem{path: elem.path + "." + key.String(), result: value})
+						return true
+					})
+				} else {
+					elem.result.ForEach(func(key, value gjson.Result) bool {
+						stack = append(stack, GjsonElem{path: key.String(), result: value})
+						return true
+					})
+				}
+			} else if !yield(elem.path, elem.result) {
+				return
+			}
+		}
+	}
+}

util/iterutil/gjson_test.go

@@ -0,0 +1,54 @@
+package iterators
+
+import (
+	"slices"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/tidwall/gjson"
+)
+
+func TestWalkGjsonLeaves(t *testing.T) {
+	tests := []struct {
+		name      string
+		json      string
+		wantPaths []string
+	}{
+		{
+			name:      "true",
+			json:      "true",
+			wantPaths: []string{""},
+		},
+		{
+			name:      "three fields",
+			json:      `{"foo": "bar", "num": 42, "bool": true}`,
+			wantPaths: []string{"bool", "foo", "num"},
+		},
+		{
+			name:      "deep struct",
+			json:      `{"foo": {"bar": {"baz": "qux"}}, "num": 42, "bool": true}`,
+			wantPaths: []string{"bool", "foo.bar.baz", "num"},
+		},
+		{
+			name: "deep struct with array",
+			json: `{"foo": {"bar": {"baz": "qux", "biz": [1,2,3,4]}}, "num": 42, "bool": true}`,
+			wantPaths: []string{
+				"bool",
+				"foo.bar.baz",
+				"foo.bar.biz.0",
+				"foo.bar.biz.1",
+				"foo.bar.biz.2",
+				"foo.bar.biz.3",
+				"num",
+			},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			gjsonResult := gjson.Parse(tt.json)
+			gotPaths := slices.Collect(Firsts(WalkGjsonLeaves(gjsonResult)))
+			slices.Sort(gotPaths)
+			assert.Equal(t, tt.wantPaths, gotPaths)
+		})
+	}
+}

util/iterutil/iter.go

@@ -0,0 +1,25 @@
+package iterators
+
+import "iter"
+
+// Firsts returns an iterator that yields the first elements of a sequence of pairs.
+func Firsts[T, U any](seq2 iter.Seq2[T, U]) iter.Seq[T] {
+	return func(yield func(T) bool) {
+		for t := range seq2 {
+			if !yield(t) {
+				return
+			}
+		}
+	}
+}
+
+// Seconds returns an iterator that yields the second elements of a sequence of pairs.
+func Seconds[T, U any](seq2 iter.Seq2[T, U]) iter.Seq[U] {
+	return func(yield func(U) bool) {
+		for _, u := range seq2 {
+			if !yield(u) {
+				return
+			}
+		}
+	}
+}

util/iterutil/iter_test.go

@@ -0,0 +1,37 @@
+package iterators
+
+import (
+	"iter"
+	"slices"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func testIter() iter.Seq2[int, string] {
+	return func(yield func(int, string) bool) {
+		if !yield(1, "one") {
+			return
+		}
+		if !yield(2, "two") {
+			return
+		}
+		if !yield(3, "three") {
+			return
+		}
+	}
+}
+
+func TestFirsts(t *testing.T) {
+	firsts := Firsts(testIter())
+	want := []int{1, 2, 3}
+	got := slices.Collect(firsts)
+	assert.Equal(t, want, got)
+}
+
+func TestSeconds(t *testing.T) {
+	seconds := Seconds(testIter())
+	want := []string{"one", "two", "three"}
+	got := slices.Collect(seconds)
+	assert.Equal(t, want, got)
+}

util/iterutil/slices.go

@@ -0,0 +1,25 @@
+package iterators
+
+import "iter"
+
+// SlicePointers returns an iterator that yields indices and pointers to the elements of a slice.
+func SlicePointers[Slice ~[]T, T any](s Slice) iter.Seq2[int, *T] {
+	return func(yield func(int, *T) bool) {
+		for i := range s {
+			if !yield(i, &s[i]) {
+				return
+			}
+		}
+	}
+}
+
+// SlicePointerValues returns an iterator that yields pointers to the elements of a slice.
+func SlicePointerValues[Slice ~[]T, T any](s Slice) iter.Seq[*T] {
+	return func(yield func(*T) bool) {
+		for i := range s {
+			if !yield(&s[i]) {
+				return
+			}
+		}
+	}
+}

util/iterutil/slices_test.go

@@ -0,0 +1,32 @@
+package iterators
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type Simple[T any] struct {
+	Value T
+}
+
+func TestSlicePointers(t *testing.T) {
+	s := []Simple[int]{{1}, {2}, {3}}
+	for i, v := range SlicePointers(s) {
+		v.Value = 99999
+		assert.EqualValues(t, 99999, s[i].Value)
+	}
+	for _, v := range s {
+		assert.EqualValues(t, 99999, v.Value)
+	}
+}
+
+func TestSlicePointerValues(t *testing.T) {
+	s := []Simple[int]{{1}, {2}, {3}}
+	for v := range SlicePointerValues(s) {
+		v.Value = 99999
+	}
+	for _, v := range s {
+		assert.EqualValues(t, 99999, v.Value)
+	}
+}

util/sliceutil/slices.go

@@ -0,0 +1,27 @@
+package sliceutil
+
+// IndexPointerFunc returns the index of the first element in the slice for which the function f returns true.
+func IndexPointerFunc[Slice ~[]T, T any](s Slice, f func(*T) bool) int {
+	for i := range s {
+		if f(&s[i]) {
+			return i
+		}
+	}
+	return -1
+}
+
+// DeletePointerFunc deletes all elements from the slice for which the function f returns true.
+func DeletePointerFunc[Slice ~[]T, T any](s Slice, f func(*T) bool) Slice {
+	i := IndexPointerFunc(s, f)
+	if i == -1 {
+		return s
+	}
+	for j := i + 1; j < len(s); j++ {
+		if v := &s[j]; !f(v) {
+			s[i] = *v
+			i++
+		}
+	}
+	clear(s[i:]) // zero/nil out the obsolete elements, for GC
+	return s[:i]
+}

util/sliceutil/slices_test.go

@@ -0,0 +1,29 @@
+package sliceutil
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+type Simple[T any] struct {
+	Value T
+}
+
+func TestIndexPointerFunc(t *testing.T) {
+	s := []Simple[int]{{1}, {2}, {3}}
+	assert.Equal(t, 1, IndexPointerFunc(s, func(v *Simple[int]) bool { return v.Value == 2 }))
+	assert.Equal(t, -1, IndexPointerFunc(s, func(v *Simple[int]) bool { return v.Value == 4 }))
+}
+
+func TestDeletePointerFunc(t *testing.T) {
+	s := []Simple[int]{{1}, {2}, {3}, {4}, {5}}
+	s = DeletePointerFunc(s, func(v *Simple[int]) bool { return v.Value%2 == 0 })
+	assert.Equal(t, []Simple[int]{{1}, {3}, {5}}, s)
+
+	s = DeletePointerFunc(s, func(v *Simple[int]) bool { return v.Value > 10 })
+	assert.Equal(t, []Simple[int]{{1}, {3}, {5}}, s)
+
+	s = DeletePointerFunc(s, func(v *Simple[int]) bool { return true })
+	assert.Empty(t, s)
+}