diff --git a/transform.go b/transform.go index 94ab977..3131c07 100644 --- a/transform.go +++ b/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 } } diff --git a/transform_test.go b/transform_test.go index c1c285c..104bb2e 100644 --- a/transform_test.go +++ b/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 + }) + } }) } } diff --git a/zip2pull_test.go b/zip2pull_test.go new file mode 100644 index 0000000..217346b --- /dev/null +++ b/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 + } + } + } +}