From 90bf3282a31cbc276b1b8ac94af37f0b114d0ed9 Mon Sep 17 00:00:00 2001
From: Maxim Voloshin <pyrozhog@gmail.com>
Date: Thu, 8 Jun 2023 16:23:58 +0300
Subject: [PATCH] Add Zip and ZipLongest function (#193)

Zip will return a new slice containing pairs with elements from input slices.
If input slices have different length, the output slice will be truncated to
the length of the smallest input slice.

ZipLongest will return a new slice containing pairs with elements from input slices.
If input slices have diffrent length, missing elements will be padded with default values.

These are the same as zip() and itertools.zip_longest() in Python.
---
 v2/zip.go              | 27 ++++++++++++++++++
 v2/zip_longest.go      | 37 +++++++++++++++++++++++++
 v2/zip_longest_test.go | 18 ++++++++++++
 v2/zip_test.go         | 62 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 144 insertions(+)
 create mode 100644 v2/zip.go
 create mode 100644 v2/zip_longest.go
 create mode 100644 v2/zip_longest_test.go
 create mode 100644 v2/zip_test.go

diff --git a/v2/zip.go b/v2/zip.go
new file mode 100644
index 0000000..da1bbd6
--- /dev/null
+++ b/v2/zip.go
@@ -0,0 +1,27 @@
+package pie
+
+// A pair struct containing two zipped values.
+type Zipped[T1, T2 any] struct {
+	A T1
+	B T2
+}
+
+// Zip will return a new slice containing pairs with elements from input slices.
+// If input slices have diffrent length, the output slice will be truncated to
+// the length of the smallest input slice.
+func Zip[T1, T2 any](ss1 []T1, ss2 []T2) []Zipped[T1, T2] {
+	var minLen int
+
+	if len(ss1) <= len(ss2) {
+		minLen = len(ss1)
+	} else {
+		minLen = len(ss2)
+	}
+
+	ss3 := []Zipped[T1, T2]{}
+	for i := 0; i < minLen; i++ {
+		ss3 = append(ss3, Zipped[T1, T2]{ss1[i], ss2[i]})
+	}
+
+	return ss3
+}
diff --git a/v2/zip_longest.go b/v2/zip_longest.go
new file mode 100644
index 0000000..8dc9004
--- /dev/null
+++ b/v2/zip_longest.go
@@ -0,0 +1,37 @@
+package pie
+
+// ZipLongest will return a new slice containing pairs with elements from input slices.
+// If input slices have diffrent length, missing elements will be padded with default values.
+func ZipLongest[T1, T2 any](ss1 []T1, ss2 []T2) []Zipped[T1, T2] {
+	var minLen, maxLen int
+	var small int8
+
+	if len(ss1) <= len(ss2) {
+		small = 1
+		minLen = len(ss1)
+		maxLen = len(ss2)
+	} else {
+		small = 2
+		minLen = len(ss2)
+		maxLen = len(ss1)
+	}
+
+	ss3 := []Zipped[T1, T2]{}
+	for i := 0; i < minLen; i++ {
+		ss3 = append(ss3, Zipped[T1, T2]{ss1[i], ss2[i]})
+	}
+
+	if small == 1 {
+		var t T1
+		for i := minLen; i < maxLen; i++ {
+			ss3 = append(ss3, Zipped[T1, T2]{t, ss2[i]})
+		}
+	} else {
+		var t T2
+		for i := minLen; i < maxLen; i++ {
+			ss3 = append(ss3, Zipped[T1, T2]{ss1[i], t})
+		}
+	}
+
+	return ss3
+}
diff --git a/v2/zip_longest_test.go b/v2/zip_longest_test.go
new file mode 100644
index 0000000..24c966b
--- /dev/null
+++ b/v2/zip_longest_test.go
@@ -0,0 +1,18 @@
+package pie_test
+
+import (
+	"github.com/elliotchance/pie/v2"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+func TestZipLongest(t *testing.T) {
+	for _, test := range zipTests {
+
+		t.Run("", func(t *testing.T) {
+			c := pie.ZipLongest(test.ss1, test.ss2)
+
+			assert.Equal(t, c, test.expectedLong)
+		})
+	}
+}
diff --git a/v2/zip_test.go b/v2/zip_test.go
new file mode 100644
index 0000000..95d09f8
--- /dev/null
+++ b/v2/zip_test.go
@@ -0,0 +1,62 @@
+package pie_test
+
+import (
+	"github.com/elliotchance/pie/v2"
+	"github.com/stretchr/testify/assert"
+	"testing"
+)
+
+var zipTests = []struct {
+	ss1           []int
+	ss2           []float32
+	expectedShort []pie.Zipped[int, float32]
+	expectedLong  []pie.Zipped[int, float32]
+}{
+	{
+		[]int{},
+		[]float32{},
+		[]pie.Zipped[int, float32]{},
+		[]pie.Zipped[int, float32]{},
+	},
+	{
+		[]int{1, 2, 3, 4, 5},
+		[]float32{},
+		[]pie.Zipped[int, float32]{},
+		[]pie.Zipped[int, float32]{{1, 0}, {2, 0}, {3, 0}, {4, 0}, {5, 0}},
+	},
+	{
+		[]int{},
+		[]float32{1.0, 2.0, 3.0, 4.0, 5.0},
+		[]pie.Zipped[int, float32]{},
+		[]pie.Zipped[int, float32]{{0, 1.0}, {0, 2.0}, {0, 3.0}, {0, 4.0}, {0, 5.0}},
+	},
+	{
+		[]int{1, 2, 3, 4, 5},
+		[]float32{1.0, 2.0, 3.0, 4.0, 5.0},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {4, 4.0}, {5, 5.0}},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {4, 4.0}, {5, 5.0}},
+	},
+	{
+		[]int{1, 2, 3, 4, 5, 6, 7, 8},
+		[]float32{1.0, 2.0, 3.0, 4.0, 5.0},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {4, 4.0}, {5, 5.0}},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {4, 4.0}, {5, 5.0}, {6, 0}, {7, 0}, {8, 0}},
+	},
+	{
+		[]int{1, 2, 3},
+		[]float32{1.0, 2.0, 3.0, 4.0, 5.0},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}},
+		[]pie.Zipped[int, float32]{{1, 1.0}, {2, 2.0}, {3, 3.0}, {0, 4.0}, {0, 5.0}},
+	},
+}
+
+func TestZip(t *testing.T) {
+	for _, test := range zipTests {
+
+		t.Run("", func(t *testing.T) {
+			c := pie.Zip(test.ss1, test.ss2)
+
+			assert.Equal(t, c, test.expectedShort)
+		})
+	}
+}