Implement Count, Take, Lift and Collect for rangefunc

Author: BooleanCat

v2/internal/assert/assert.go

@@ -0,0 +1,66 @@
+package assert
+
+import "testing"
+
+func Equal[T comparable](t *testing.T, a, b T) {
+	t.Helper()
+
+	if a != b {
+		t.Errorf("expected `%v` to equal `%v`", a, b)
+	}
+}
+
+func True(t *testing.T, b bool) {
+	t.Helper()
+
+	if !b {
+		t.Error("expected `false` to be `true`")
+	}
+}
+
+func False(t *testing.T, b bool) {
+	t.Helper()
+
+	if b {
+		t.Error("expected `true` to be `false`")
+	}
+}
+
+func Nil(t *testing.T, v interface{}) {
+	t.Helper()
+
+	if v != nil {
+		t.Errorf("expected `%v` to equal `nil`", v)
+	}
+}
+
+func NotNil(t *testing.T, v interface{}) {
+	t.Helper()
+
+	if v == nil {
+		t.Error("expected `nil` not to equal `nil`")
+	}
+}
+
+func SliceEqual[T comparable](t *testing.T, a, b []T) {
+	t.Helper()
+
+	if len(a) != len(b) {
+		t.Errorf("expected `%v` to equal `%v` but lengths differ", a, b)
+		return
+	}
+
+	for i, v := range a {
+		if v != b[i] {
+			t.Errorf("expected `%v` to equal `%v`", a, b)
+		}
+	}
+}
+
+func Empty[S any, T []S | ~string](t *testing.T, items T) {
+	t.Helper()
+
+	if len(items) != 0 {
+		t.Errorf("expected `%v` to be empty", items)
+	}
+}

v2/iter/count.go

@@ -0,0 +1,14 @@
+package iter
+
+import "iter"
+
+// Count yields an infinite sequence of integers, starting from 0.
+func Count() Iterator[int] {
+	return Iterator[int](iter.Seq[int](func(yield func(int) bool) {
+		for i := 0; ; i++ {
+			if !yield(i) {
+				return
+			}
+		}
+	}))
+}

v2/iter/count_test.go

@@ -0,0 +1,25 @@
+package iter_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/BooleanCat/go-functional/v2/internal/assert"
+	"github.com/BooleanCat/go-functional/v2/iter"
+)
+
+func ExampleCount() {
+	for i := range iter.Count().Take(3) {
+		fmt.Println(i)
+	}
+
+	// Output:
+	// 0
+	// 1
+	// 2
+}
+
+func TestCount(t *testing.T) {
+	numbers := iter.Count().Take(3).Collect()
+	assert.SliceEqual(t, []int{0, 1, 2}, numbers)
+}

v2/iter/iter.go

@@ -2,4 +2,33 @@ package iter
 
 import "iter"
 
-type Iterator[T any] iter.Seq[T]
+// Iterator is a wrapper around the iter.Seq type that allowed for method
+// chaining of iterators found in this package.
+type Iterator[V any] iter.Seq[V]
+
+// Lift yields all items in the provided slice.
+func Lift[V any](slice []V) Iterator[V] {
+	return Iterator[V](iter.Seq[V](func(yield func(V) bool) {
+		for _, item := range slice {
+			if !yield(item) {
+				return
+			}
+		}
+	}))
+}
+
+// Collect consumes an iterator and returns a slice of all items yielded.
+func Collect[V any](iter Iterator[V]) []V {
+	collection := make([]V, 0)
+
+	for item := range iter {
+		collection = append(collection, item)
+	}
+
+	return collection
+}
+
+// Collect is a convenience method for chaining [Collect] after an iterator.
+func (iter Iterator[V]) Collect() []V {
+	return Collect[V](iter)
+}

v2/iter/iter_test.go

@@ -1,9 +1,54 @@
 package iter_test
 
-import "testing"
+import (
+	"fmt"
+	"testing"
 
-func TestNothing(t *testing.T) {
-	if false {
-		t.Errorf("This should not happen")
+	"github.com/BooleanCat/go-functional/v2/internal/assert"
+	"github.com/BooleanCat/go-functional/v2/iter"
+)
+
+func ExampleLift() {
+	for i := range iter.Lift([]int{1, 2}) {
+		fmt.Println(i)
+	}
+
+	// Output:
+	// 1
+	// 2
+}
+
+func TestLift(t *testing.T) {
+	number := 0
+
+	for i := range iter.Lift([]int{1, 2}) {
+		assert.Equal(t, number+1, i)
+		number++
+	}
+}
+
+func TestLiftEmpty(t *testing.T) {
+	for _ = range iter.Lift([]int{}) {
+		t.Error("expected no iteration")
 	}
 }
+
+func ExampleCollect() {
+	fmt.Println(iter.Collect(iter.Lift([]int{1, 2})))
+	// Output: [1 2]
+}
+
+func ExampleCollect_method() {
+	fmt.Println(iter.Lift([]int{1, 2}).Collect())
+	// Output: [1 2]
+}
+
+func TestCollect(t *testing.T) {
+	numbers := iter.Lift([]int{1, 2, 3}).Collect()
+	assert.SliceEqual(t, []int{1, 2, 3}, numbers)
+}
+
+func TestCollectEmpty(t *testing.T) {
+	numbers := iter.Lift([]int{}).Collect()
+	assert.Empty[int](t, numbers)
+}

v2/iter/take.go

@@ -0,0 +1,30 @@
+package iter
+
+import "iter"
+
+// Take limits the number of elements yielded by a delegate iterator to a
+// maximum limit.
+func Take[V any](delegate Iterator[V], limit int) Iterator[V] {
+	return Iterator[V](iter.Seq[V](func(yield func(V) bool) {
+		next, stop := iter.Pull(iter.Seq[V](delegate))
+		defer stop()
+
+		for {
+			limit--
+
+			if limit < 0 {
+				return
+			}
+
+			v, ok := next()
+			if !ok || !yield(v) {
+				return
+			}
+		}
+	}))
+}
+
+// Take is a convenience method for chaining [Take] after an iterator.
+func (iter Iterator[V]) Take(limit int) Iterator[V] {
+	return Take[V](iter, limit)
+}

v2/iter/take_test.go

@@ -0,0 +1,39 @@
+package iter_test
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/BooleanCat/go-functional/v2/internal/assert"
+	"github.com/BooleanCat/go-functional/v2/iter"
+)
+
+func ExampleTake() {
+	for i := range iter.Take(iter.Lift([]int{1, 2, 3}), 2) {
+		fmt.Println(i)
+	}
+
+	// Output:
+	// 1
+	// 2
+}
+
+func ExampleTake_method() {
+	for i := range iter.Lift([]int{1, 2, 3}).Take(2) {
+		fmt.Println(i)
+	}
+
+	// Output:
+	// 1
+	// 2
+}
+
+func TestTake(t *testing.T) {
+	numbers := iter.Lift([]int{1, 2, 3}).Take(2).Collect()
+	assert.SliceEqual(t, []int{1, 2}, numbers)
+}
+
+func TestTakeEmpty(t *testing.T) {
+	numbers := iter.Lift([]int{}).Take(2).Collect()
+	assert.Empty[int](t, numbers)
+}