zip.go

package iterutils

import "iter"

type Tuple[U, V any] struct {
	A U
	B V
}

type Tuple3[U, V, W any] struct {
	A U
	B V
	C W
}

// Zip zips two sequences into one. It stops when the shortest sequence is exhausted.
func Zip[U, V any](seq1 iter.Seq[U], seq2 iter.Seq[V]) iter.Seq[Tuple[U, V]] {
	return func(yield func(Tuple[U, V]) bool) {
		next, stop := iter.Pull(seq2)
		defer stop()
		for u := range seq1 {
			v, ok := next()
			if !ok || !yield(Tuple[U, V]{u, v}) {
				return
			}
		}
	}
}

// Zip3 zips three sequences into one. It stops when the shortest sequence is exhausted.
func Zip3[U, V, W any](seq1 iter.Seq[U], seq2 iter.Seq[V], seq3 iter.Seq[W]) iter.Seq[Tuple3[U, V, W]] {
	return func(yield func(Tuple3[U, V, W]) bool) {
		next2, stop2 := iter.Pull(seq2)
		defer stop2()
		next3, stop3 := iter.Pull(seq3)
		defer stop3()
		for u := range seq1 {
			v, ok2 := next2()
			w, ok3 := next3()
			if !ok2 || !ok3 || !yield(Tuple3[U, V, W]{u, v, w}) {
				return
			}
		}
	}
}

// Unzip unzips a sequence into two.
func Unzip[U, V any](seq iter.Seq[Tuple[U, V]]) (iter.Seq[U], iter.Seq[V]) {
	seq1 := func(yield func(U) bool) {
		for t := range seq {
			if !yield(t.A) {
				return
			}
		}
	}
	seq2 := func(yield func(V) bool) {
		for t := range seq {
			if !yield(t.B) {
				return
			}
		}
	}
	return seq1, seq2
}

// Unzip3 unzips a sequence into three.
func Unzip3[U, V, W any](seq iter.Seq[Tuple3[U, V, W]]) (iter.Seq[U], iter.Seq[V], iter.Seq[W]) {
	seq1 := func(yield func(U) bool) {
		for t := range seq {
			if !yield(t.A) {
				return
			}
		}
	}
	seq2 := func(yield func(V) bool) {
		for t := range seq {
			if !yield(t.B) {
				return
			}
		}
	}
	seq3 := func(yield func(W) bool) {
		for t := range seq {
			if !yield(t.C) {
				return
			}
		}
	}
	return seq1, seq2, seq3
}