diff --git a/transform.go b/transform.go index d6ca43a..dcf9976 100644 --- a/transform.go +++ b/transform.go @@ -101,19 +101,35 @@ type Zipped[T1, T2 any] struct { // simultaneously. func Zip[T1, T2 any](seq1 Seq[T1], seq2 Seq[T2]) Seq[Zipped[T1, T2]] { return func(yield func(Zipped[T1, T2]) bool) { - p1, stop := Pull(seq1) - defer stop() - p2, stop := Pull(seq2) - defer stop() + p2, stop2 := Pull(seq2) + defer stop2() + done := false - for { + f := func(v T1) bool { var val Zipped[T1, T2] - val.V1, val.OK1 = p1() + val.V1, val.OK1 = v, true val.V2, val.OK2 = p2() - if (!val.OK1 && !val.OK2) || !yield(val) { + if !yield(val) { + done = true + return false + } + return true + } + seq1(f) + if done { + return + } + // seq1 is exhausted + for v2, ok2 := p2(); ok2; v2, ok2 = p2() { + var v1 T1 + var val Zipped[T1, T2] + val.V1, val.OK1 = v1, false + val.V2, val.OK2 = v2, true + if !yield(val) { return } } + return } } diff --git a/transform_test.go b/transform_test.go index 71dd676..9b85459 100644 --- a/transform_test.go +++ b/transform_test.go @@ -63,6 +63,35 @@ func TestZip(t *testing.T) { }) } +func TestZipShort1(t *testing.T) { + s1 := OfSlice([]int{1, 2, 3, 4}) + s2 := OfSlice([]int{2, 3, 4, 5, 1}) + seq := Zip(s1, s2) + seq(func(v Zipped[int, int]) bool { + if v.V2-v.V1 != 1 { + t.Fatalf("unexpected values: %+v", v) + } + return true + }) + t.Logf("Greetings from TestZipShort1") +} + +func TestZipShort2(t *testing.T) { + s1 := OfSlice([]int{1, 2, 3, 4, -1, -2}) + s2 := OfSlice([]int{2, 3, 4, 5}) + seq := Zip(s1, s2) + seq(func(v Zipped[int, int]) bool { + if v.V1 == -2 { + return false + } + if v.V2-v.V1 != 1 { + t.Fatalf("unexpected values: %+v", v) + } + return true + }) + t.Logf("Greetings from TestZipShort2") +} + func BenchmarkZip(b *testing.B) { slice1 := []int{1, 2, 3, 4, 5} slice2 := []int{2, 3, 4, 5, 6} diff --git a/zip2pull_test.go b/zip2pull_test.go new file mode 100644 index 0000000..576baf9 --- /dev/null +++ b/zip2pull_test.go @@ -0,0 +1,38 @@ +package xiter + +import "testing" + +func BenchmarkZip2Pull(b *testing.B) { + b.ReportAllocs() + slice1 := []int{1, 2, 3, 4, 5} + slice2 := []int{2, 3, 4, 5, 6} + + for i := 0; i < b.N; i++ { + s1 := OfSlice(slice1) + s2 := OfSlice(slice2) + seq := zip2Pull(s1, s2) + seq(func(v Zipped[int, int]) bool { + return true + }) + } +} + +// zip2Pull returns a new Seq that yields the values of seq1 and seq2 +// simultaneously. This is the straightforward 2-Pull version. +func zip2Pull[T1, T2 any](seq1 Seq[T1], seq2 Seq[T2]) Seq[Zipped[T1, T2]] { + return func(yield func(Zipped[T1, T2]) bool) { + p1, stop := Pull(seq1) + defer stop() + p2, stop := Pull(seq2) + defer stop() + + for { + var val Zipped[T1, T2] + val.V1, val.OK1 = p1() + val.V2, val.OK2 = p2() + if (!val.OK1 && !val.OK2) || !yield(val) { + return + } + } + } +}