add `iterator.map2`

3 files changed, 79 insertions, 0 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8e52b72..a90d16a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -2,6 +2,7 @@
 
 ## v0.31.0 - 2023-09-25
 
+- The `iterator` module gains the `map2` function.
 - `list.flatten` is no longer deprecated and is kept as a synonym of
   `list.concat`
 - The `iterator` module gains the `concat` function.
diff --git a/src/gleam/iterator.gleam b/src/gleam/iterator.gleam
index 5865646..c13a10e 100644
--- a/src/gleam/iterator.gleam
+++ b/src/gleam/iterator.gleam
@@ -386,6 +386,56 @@ pub fn map(over iterator: Iterator(a), with f: fn(a) -> b) -> Iterator(b) {
   |> Iterator
 }
 
+fn do_map2(
+  continuation1: fn() -> Action(a),
+  continuation2: fn() -> Action(b),
+  with fun: fn(a, b) -> c,
+) -> fn() -> Action(c) {
+  fn() {
+    case continuation1() {
+      Stop -> Stop
+      Continue(a, next_a) ->
+        case continuation2() {
+          Stop -> Stop
+          Continue(b, next_b) ->
+            Continue(fun(a, b), do_map2(next_a, next_b, fun))
+        }
+    }
+  }
+}
+
+/// Combines two interators into a single one using the given function.
+/// 
+/// If an iterator is longer than the other the extra elements are dropped.
+/// 
+/// This function does not evaluate the elements of the two iterators, the
+/// computation is performed when the resulting iterator is later run.
+/// 
+/// ## Examples
+/// 
+/// ```gleam
+/// let first = from_list([1, 2, 3])
+/// let second = from_list([4, 5, 6])
+/// map2(first, second, fn(x, y) { x + y }) |> to_list
+/// // -> [5, 7, 9]
+/// ```
+/// 
+/// ```gleam
+/// let first = from_list([1, 2])
+/// let second = from_list(["a", "b", "c"])
+/// map2(first, second, fn(i, x) { #(i, x) }) |> to_list
+/// // -> [#(1, "a"), #(2, "b")]
+/// ```
+/// 
+pub fn map2(
+  iterator1: Iterator(a),
+  iterator2: Iterator(b),
+  with fun: fn(a, b) -> c,
+) -> Iterator(c) {
+  do_map2(iterator1.continuation, iterator2.continuation, fun)
+  |> Iterator
+}
+
 fn do_append(first: fn() -> Action(a), second: fn() -> Action(a)) -> Action(a) {
   case first() {
     Continue(e, first) -> Continue(e, fn() { do_append(first, second) })
diff --git a/test/gleam/iterator_test.gleam b/test/gleam/iterator_test.gleam
index 66e8bd4..adb4802 100644
--- a/test/gleam/iterator_test.gleam
+++ b/test/gleam/iterator_test.gleam
@@ -155,6 +155,34 @@ pub fn map_test() {
   test([1, 2, 3, 4, 5, 6, 7, 8], f)
 }
 
+// map2(from_list(a), from_list(b), f)  == list.map2(a, b, f)
+pub fn map2_test() {
+  let test = fn(one, other, f) {
+    iterator.map2(iterator.from_list(one), iterator.from_list(other), f)
+    |> iterator.to_list
+    |> should.equal(list.map2(one, other, f))
+  }
+
+  let f = fn(a, b) { a / b }
+  test([], [], f)
+  test([], [2, 10, 3], f)
+  test([10], [2, 10, 3], f)
+  test([10, 20], [2, 10, 3], f)
+  test([10, 20, 30], [2, 10, 3], f)
+  test([10, 20, 30], [2, 10], f)
+  test([10, 20, 30], [2], f)
+  test([10, 20, 30], [], f)
+}
+
+pub fn map2_is_lazy_test() {
+  let one = iterator.from_list([])
+  let other = iterator.once(fn() { panic as "unreachable" })
+
+  iterator.map2(one, other, fn(x, y) { x + y })
+  |> iterator.to_list
+  |> should.equal([])
+}
+
 // a |> from_list |> flat_map(f) |> to_list ==
 //   a |> list.map(f) |> list.map(to_list) |> list.concat
 pub fn flat_map_test() {