WIP: iterators

README.md

@@ -76,6 +76,7 @@ Genesis contains the following packages:
 + [🗺 maps](https://pkg.go.dev/github.com/life4/genesis/maps): generic functions for maps (`map[K]V`).
 + [📺 channels](https://pkg.go.dev/github.com/life4/genesis/channels): generic function for channels (`chan T`).
 + [⚙️ sets](https://pkg.go.dev/github.com/life4/genesis/sets): generic function for sets (`map[T]struct{}`).
++ [🦥 iters](https://pkg.go.dev/github.com/life4/genesis/iters): generic function for lazy iteration.
 + [🛟 lambdas](https://pkg.go.dev/github.com/life4/genesis/lambdas): helper generic functions to work with `slices.Map` and similar.
 
 See [📄 DOCUMENTATION](https://pkg.go.dev/github.com/life4/genesis) for more info.

channels/channel_ctx.go

@@ -94,7 +94,6 @@ func CountC[T comparable](ctx context.Context, c <-chan T, el T) int {
 }
 
 // Drop drops first n elements from channel c and returns a new channel with the rest.
-// It returns channel do be unblocking. If you want array instead, wrap result into TakeAll.
 func DropC[T any](ctx context.Context, c <-chan T, n int) chan T {
 	result := make(chan T)
 	go func() {

doc.go

@@ -4,6 +4,7 @@ package genesis
 // required for being discovered by pkg.go.dev
 import (
 	_ "github.com/life4/genesis/channels"
+	_ "github.com/life4/genesis/iters"
 	_ "github.com/life4/genesis/lambdas"
 	_ "github.com/life4/genesis/maps"
 	_ "github.com/life4/genesis/sets"

iters/README.md

@@ -0,0 +1,3 @@
+# genesis/iters
+
+Package iters provides generic functions for lazy iteration.

iters/allocs_test.go

@@ -0,0 +1,27 @@
+package iters_test
+
+import (
+	"testing"
+
+	"github.com/life4/genesis/iters"
+	"github.com/matryer/is"
+)
+
+func assertAllocs(t *testing.T, expected uint64, f func()) {
+	res := testing.Benchmark(func(b *testing.B) {
+		b.ReportAllocs()
+		f()
+	})
+	is := is.New(t)
+	is.Equal(res.MemAllocs, expected)
+}
+
+func TestFromSlice_Allocs(t *testing.T) {
+	s := make([]int, 1000)
+	assertAllocs(t, 2, func() {
+		next := iters.FromSlice(s)
+		next()
+		next()
+		next()
+	})
+}

iters/doc.go

@@ -0,0 +1,7 @@
+// 🦥 Package iters provides generic functions for lazy iteration.
+//
+// Iterators are useful for single-threaded processing of large amount of items.
+// For a small collection, slices should give better performance for the price of
+// higher memory consumption. For big collection with possibility of concurrency,
+// you can get better performance by using channels and goroutines.
+package iters

iters/iters.go

@@ -0,0 +1,146 @@
+package iters
+
+import c "github.com/life4/genesis/constraints"
+
+// Next returns the next element from the iterator.
+//
+// The second return value indicates if there are more values to pull.
+// If the iterator is exhausted, the first value is the default value
+// of the type and second is false. When the iterator is exhausted,
+// repeated attempts to get the next value should produce the same
+// default+false result.
+//
+// In other words, it should behave like pulling from a (closed) channel.
+//
+// The code using an iterator doesn't guarantee to exhaust it.
+// For example, [Take] only takes the number of elements it needs
+// and never calls Next again. Hence you shouldn't rely on Next
+// for closing connections and cleaning up unused resources.
+// If your iterator needs to provide logic like this, you should
+// implement a Close method and defer it.
+//
+// An iterator is allowed to be infinite and never return false.
+type Next[T any] func() (T, bool)
+
+// Drop returns an iterator dropping the first n items from the given iterator.
+//
+// When the resulting iterator is called for the first time,
+// it will drop the first n item from the input iterator.
+// All consecutive calls to the iterator will be forwarded
+// to the input iterator.
+func Drop[T any, I c.Integer](next Next[T], n I) Next[T] {
+	return func() (T, bool) {
+		if n != 0 {
+			for ; n > 0; n-- {
+				val, more := next()
+				if !more {
+					return val, more
+				}
+			}
+		}
+		return next()
+	}
+}
+
+// Filter returns an iterator of elements from the given iterator for which the function returns true.
+func Filter[T any](next Next[T], f func(T) bool) Next[T] {
+	return func() (T, bool) {
+		for {
+			val, more := next()
+			if !more {
+				return val, false
+			}
+			if f(val) {
+				return val, true
+			}
+		}
+	}
+}
+
+// FromChannel produces an iterator returning elements from the given channel.
+//
+// Each call to Iter will pull from the channel, which means
+// you have to make sure it won't block forever. It's a good idea
+// to make the channel cancelable by using channels.WithContext.
+func FromChannel[T any](ch <-chan T) Next[T] {
+	return func() (T, bool) {
+		v, ok := <-ch
+		return v, ok
+	}
+}
+
+// FromSlice produces an iterator returning elements from the given slice.
+func FromSlice[S ~[]T, T any](slice S) Next[T] {
+	next := 0
+	return func() (T, bool) {
+		if next >= len(slice) {
+			return *new(T), false
+		}
+		v := slice[next]
+		next += 1
+		return v, true
+	}
+}
+
+// Map returns an iterator of results of applying the function to each element of the given iterator.
+func Map[T, R any](next Next[T], f func(T) R) Next[R] {
+	return func() (R, bool) {
+		val, more := next()
+		if !more {
+			var res R
+			return res, false
+		}
+		res := f(val)
+		return res, true
+	}
+}
+
+// Reduce applies the function to acc and every iterator element and returns the acc.
+func Reduce[T, R any](next Next[T], acc R, f func(T, R) R) R {
+	for {
+		val, more := next()
+		if !more {
+			return acc
+		}
+		acc = f(val, acc)
+	}
+}
+
+// Take returns an iterator returning only the first n items from the given iterator.
+//
+// When n items are consumed, Take will not call Next on the input iterator again.
+// So, it's possible for the input iterator to not be fully exhausted.
+//
+// If the input iterator returns fewer than n items, Take will just stop and
+// not generate additional items.
+func Take[T any, I c.Integer](next Next[T], n I) Next[T] {
+	return func() (T, bool) {
+		if n <= 0 {
+			var val T
+			return val, false
+		}
+		n -= 1
+		return next()
+	}
+}
+
+// ToSlice converts the given iterator to a slice.
+//
+// The function returns only when there are no more elements to consume
+// from the iterator. It's a good idea to use [Take] to limit the number
+// of elements if it's possible for the iterator to be infinite or just too big.
+//
+// Also, you should make sure that the iterator doesn't block forever.
+// In particular, when creating an iterator from a channel using [FromChannel],
+// you may want to use channels.WithContext and set a deadline or cancelation
+// on that context.
+func ToSlice[T any](next Next[T]) []T {
+	res := make([]T, 0)
+	for {
+		val, more := next()
+		if !more {
+			return res
+		}
+		res = append(res, val)
+	}
+}

iters/iters_test.go

@@ -0,0 +1,85 @@
+package iters_test
+
+import (
+	"testing"
+
+	"github.com/life4/genesis/iters"
+	"github.com/matryer/is"
+)
+
+var ts = iters.ToSlice[int]
+
+func new[T any](vs ...T) iters.Next[T] {
+	return iters.FromSlice(vs)
+}
+
+func TestDrop(t *testing.T) {
+	is := is.NewRelaxed(t)
+	is.Equal(ts(iters.Drop(new(3, 4, 5), 1)), []int{4, 5})
+	is.Equal(ts(iters.Drop(new(3, 4, 5), 2)), []int{5})
+	is.Equal(ts(iters.Drop(new(3, 4, 5), 0)), []int{3, 4, 5})
+	is.Equal(ts(iters.Drop(new(3, 4, 5), -5)), []int{3, 4, 5})
+	is.Equal(ts(iters.Drop(new(3, 4, 5), 3)), []int{})
+	is.Equal(ts(iters.Drop(new(3, 4, 5), 6)), []int{})
+}
+
+func TestFilter(t *testing.T) {
+	is := is.NewRelaxed(t)
+	even := func(x int) bool { return x%2 == 0 }
+	is.Equal(ts(iters.Filter(new(3, 4, 5, 6), even)), []int{4, 6})
+	is.Equal(ts(iters.Filter(new(4, 6), even)), []int{4, 6})
+	is.Equal(ts(iters.Filter(new(3, 5), even)), []int{})
+	is.Equal(ts(iters.Filter(new[int](), even)), []int{})
+}
+
+func TestFromChannel(t *testing.T) {
+	is := is.New(t)
+	ch := make(chan int)
+	go func() {
+		ch <- 3
+		ch <- 4
+		ch <- 5
+		close(ch)
+	}()
+	is.Equal(ts(iters.FromChannel(ch)), []int{3, 4, 5})
+}
+
+func TestFromSlice(t *testing.T) {
+	is := is.NewRelaxed(t)
+	is.Equal(ts(iters.FromSlice([]int{3, 4, 5})), []int{3, 4, 5})
+	is.Equal(ts(iters.FromSlice([]int{3})), []int{3})
+	is.Equal(ts(iters.FromSlice([]int{})), []int{})
+	is.Equal(ts(iters.FromSlice([]int(nil))), []int{})
+}
+
+func TestMap(t *testing.T) {
+	is := is.NewRelaxed(t)
+	double := func(x int) int { return x * 2 }
+	is.Equal(ts(iters.Map(new(3, 4, 5), double)), []int{6, 8, 10})
+	is.Equal(ts(iters.Map(new[int](), double)), []int{})
+}
+
+func TestReduce(t *testing.T) {
+	is := is.NewRelaxed(t)
+	add := func(x, a int) int { return x + a }
+	is.Equal(iters.Reduce(new(3, 4, 5), 0, add), 12)
+	is.Equal(iters.Reduce(new(3, 4, 5), 3, add), 15)
+	is.Equal(iters.Reduce(new[int](), 0, add), 0)
+	is.Equal(iters.Reduce(new[int](), 7, add), 7)
+}
+
+func TestTake(t *testing.T) {
+	is := is.NewRelaxed(t)
+	is.Equal(ts(iters.Take(new(3, 4, 5), 2)), []int{3, 4})
+	is.Equal(ts(iters.Take(new(3, 4, 5), 1)), []int{3})
+	is.Equal(ts(iters.Take(new(3, 4, 5), 10)), []int{3, 4, 5})
+	is.Equal(ts(iters.Take(new(3, 4, 5), 0)), []int{})
+	is.Equal(ts(iters.Take(new(3, 4, 5), -1)), []int{})
+	is.Equal(ts(iters.Take(new(3, 4, 5), -10)), []int{})
+}
+
+func TestToSlice(t *testing.T) {
+	is := is.NewRelaxed(t)
+	is.Equal(iters.ToSlice(new(3, 4, 5)), []int{3, 4, 5})
+	is.Equal(iters.ToSlice(new[int]()), []int{})
+}

slices/examples_test.go

@@ -3,6 +3,7 @@ package slices_test
 import (
 	"errors"
 	"fmt"
+	"sync/atomic"
 
 	"github.com/life4/genesis/channels"
 	"github.com/life4/genesis/slices"
@@ -205,10 +206,10 @@ func ExampleEach() {
 }
 
 func ExampleEachAsync() {
-	s := []int{4, 5, 6}
-	sum := 0
-	slices.EachAsync(s, 0, func(x int) {
-		sum += x
+	s := []uint32{4, 5, 6}
+	var sum uint32
+	slices.EachAsync(s, 0, func(x uint32) {
+		atomic.AddUint32(&sum, x)
 	})
 	fmt.Println(sum)
 	// Output: 15