seq.go

package iterutils

import "iter"

// First returns an [iter.Seq] with the first values of the pairs in an [iter.Seq2].
func First[U, V any](seq iter.Seq2[U, V]) iter.Seq[U] {
	return func(yield func(U) bool) {
		for u := range seq {
			if !yield(u) {
				return
			}
		}
	}
}

// Second returns an [iter.Seq] with the second values of the pairs in an [iter.Seq2].
func Second[U, V any](seq iter.Seq2[U, V]) iter.Seq[V] {
	return func(yield func(V) bool) {
		for _, v := range seq {
			if !yield(v) {
				return
			}
		}
	}
}

// Split returns two [iter.Seq]s (see [First] and [Second]).
func Split[U, V any](seq iter.Seq2[U, V]) (iter.Seq[U], iter.Seq[V]) {
	return First(seq), Second(seq)
}

// Merge returns an [iter.Seq2] combining two [iter.Seq]s. It stops when the shortest sequence is exhausted.
func Merge[U, V any](seq1 iter.Seq[U], seq2 iter.Seq[V]) iter.Seq2[U, V] {
	return func(yield func(U, V) bool) {
		next, stop := iter.Pull(seq2)
		defer stop()
		for u := range seq1 {
			if v, ok := next(); !ok || !yield(u, v) {
				return
			}
		}
	}
}

// Chain returns an [iter.Seq] that combines the given sequences.
func Chain[T any](seqs ...iter.Seq[T]) iter.Seq[T] {
	return func(yield func(T) bool) {
		for _, seq := range seqs {
			for v := range seq {
				if !yield(v) {
					return
				}
			}
		}
	}
}