optimize Zip to only use one Pull, and test for coverage

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
 	}
 }
 

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}

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
+			}
+		}
+	}
+}