mediocre-go-lib/miter/iter.go

// 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)
	})
}