From 01aafb4cee88e6abcd982dc6ca6b0aa184cb15a1 Mon Sep 17 00:00:00 2001
From: myfreess <neomyfreess@outlook.com>
Date: Sun, 26 Jan 2025 11:54:58 +0800
Subject: [PATCH] add @array.zip and @array.zip_with

---
 array/array.mbt      | 63 ++++++++++++++++++++++++++++++++++++++++++++
 array/array.mbti     |  3 +++
 array/array_test.mbt | 40 ++++++++++++++++++++++++++++
 3 files changed, 106 insertions(+)

diff --git a/array/array.mbt b/array/array.mbt
index accb04b9d..d8de30ec2 100644
--- a/array/array.mbt
+++ b/array/array.mbt
@@ -202,6 +202,69 @@ pub fn last[A](self : Array[A]) -> A? {
   }
 }
 
+///| Zips two arrays into a single array of tuples.
+///
+/// Parameters:
+///
+/// * `l` : The first array.
+/// * `r` : The second array.
+///
+/// Returns an array of tuples, where each tuple contains corresponding elements
+/// from the two input arrays.
+///
+/// Example:
+///
+/// ```moonbit
+/// test "zip" {
+///   let arr1 = [1, 2, 3]
+///   let arr2 = ['a', 'b', 'c']
+///   inspect!(zip(arr1, arr2), content="[(1, 'a'), (2, 'b'), (3, 'c')]")
+/// }
+/// ```
+pub fn zip[A, B](l : Array[A], r : Array[B]) -> Array[(A, B)] {
+  let length = if l.length() < r.length() { l.length() } else { r.length() }
+  let arr = Array::new(capacity=length)
+  for i = 0; i < length; i = i + 1 {
+    arr.push((l[i], r[i]))
+  } else {
+    return arr
+  }
+}
+
+///| Zips two arrays into a single array by applying a function to each pair of elements.
+///
+/// Parameters:
+///
+/// * `l` : The first array.
+/// * `r` : The second array.
+/// * `merge` : A function that takes two arguments, one from each array, and returns a value.
+///
+/// Returns an array containing the results of applying the function to each pair of elements.
+///
+/// Example:
+///
+/// ```moonbit
+/// test "zip_with" {
+///   let arr1 = [1, 2, 3]
+///   let arr2 = [4, 5, 6]
+///   let add = fn(a, b) { a + b }
+///   inspect!(zip_with(arr1, arr2, merge=add), content="[5, 7, 9]")
+/// }
+/// ```
+pub fn zip_with[A, B, C](
+  l : Array[A],
+  r : Array[B],
+  merge~ : (A, B) -> C
+) -> Array[C] {
+  let length = if l.length() < r.length() { l.length() } else { r.length() }
+  let arr = Array::new(capacity=length)
+  for i = 0; i < length; i = i + 1 {
+    arr.push(merge(l[i], r[i]))
+  } else {
+    return arr
+  }
+}
+
 ///|
 pub impl[X : @quickcheck.Arbitrary] @quickcheck.Arbitrary for Array[X] with arbitrary(
   size,
diff --git a/array/array.mbti b/array/array.mbti
index 80a142f60..e5100400f 100644
--- a/array/array.mbti
+++ b/array/array.mbti
@@ -3,6 +3,9 @@ package moonbitlang/core/array
 alias @moonbitlang/core/quickcheck as @quickcheck
 
 // Values
+fn zip[A, B](Array[A], Array[B]) -> Array[(A, B)]
+
+fn zip_with[A, B, C](Array[A], Array[B], merge~ : (A, B) -> C) -> Array[C]
 
 // Types and methods
 impl FixedArray {
diff --git a/array/array_test.mbt b/array/array_test.mbt
index 33d730e23..5ba56ca92 100644
--- a/array/array_test.mbt
+++ b/array/array_test.mbt
@@ -555,6 +555,46 @@ test "Array::last" {
   inspect!([1].last(), content="Some(1)")
 }
 
+test "zip" {
+  // Test with two non-empty arrays
+  let arr1 = [1, 2, 3]
+  let arr2 = ['a', 'b', 'c']
+  inspect!(zip(arr1, arr2), content="[(1, 'a'), (2, 'b'), (3, 'c')]")
+
+  // Test with arrays of different lengths
+  let arr3 = [1, 2]
+  let arr4 = ["a", "b", "c"]
+  inspect!(
+    zip(arr3, arr4),
+    content=
+      #|[(1, "a"), (2, "b")]
+    ,
+  )
+
+  // Test with an empty array
+  let arr5 : Array[Int] = []
+  let arr6 = ["a", "b", "c"]
+  inspect!(zip(arr5, arr6), content="[]")
+}
+
+test "zip_with" {
+  // Test with two non-empty arrays and a function
+  let arr1 = [1, 2, 3]
+  let arr2 = [4, 5, 6]
+  let add = fn(a, b) { a + b }
+  inspect!(zip_with(arr1, arr2, merge=add), content="[5, 7, 9]")
+
+  // Test with arrays of different lengths and a function
+  let arr3 = [1, 2]
+  let arr4 = [4, 5, 6]
+  inspect!(zip_with(arr3, arr4, merge=add), content="[5, 7]")
+
+  // Test with an empty array and a function
+  let arr5 : Array[Int] = []
+  let arr6 = [4, 5, 6]
+  inspect!(zip_with(arr5, arr6, merge=add), content="[]")
+}
+
 test "arbitrary" {
   let arr : Array[Array[Int]] = @quickcheck.samples(20)
   inspect!(arr[5:9], content="[[0], [0], [0], [0, 0, 0]]")