iterator.find_map (#573)

3 files changed, 81 insertions, 0 deletions

diff --git a/CHANGELOG.md b/CHANGELOG.md
index fdae968..a61e116 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -3,6 +3,7 @@
 ## Unreleased
 
 - The `list` module gains the `wrap` function.
+- The `iterator` module gains the `find_map` function.
 
 ## v0.37.0 - 2024-04-19
 
diff --git a/src/gleam/iterator.gleam b/src/gleam/iterator.gleam
index c5cd0af..bb6be98 100644
--- a/src/gleam/iterator.gleam
+++ b/src/gleam/iterator.gleam
@@ -712,6 +712,51 @@ pub fn find(
   |> do_find(is_desired)
 }
 
+fn do_find_map(
+  continuation: fn() -> Action(a),
+  f: fn(a) -> Result(b, c),
+) -> Result(b, Nil) {
+  case continuation() {
+    Stop -> Error(Nil)
+    Continue(e, next) ->
+      case f(e) {
+        Ok(e) -> Ok(e)
+        Error(_) -> do_find_map(next, f)
+      }
+  }
+}
+
+/// Finds the first element in a given iterator
+/// for which the given function returns `Ok(new_value)`,
+/// then returns the wrapped `new_value`.
+///
+/// Returns `Error(Nil)` if no such element is found.
+///
+/// ## Examples
+///
+/// ```gleam
+/// find_map(from_list([1, 2, 3]), first)
+/// // -> Ok(1)
+/// ```
+///
+/// ```gleam
+/// find_map(from_list([]), first)
+/// // -> Error(Nil)
+/// ```
+///
+/// ```gleam
+/// find(empty(), first)
+/// // -> Error(Nil)
+/// ```
+///
+pub fn find_map(
+  in haystack: Iterator(a),
+  one_that is_desired: fn(a) -> Result(b, c),
+) -> Result(b, Nil) {
+  haystack.continuation
+  |> do_find_map(is_desired)
+}
+
 fn do_index(
   continuation: fn() -> Action(element),
   next: Int,
diff --git a/test/gleam/iterator_test.gleam b/test/gleam/iterator_test.gleam
index cd58035..186ea69 100644
--- a/test/gleam/iterator_test.gleam
+++ b/test/gleam/iterator_test.gleam
@@ -376,6 +376,41 @@ pub fn find_test() {
   |> should.equal(Ok(Cat(id: 10)))
 }
 
+pub fn find_map_test() {
+  iterator.range(0, 10)
+  |> iterator.find_map(fn(e) {
+    case e == 5 {
+      True -> Ok(e)
+      False -> Error(Nil)
+    }
+  })
+  |> should.equal(Ok(5))
+
+  iterator.range(0, 10)
+  |> iterator.find_map(fn(e) {
+    case e > 10 {
+      True -> Ok(e)
+      False -> Error(Nil)
+    }
+  })
+  |> should.equal(Error(Nil))
+
+  iterator.empty()
+  |> iterator.find_map(fn(_x) { Ok(True) })
+  |> should.equal(Error(Nil))
+
+  iterator.unfold(Cat(id: 1), fn(cat: Cat) {
+    iterator.Next(cat, Cat(id: cat.id + 1))
+  })
+  |> iterator.find_map(fn(cat: Cat) {
+    case cat.id == 10 {
+      True -> Ok(cat)
+      False -> Error(Nil)
+    }
+  })
+  |> should.equal(Ok(Cat(id: 10)))
+}
+
 pub fn index_test() {
   iterator.from_list(["a", "b", "c"])
   |> iterator.index