From 689644acda08e7a5bb393ddedf16643d24553193 Mon Sep 17 00:00:00 2001
From: myfreess <neomyfreess@outlook.com>
Date: Tue, 18 Feb 2025 15:55:42 +0800
Subject: [PATCH] add @array.zip and @array.zip_with

---
 array/array.mbt      | 67 ++++++++++++++++++++++++++++++++++++++++++++
 array/array.mbti     |  2 ++
 array/array_test.mbt | 42 +++++++++++++++++++++++++++
 3 files changed, 111 insertions(+)

diff --git a/array/array.mbt b/array/array.mbt
index 958d255a6..f7998e68c 100644
--- a/array/array.mbt
+++ b/array/array.mbt
@@ -202,6 +202,73 @@ pub fn last[A](self : Array[A]) -> A? {
   }
 }
 
+///| Zips two arrays into a single array of tuples.
+///
+/// Parameters:
+///
+/// * `self` : The first array.
+/// * `other` : 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!(arr1.zip(arr2), content="[(1, a), (2, b), (3, c)]")
+/// }
+/// ```
+pub fn zip[A, B](self : Array[A], other : Array[B]) -> Array[(A, B)] {
+  let length = if self.length() < other.length() {
+    self.length()
+  } else {
+    other.length()
+  }
+  let arr = Array::new(capacity=length)
+  for i = 0; i < length; i = i + 1 {
+    arr.push((self[i], other[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 7ccd849a8..b3287a60a 100644
--- a/array/array.mbti
+++ b/array/array.mbti
@@ -4,6 +4,7 @@ alias @moonbitlang/core/bytes as @bytes
 alias @moonbitlang/core/quickcheck as @quickcheck
 
 // Values
+fn zip_with[A, B, C](Array[A], Array[B], merge~ : (A, B) -> C) -> Array[C]
 
 // Types and methods
 impl FixedArray {
@@ -62,6 +63,7 @@ impl Array {
   sort[T : Compare](Self[T]) -> Unit
   sort_by[T](Self[T], (T, T) -> Int) -> Unit
   sort_by_key[T, K : Compare](Self[T], (T) -> K) -> Unit
+  zip[A, B](Self[A], Self[B]) -> Self[(A, B)]
 }
 impl[X : @quickcheck.Arbitrary] @quickcheck.Arbitrary for Array[X]
 
diff --git a/array/array_test.mbt b/array/array_test.mbt
index bfca078e5..35d7ac72b 100644
--- a/array/array_test.mbt
+++ b/array/array_test.mbt
@@ -618,6 +618,48 @@ 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!(arr1.zip(arr2), content="[(1, 'a'), (2, 'b'), (3, 'c')]")
+
+  // Test with arrays of different lengths
+  let arr3 = [1, 2]
+  let arr4 = ["a", "b", "c"]
+  inspect!(
+    arr3.zip(arr4),
+    content=
+      #|[(1, "a"), (2, "b")]
+    ,
+  )
+
+  // Test with an empty array
+  let arr5 : Array[Int] = []
+  let arr6 = ["a", "b", "c"]
+  inspect!(arr5.zip(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)