diff --git a/miter/iter.go b/miter/iter.go new file mode 100644 index 0000000..5bc0134 --- /dev/null +++ b/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) + }) +} diff --git a/miter/util.go b/miter/util.go new file mode 100644 index 0000000..8d79ce4 --- /dev/null +++ b/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 + } + } + }) +}