You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
136 lines
3.3 KiB
136 lines
3.3 KiB
// 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)
|
|
})
|
|
}
|
|
|