xiter: optimize Zip()

transform.go

@@ -112,16 +112,23 @@ type Zipped[T1, T2 any] struct {
 // Must not be a method. See https://github.com/golang/go/issues/80172.
 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 := iter.Pull(seq1.Seq())
-		defer stop()
-		p2, stop := iter.Pull(seq2.Seq())
-		defer stop()
+		next2, stop2 := iter.Pull(seq2.Seq())
+		defer stop2()
+
+		for v1 := range seq1 {
+			v2, ok2 := next2()
+			val := Zipped[T1, T2]{
+				V1: v1, OK1: true,
+				V2: v2, OK2: ok2,
+			}
+			if !yield(val) {
+				return
+			}
+		}
 
-		for {
-			var val Zipped[T1, T2]
-			val.V1, val.OK1 = p1()
-			val.V2, val.OK2 = p2()
-			if (!val.OK1 && !val.OK2) || !yield(val) {
+		for v2, ok2 := next2(); ok2; v2, ok2 = next2() {
+			val := Zipped[T1, T2]{V2: v2, OK2: ok2}
+			if !yield(val) {
 				return
 			}
 		}

transform_test.go

@@ -50,27 +50,133 @@ func TestConcat(t *testing.T) {
 }
 
 func TestZip(t *testing.T) {
-	s1 := Of(1, 2, 3, 4, 5)
-	s2 := Of(2, 3, 4, 5, 6)
-	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
-	})
+	tests := []struct {
+		name  string
+		s1    []int
+		s2    []int
+		limit int
+		want  []Zipped[int, int]
+	}{
+		{
+			name: "equal",
+			s1:   []int{1, 2, 3, 4, 5},
+			s2:   []int{2, 3, 4, 5, 6},
+			want: []Zipped[int, int]{
+				{V1: 1, OK1: true, V2: 2, OK2: true},
+				{V1: 2, OK1: true, V2: 3, OK2: true},
+				{V1: 3, OK1: true, V2: 4, OK2: true},
+				{V1: 4, OK1: true, V2: 5, OK2: true},
+				{V1: 5, OK1: true, V2: 6, OK2: true},
+			},
+		},
+		{
+			name: "seq1 shorter",
+			s1:   []int{1, 2, 3, 4},
+			s2:   []int{2, 3, 4, 5, 10},
+			want: []Zipped[int, int]{
+				{V1: 1, OK1: true, V2: 2, OK2: true},
+				{V1: 2, OK1: true, V2: 3, OK2: true},
+				{V1: 3, OK1: true, V2: 4, OK2: true},
+				{V1: 4, OK1: true, V2: 5, OK2: true},
+				{V2: 10, OK2: true},
+			},
+		},
+		{
+			name: "seq2 shorter",
+			s1:   []int{1, 2, 3, 4, 10},
+			s2:   []int{2, 3, 4, 5},
+			want: []Zipped[int, int]{
+				{V1: 1, OK1: true, V2: 2, OK2: true},
+				{V1: 2, OK1: true, V2: 3, OK2: true},
+				{V1: 3, OK1: true, V2: 4, OK2: true},
+				{V1: 4, OK1: true, V2: 5, OK2: true},
+				{V1: 10, OK1: true},
+			},
+		},
+		{
+			name: "empty seq1",
+			s2:   []int{1, 2, 3},
+			want: []Zipped[int, int]{
+				{V2: 1, OK2: true},
+				{V2: 2, OK2: true},
+				{V2: 3, OK2: true},
+			},
+		},
+		{
+			name: "empty seq2",
+			s1:   []int{1, 2, 3},
+			want: []Zipped[int, int]{
+				{V1: 1, OK1: true},
+				{V1: 2, OK1: true},
+				{V1: 3, OK1: true},
+			},
+		},
+		{
+			name:  "stop with seq2 remaining",
+			s1:    []int{1, 2},
+			s2:    []int{10, 20, 30, 40},
+			limit: 1,
+			want: []Zipped[int, int]{
+				{V1: 1, OK1: true, V2: 10, OK2: true},
+			},
+		},
+		{
+			name:  "stop leftover",
+			s1:    []int{1},
+			s2:    []int{10, 20, 30},
+			limit: 2,
+			want: []Zipped[int, int]{
+				{V1: 1, OK1: true, V2: 10, OK2: true},
+				{V2: 20, OK2: true},
+			},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			seq := Zip(Of(tt.s1...), Of(tt.s2...))
+			if tt.limit > 0 {
+				seq = seq.Limit(tt.limit)
+			}
+			got := seq.Collect()
+			if !slices.Equal(got, tt.want) {
+				t.Fatal(got)
+			}
+		})
+	}
 }
 
 func BenchmarkZip(b *testing.B) {
-	slice1 := []int{1, 2, 3, 4, 5}
-	slice2 := []int{2, 3, 4, 5, 6}
-
-	for b.Loop() {
-		s1 := S(slices.Values(slice1))
-		s2 := S(slices.Values(slice2))
-		seq := Zip(s1, s2)
-		seq(func(v Zipped[int, int]) bool {
-			return true
+	benchmarkZip(b, Zip)
+}
+
+func benchmarkZip(b *testing.B, zip func(Seq[int], Seq[int]) Seq[Zipped[int, int]]) {
+	shapes := []struct {
+		name   string
+		n1, n2 int
+	}{
+		{"n5_equal", 5, 5},
+		{"n1000_equal", 1000, 1000},
+		{"seq1_short", 10, 1000},
+	}
+
+	for _, shape := range shapes {
+		b.Run(shape.name, func(b *testing.B) {
+			s1 := make([]int, shape.n1)
+			s2 := make([]int, shape.n2)
+			for i := range s1 {
+				s1[i] = i
+			}
+			for i := range s2 {
+				s2[i] = i + 1
+			}
+
+			for b.Loop() {
+				seq := zip(S(slices.Values(s1)), S(slices.Values(s2)))
+				seq(func(v Zipped[int, int]) bool {
+					return true
+				})
+			}
 		})
 	}
 }

zip2pull_test.go

@@ -0,0 +1,30 @@
+package xiter
+
+import (
+	"iter"
+	"testing"
+)
+
+func BenchmarkZip2Pull(b *testing.B) {
+	benchmarkZip(b, zip2Pull)
+}
+
+// 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 := iter.Pull(seq1.Seq())
+		defer stop()
+		p2, stop := iter.Pull(seq2.Seq())
+		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
+			}
+		}
+	}
+}