From c5a5e654d4ac65ec57dd66a1200095bac6232951 Mon Sep 17 00:00:00 2001
From: Joseph Fourny <joseph.fourny@shopify.com>
Date: Thu, 29 Feb 2024 21:54:58 -0500
Subject: [PATCH] Add `UnzipFirst` and `UnzipSecond` stream functions.

---
 pkg/stream/combine.go      | 34 ++++++++++++++++++++++++++++++++++
 pkg/stream/combine_test.go | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 72 insertions(+)

diff --git a/pkg/stream/combine.go b/pkg/stream/combine.go
index 7da35bf..26680a1 100644
--- a/pkg/stream/combine.go
+++ b/pkg/stream/combine.go
@@ -120,3 +120,37 @@ func Zip[E, F any](s1 Stream[E], s2 Stream[F]) Stream[pair.Pair[E, F]] {
 func ZipWithIndex[E any, I constraint.Integer](s Stream[E], offset I) Stream[pair.Pair[E, I]] {
 	return Zip(s, Walk(offset, pred.True[I](), mapper.Increment[I](1)))
 }
+
+// UnzipFirst returns a stream that contains the first elements of each pair in the input stream.
+//
+// Example usage:
+//
+//	s := stream.UnzipFirst(
+//	  stream.Of(
+//	    pair.Of(1, "foo"),
+//	    pair.Of(2, "bar"),
+//	  ),
+//	)
+//	out := stream.DebugString(s) // "<1, 2>"
+func UnzipFirst[E, F any](s Stream[pair.Pair[E, F]]) Stream[E] {
+	return Map(s, func(p pair.Pair[E, F]) E {
+		return p.First()
+	})
+}
+
+// UnzipSecond returns a stream that contains the second elements of each pair in the input stream.
+//
+// Example usage:
+//
+//	s := stream.UnzipSecond(
+//	  stream.Of(
+//	    pair.Of(1, "foo"),
+//	    pair.Of(2, "bar"),
+//	  ),
+//	)
+//	out := stream.DebugString(s) // "<foo, bar>"
+func UnzipSecond[E, F any](s Stream[pair.Pair[E, F]]) Stream[F] {
+	return Map(s, func(p pair.Pair[E, F]) F {
+		return p.Second()
+	})
+}
diff --git a/pkg/stream/combine_test.go b/pkg/stream/combine_test.go
index 9131dbc..f2d2f63 100644
--- a/pkg/stream/combine_test.go
+++ b/pkg/stream/combine_test.go
@@ -153,3 +153,41 @@ func TestZipWithIndex(t *testing.T) {
 		assert.ElementsMatch(t, got, want)
 	})
 }
+
+func TestUnzipFirst(t *testing.T) {
+	t.Run("empty", func(t *testing.T) {
+		s := UnzipFirst(Empty[pair.Pair[int, string]]())
+		got := CollectSlice(s)
+		var want []int
+		assert.ElementsMatch(t, got, want)
+	})
+
+	t.Run("non-empty", func(t *testing.T) {
+		s := UnzipFirst(Of(
+			pair.Of(1, "foo"),
+			pair.Of(2, "bar"),
+		))
+		got := CollectSlice(s)
+		want := []int{1, 2}
+		assert.ElementsMatch(t, got, want)
+	})
+}
+
+func TestUnzipSecond(t *testing.T) {
+	t.Run("empty", func(t *testing.T) {
+		s := UnzipSecond(Empty[pair.Pair[int, string]]())
+		got := CollectSlice(s)
+		var want []string
+		assert.ElementsMatch(t, got, want)
+	})
+
+	t.Run("non-empty", func(t *testing.T) {
+		s := UnzipSecond(Of(
+			pair.Of(1, "foo"),
+			pair.Of(2, "bar"),
+		))
+		got := CollectSlice(s)
+		want := []string{"foo", "bar"}
+		assert.ElementsMatch(t, got, want)
+	})
+}