iterutils.go

package iterutils

import (
	"iter"

	"golang.org/x/exp/constraints"
)

// Predicate function
type Pred[T any] func(T) bool

// [Pred] function for bool values
func Bool[T bool](v bool) bool {
	return v
}

// Equal function
type Eq[T any] func(T, T) bool

// [Eq] function for comparable values
func Equ[T comparable](v1, v2 T) bool {
	return v1 == v2
}

// Compare function.
//
// Returns -1 if the first value is less then the second, 0 if both are equal, 1 if the first value is greater then the second.
type Cmp[T any] func(T, T) int

// Function type for use with [Create]. To stop the creation of more values, the bool return value must be true.
type CreateFunc[T any] func() (T, bool)

// New returns a new [iter.Seq] from the given values.
func New[T any](vs ...T) iter.Seq[T] {
	return func(yield func(T) bool) {
		for _, v := range vs {
			if !yield(v) {
				return
			}
		}
	}
}

// Create returns a new [iter.Seq] with the values created bei the given [CreateFunc] function.
//
// Example: create an endless sequence with random ints.
//
//	func RandomInt(n int) iter.Seq[int] {
//	  if n == 0 {
//	    return Create(func() (int, bool) { return rand.Int(), false })
//	  }
//	  return Create(func() (int, bool) { return rand.IntN(n), false })
//	}
func Create[T any](fn CreateFunc[T]) iter.Seq[T] {
	return func(yield func(T) bool) {
		for {
			v, stop := fn()
			if stop || !yield(v) {
				return
			}
		}
	}
}

// Cycle returns an [iter.Seq] that cycles through the given values.
func Cycle[T any](v T, vs ...T) iter.Seq[T] {
	return func(yield func(T) bool) {
		sl := []T{v}
		sl = append(sl, vs...)
		i := 0
		l := len(sl)
		for {
			if !yield(sl[i%l]) {
				return
			}
			i++
		}
	}
}

// Repeat returns an [iter.Seq] that returns v n times (n == 0 -> no limit).
func Repeat[T any](v T, n int) iter.Seq[T] {
	if n == 0 {
		return func(yield func(T) bool) {
			for {
				if !yield(v) {
					return
				}
			}
		}
	}

	return func(yield func(T) bool) {
		for range n {
			if !yield(v) {
				return
			}
		}
	}
}

// Count returns an [iter.Seq] that counts from the start value upward.
func Count[T constraints.Integer](start T) iter.Seq[T] {
	return func(yield func(T) bool) {
		n := start
		for {
			if !yield(n) {
				return
			}
			n++
		}
	}
}

// Equal returns true if both sequences contain the same values in the same order.
func Equal[T any](seq1, seq2 iter.Seq[T], eq Eq[T]) bool {
	next, stop := iter.Pull(seq2)
	defer stop()
	for v1 := range seq1 {
		v2, ok := next()
		if !ok {
			return false
		}
		if !eq(v1, v2) {
			return false
		}
	}
	_, ok := next()
	return !ok
}

// Contains returns wether the value is contained in the sequence.
func Contains[T any](seq iter.Seq[T], value T, eq Eq[T]) bool {
	for v := range seq {
		if eq(v, value) {
			return true
		}
	}
	return false
}

// Find returns an [iter.Seq2] with all found values and their position in the sequence.
func Find[T any](seq iter.Seq[T], value T, eq Eq[T]) iter.Seq2[int, T] {
	return func(yield func(int, T) bool) {
		i := 0
		for v := range seq {
			if eq(v, value) && !yield(i, v) {
				return
			}
			i++
		}
	}
}

// Enumerate returns an [iter.Seq2] with all values and their position in the sequence.
func Enumerate[T any](seq iter.Seq[T], start int) iter.Seq2[int, T] {
	return func(yield func(int, T) bool) {
		i := start
		for v := range seq {
			if !yield(i, v) {
				return
			}
			i++
		}
	}
}