Add miter package

miter/iter.go

@@ -0,0 +1,136 @@
+// Package miter implements a simple iterator type, as well as helper types
+// around it.
+package miter
+
+import (
+	"context"
+	"errors"
+)
+
+// ErrEnd is returned from Iterator's Next method when the iterator has
+// completed.
+var ErrEnd = errors.New("end")
+
+// Iterator is used to visit a sequence of elements one by one.
+type Iterator[T any] interface {
+
+	// Next will block until:
+	//
+	// - It returns the next item in the sequence
+	// - The Context is canceled and ctx.Err() is returned
+	// - ErrEnd is returned, indicating there are no more items in the sequence
+	// - Some other error is returned, indicating an unexpected error occurred
+	//
+	// If any non-context error is returned then Next should not be called
+	// again.
+	Next(ctx context.Context) (T, error)
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+type iterFn[T any] func(context.Context) (T, error)
+
+// FromFunc wraps a function such that it implements Iterator.
+func FromFunc[T any](fn func(context.Context) (T, error)) Iterator[T] {
+	return iterFn[T](fn)
+}
+
+func (f iterFn[T]) Next(ctx context.Context) (T, error) { return f(ctx) }
+
+// Error returns an Iterator which will always return the given error as the
+// error result.
+func Error[T any](err error) Iterator[T] {
+	return FromFunc(func(context.Context) (T, error) {
+		var zero T
+		return zero, err
+	})
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+// ToSlice consumes all items off the given Iterator until ErrEnd, creating and
+// returning a slice. If the Iterator ever returns any other error then that
+// error is returned, along with all consumed items so far.
+func ToSlice[T any](ctx context.Context, i Iterator[T]) ([]T, error) {
+	var s []T
+	for {
+		v, err := i.Next(ctx)
+		if errors.Is(err, ErrEnd) {
+			return s, nil
+		} else if err != nil {
+			return s, err
+		}
+
+		s = append(s, v)
+	}
+}
+
+type iterSlice[T any] []T
+
+// FromSlice wraps a slice such that it implements Iterator.
+func FromSlice[T any](s []T) Iterator[T] {
+	it := iterSlice[T](s)
+	return &it
+}
+
+func (s *iterSlice[T]) Next(context.Context) (T, error) {
+	if len(*s) == 0 {
+		var zero T
+		return zero, ErrEnd
+	}
+
+	next := (*s)[0]
+	*s = (*s)[1:]
+
+	return next, nil
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+type iterCh[T any] <-chan T
+
+// FromChannel wraps a channel such that it implements Iterator.
+func FromChannel[T any](ch <-chan T) Iterator[T] {
+	return iterCh[T](ch)
+}
+
+func (ch iterCh[T]) Next(ctx context.Context) (T, error) {
+	var (
+		v  T
+		ok bool
+	)
+
+	select {
+	case v, ok = <-ch:
+		if !ok {
+			return v, ErrEnd
+		}
+		return v, nil
+
+	case <-ctx.Done():
+		return v, ctx.Err()
+	}
+}
+
+////////////////////////////////////////////////////////////////////////////////
+
+// Lazily returns an Iterator which will initialize itself using the given
+// callback upon the first call to Next. If the callback returns an error then
+// this error is returned from all calls to Next.
+//
+// Lazily is useful for cases where an Iterator requires a Context to
+// initialize, but no Context is available.
+func Lazily[T any](
+	fn func(ctx context.Context) (Iterator[T], error),
+) Iterator[T] {
+	var i Iterator[T]
+	return FromFunc(func(ctx context.Context) (T, error) {
+		if i == nil {
+			var err error
+			if i, err = fn(ctx); err != nil {
+				i = Error[T](err)
+			}
+		}
+		return i.Next(ctx)
+	})
+}

miter/util.go

@@ -0,0 +1,102 @@
+package miter
+
+import (
+	"context"
+	"errors"
+)
+
+// Empty returns the empty Iterator, i.e. one which will immediately produce
+// ErrEnd.
+func Empty[T any]() Iterator[T] { return Error[T](ErrEnd) }
+
+// ForEach calls fn with each item read off the iterator, returning nil once
+// ErrEnd is returned from the Iterator or function. If the Iterator or function
+// return any other error then that is returned instead.
+func ForEach[T any](
+	ctx context.Context, i Iterator[T], fn func(T) error,
+) error {
+	for {
+		v, err := i.Next(ctx)
+		if errors.Is(err, ErrEnd) {
+			return nil
+		} else if err != nil {
+			return err
+		}
+
+		if err := fn(v); errors.Is(err, ErrEnd) {
+			return nil
+		} else if err != nil {
+			return err
+		}
+	}
+}
+
+// Map will read items from the given Iterator, pass them through the given
+// mapping function, and produce the mapped items from the returned Iterator.
+//
+// If the mapping function returns an error then that is returned from returned
+// Iterator, and no more calls should be made to it.
+func Map[T1, T2 any](i Iterator[T1], fn func(T1) (T2, error)) Iterator[T2] {
+	return FromFunc(func(ctx context.Context) (T2, error) {
+		v, err := i.Next(ctx)
+		if err != nil {
+			var zero T2
+			return zero, err
+		}
+
+		return fn(v)
+	})
+}
+
+// Concat concats all the given Iterators together into a single larger one.
+// Each Iterator will be consumed until ErrEnd in turn. Any other errors from
+// the inner Iterators will be returned as-is from the outer one.
+func Concat[T any](iters ...Iterator[T]) Iterator[T] {
+	var (
+		i    int
+		zero T
+	)
+
+	return FromFunc(func(ctx context.Context) (T, error) {
+		for {
+			if i >= len(iters) {
+				return zero, ErrEnd
+			}
+
+			v, err := iters[i].Next(ctx)
+			if errors.Is(err, ErrEnd) {
+				i++
+				continue
+			}
+
+			return v, err
+		}
+	})
+}
+
+// Filter returns an Iterator which will produce all items from the given
+// Iterator for which the function returns true. If the function returns any
+// error then that error is returned as-is.
+func Filter[T any](
+	i Iterator[T], fn func(context.Context, T) (bool, error),
+) Iterator[T] {
+	var zero T
+
+	return FromFunc(func(ctx context.Context) (T, error) {
+		for {
+			v, err := i.Next(ctx)
+			if err != nil {
+				return zero, err
+			}
+
+			keep, err := fn(ctx, v)
+			if err != nil {
+				return zero, err
+			}
+
+			if keep {
+				return v, nil
+			}
+		}
+	})
+}