feat: composable iterator transforms

shils · shils · commit d939914f0051 · 2024-12-02T09:40:49.000-05:00
diff --git a/iterator/iter.go b/iterator/iter.go
@@ -0,0 +1,134 @@
+package iterator
+
+import (
+	"iter"
+)
+
+func Map[T, V any](it iter.Seq[T], fn func(T) V) iter.Seq[V] {
+	return func(yield func(V) bool) {
+		for t := range it {
+			if !yield(fn(t)) {
+				break
+			}
+		}
+	}
+}
+
+func Map2[T, V, W any](it iter.Seq[T], fn func(T) (V, W)) iter.Seq2[V, W] {
+	return func(yield func(V, W) bool) {
+		for t := range it {
+			if !yield(fn(t)) {
+				break
+			}
+		}
+	}
+}
+
+func Filter[T any](it iter.Seq[T], fn func(T) bool) iter.Seq[T] {
+	return func(yield func(T) bool) {
+		for t := range it {
+			if fn(t) && !yield(t) {
+				break
+			}
+		}
+	}
+}
+
+func Reduce[T, V any](it iter.Seq[T], fn func(V, T) V, init V) V {
+	acc := init
+	for t := range it {
+		acc = fn(acc, t)
+	}
+
+	return acc
+}
+
+// Zip returns an iterator that yields pairs of elements from the two input iterators, terminating when either
+// iterator is exhausted
+func Zip[T, U any](it iter.Seq[T], other iter.Seq[U]) iter.Seq2[T, U] {
+	return func(yield func(T, U) bool) {
+		next, stop := iter.Pull(other)
+		defer stop()
+		for t := range it {
+			u, ok := next()
+			if !ok || !yield(t, u) {
+				break
+			}
+		}
+	}
+}
+
+// ZipWith returns an iterator that yields the results of applying a function to pairs of elements from the two input
+// iterators, terminating when either iterator is exhausted
+func ZipWith[T, U, V any](it iter.Seq[T], other iter.Seq[U], fn func(T, U) V) iter.Seq[V] {
+	return func(yield func(V) bool) {
+		next, stop := iter.Pull(other)
+		defer stop()
+		for t := range it {
+			u, ok := next()
+			if !ok || !yield(fn(t, u)) {
+				break
+			}
+		}
+	}
+}
+
+func Indexed[T any](it iter.Seq[T]) iter.Seq2[int, T] {
+	return func(yield func(int, T) bool) {
+		i := 0
+		for t := range it {
+			if !yield(i, t) {
+				break
+			}
+			i++
+		}
+	}
+}
+
+func Concat[T any](its ...iter.Seq[T]) iter.Seq[T] {
+	return func(yield func(T) bool) {
+		for _, it := range its {
+			for t := range it {
+				if !yield(t) {
+					break
+				}
+			}
+		}
+	}
+}
+
+func Compact[T comparable](it iter.Seq[T]) iter.Seq[T] {
+	return func(yield func(T) bool) {
+		var last T
+		var started bool
+		for t := range it {
+			if started && t == last {
+				continue
+			}
+
+			if !yield(t) {
+				break
+			}
+			last = t
+			started = true
+		}
+	}
+}
+
+func CompactFunc[T any](it iter.Seq[T], eq func(T, T) bool) iter.Seq[T] {
+	return func(yield func(T) bool) {
+		var last T
+		var started bool
+		for t := range it {
+			if started && eq(t, last) {
+				continue
+			}
+
+			if !yield(t) {
+				break
+			}
+			last = t
+			started = true
+		}
+	}
+}
diff --git a/iterator/iter_test.go b/iterator/iter_test.go
@@ -0,0 +1,113 @@
+package iterator
+
+import (
+	"gotest.tools/v3/assert"
+	"maps"
+	"slices"
+	"strconv"
+	"strings"
+	"testing"
+)
+
+func TestMap(t *testing.T) {
+	it := slices.Values([]int{1, 2, 3, 4})
+	toString := func(n int) string {
+		return strconv.Itoa(n)
+	}
+
+	result := slices.Collect(Map(it, toString))
+	expected := []string{"1", "2", "3", "4"}
+
+	assert.DeepEqual(t, result, expected)
+
+}
+
+func TestMap2(t *testing.T) {
+	it := slices.Values([]string{"a", "three whole words", "two tokens"})
+
+	results := maps.Collect(Map2(it, func(str string) (int, int) {
+		words := strings.Fields(str)
+		return len(words), len(str)
+	}))
+
+	expected := map[int]int{1: 1, 2: 10, 3: 17}
+	assert.DeepEqual(t, results, expected)
+}
+
+func TestFilter(t *testing.T) {
+	it := slices.Values([]int{1, 2, 3, 4, 5})
+	isEven := func(n int) bool {
+		return n%2 == 0
+	}
+
+	result := slices.Collect(Filter(it, isEven))
+	assert.DeepEqual(t, result, []int{2, 4})
+}
+
+func TestReduce(t *testing.T) {
+	it := slices.Values([]int{1, 2, 3, 4})
+	sum := func(acc, n int) int {
+		return acc + n
+	}
+
+	result := Reduce(it, sum, 0)
+	assert.Equal(t, result, 10)
+}
+
+func TestZip(t *testing.T) {
+	it1 := slices.Values([]int{1, 2, 3})
+	it2 := slices.Values([]string{"a", "b", "c"})
+
+	result := maps.Collect(Zip(it1, it2))
+	expected := map[int]string{1: "a", 2: "b", 3: "c"}
+
+	assert.DeepEqual(t, result, expected)
+}
+
+func TestZipWith(t *testing.T) {
+	it1 := slices.Values([]int{1, 2, 3})
+	it2 := slices.Values([]string{"c", "b", "a"})
+
+	result := slices.Collect(ZipWith(it1, it2, func(n int, s string) string {
+		return strings.Repeat(s, n)
+	}))
+	expected := []string{"c", "bb", "aaa"}
+
+	assert.DeepEqual(t, result, expected)
+}
+
+func TestIndexed(t *testing.T) {
+	it := slices.Values([]string{"c", "b", "a"})
+	result := maps.Collect(Indexed(it))
+
+	expected := map[int]string{0: "c", 1: "b", 2: "a"}
+	assert.DeepEqual(t, result, expected)
+}
+
+func TestConcat(t *testing.T) {
+	it1 := slices.Values([]int{1, 2})
+	it2 := slices.Values([]int{3, 4})
+
+	result := slices.Collect(Concat(it1, it2))
+
+	expected := []int{1, 2, 3, 4}
+	assert.DeepEqual(t, result, expected)
+}
+
+func TestCompact(t *testing.T) {
+	it := slices.Values([]int{2, 2, 1, 3, 3, 3})
+	result := slices.Collect(Compact(it))
+
+	expected := []int{2, 1, 3}
+	assert.DeepEqual(t, result, expected)
+}
+
+func TestCompactFunc(t *testing.T) {
+	it := slices.Values([]int{2, 7, 1, 9, 9, 9})
+	result := slices.Collect(CompactFunc(it, func(a, b int) bool {
+		return a%5 == b%5
+	}))
+
+	expected := []int{2, 1, 9}
+	assert.DeepEqual(t, result, expected)
+}
diff --git a/iterator2/iter2.go b/iterator2/iter2.go
@@ -0,0 +1,42 @@
+package iterator2
+
+import "iter"
+
+func Map[T, U, V, W any](it iter.Seq2[T, U], fn func(T, U) (V, W)) iter.Seq2[V, W] {
+	return func(yield func(V, W) bool) {
+		for t, u := range it {
+			if !yield(fn(t, u)) {
+				break
+			}
+		}
+	}
+}
+
+func Map1[T, U, V any](it iter.Seq2[T, U], fn func(T, U) V) iter.Seq[V] {
+	return func(yield func(V) bool) {
+		for t, u := range it {
+			if !yield(fn(t, u)) {
+				break
+			}
+		}
+	}
+}
+
+func Filter[T, U any](it iter.Seq2[T, U], fn func(T, U) bool) iter.Seq2[T, U] {
+	return func(yield func(T, U) bool) {
+		for t, u := range it {
+			if fn(t, u) && !yield(t, u) {
+				break
+			}
+		}
+	}
+}
+
+func Reduce[T, U, V any](it iter.Seq2[T, U], fn func(V, T, U) V, init V) V {
+	acc := init
+	for t, u := range it {
+		acc = fn(acc, t, u)
+	}
+
+	return acc
+}
diff --git a/iterator2/iter2_test.go b/iterator2/iter2_test.go
@@ -0,0 +1,57 @@
+package iterator2
+
+import (
+	"go-exp/functions/hof"
+	"gotest.tools/v3/assert"
+	"maps"
+	"slices"
+	"strconv"
+	"strings"
+	"testing"
+)
+
+func TestMap(t *testing.T) {
+	it := slices.All([]int{1, 2, 3, 4})
+	toString := func(i, n int) (string, string) {
+		return strconv.Itoa(i), strconv.Itoa(n)
+	}
+
+	result := maps.Collect(Map(it, toString))
+	expected := map[string]string{"0": "1", "1": "2", "2": "3", "3": "4"}
+
+	assert.DeepEqual(t, result, expected)
+
+}
+
+func TestMap1(t *testing.T) {
+	it := slices.All([]int{1, 2, 3, 4})
+
+	result := slices.Collect(Map1(it, func(i, n int) string {
+		return strconv.Itoa(n)
+	}))
+	expected := []string{"1", "2", "3", "4"}
+
+	assert.DeepEqual(t, result, expected)
+}
+
+func TestFilter(t *testing.T) {
+	it := slices.All([]int{1, 2, 3, 4, 5})
+	isEven := func(n int) bool {
+		return n%2 == 0
+	}
+
+	result := maps.Collect(Filter(it, hof.LiftArity1Left[int, int, bool](isEven)))
+	expected := map[int]int{1: 2, 3: 4}
+
+	assert.DeepEqual(t, result, expected)
+}
+
+func TestReduce(t *testing.T) {
+	it := slices.All([]string{"d", "c", "b", "a"})
+	expand := func(acc string, i int, s string) string {
+		return acc + strings.Repeat(s, i)
+	}
+
+	result := Reduce(it, expand, "")
+	assert.Equal(t, result, "cbbaaa")
+}
