Add flat_map and friends to iterator

The following functions are added to iterator:

* append
* flatten
* flat_map

Author: drewolson

CHANGELOG.md

@@ -10,6 +10,7 @@
 - The `dynamic` module gains the `option` function.
 - The `uri` module gains the `percent_encode` and `percent_decode` functions.
 - The `os` module gains the `erlang_timestamp` function.
+- The `iterator` module gains the `append`, `flatten` and `flat_map` functions.
 
 ## v0.11.0 - 2020-08-22
 

src/gleam/iterator.gleam

@@ -268,6 +268,83 @@ pub fn map(over iterator: Iterator(a), with f: fn(a) -> b) -> Iterator(b) {
   |> Iterator
 }
 
+fn do_append(
+  first: fn() -> Action(a),
+  second: fn() -> Action(a),
+) -> fn() -> Action(a) {
+  fn() {
+    case first() {
+      Continue(e, first) -> Continue(e, do_append(first, second))
+      Stop -> second()
+    }
+  }
+}
+
+/// Append two iterators, producing a new iterator.
+///
+/// This function does not evaluate the elements of the iterators, the
+/// computation is performed when the resulting iterator is later run.
+///
+/// ## Examples
+///
+///    > [1, 2] |> from_list |> append([3, 4] |> from_list) |> to_list
+///    [1, 2, 3, 4]
+///
+pub fn append(to first: Iterator(a), suffix second: Iterator(a)) -> Iterator(a) {
+  first.continuation
+  |> do_append(second.continuation)
+  |> Iterator
+}
+
+fn do_flatten(continuation: fn() -> Action(Iterator(a))) -> fn() -> Action(a) {
+  fn() {
+    case continuation() {
+      Continue(e, continuation) ->
+        do_append(e.continuation, do_flatten(continuation))()
+      Stop -> Stop
+    }
+  }
+}
+
+/// Flatten an iterator of iterator of iterators, creating a new iterator.
+///
+/// This function does not evaluate the elements of the iterator, the
+/// computation is performed when the iterator is later run.
+///
+/// ## Examples
+///
+///    > [[1, 2], [3, 4]] |> list.map(from_list) |> flatten |> to_list
+///    [1, 2, 3, 4]
+///
+pub fn flatten(iterator: Iterator(Iterator(a))) -> Iterator(a) {
+  iterator.continuation
+  |> do_flatten
+  |> Iterator
+}
+
+/// Create an iterator from an existing iterator and a transformation function.
+///
+/// Each element in the new iterator will be the result of calling the given
+/// function on the elements in the given iterator and then flattening the
+/// results.
+///
+/// This function does not evaluate the elements of the iterator, the
+/// computation is performed when the iterator is later run.
+///
+/// ## Examples
+///
+///    > [1, 2] |> from_list |> flat_map(fn(x) { from_list([x, x + 1]) }) |> to_list
+///    [1, 2, 2, 3]
+///
+pub fn flat_map(
+  over iterator: Iterator(a),
+  with f: fn(a) -> Iterator(b),
+) -> Iterator(b) {
+  iterator
+  |> map(f)
+  |> flatten
+}
+
 fn do_filter(
   continuation: fn() -> Action(e),
   predicate: fn(e) -> Bool,

test/gleam/iterator_test.gleam

@@ -71,6 +71,60 @@ pub fn map_test() {
   test([1, 2, 3, 4, 5, 6, 7, 8], f)
 }
 
+// a |> from_list |> flat_map(f) |> to_list ==
+//   a |> list.map(f) |> list.map(to_list) |> list.flatten
+pub fn flat_map_test() {
+  let test = fn(subject, f) {
+    subject
+    |> iterator.from_list
+    |> iterator.flat_map(f)
+    |> iterator.to_list
+    |> should.equal(
+      subject
+      |> list.map(f)
+      |> list.map(iterator.to_list)
+      |> list.flatten,
+    )
+  }
+
+  let f = fn(i) { iterator.range(i, i + 2) }
+
+  test([], f)
+  test([1], f)
+  test([1, 2], f)
+}
+
+// a |> from_list |> append(from_list(b)) |> to_list == list.flatten([a, b])
+pub fn append_test() {
+  let test = fn(left, right) {
+    left
+    |> iterator.from_list
+    |> iterator.append(iterator.from_list(right))
+    |> iterator.to_list
+    |> should.equal(list.flatten([left, right]))
+  }
+
+  test([], [])
+  test([1], [2])
+  test([1, 2], [3, 4])
+}
+
+// a |> list.map(from_list) |> flatten |> to_list == list.flatten(a)
+pub fn flatten_test() {
+  let test = fn(lists) {
+    lists
+    |> list.map(iterator.from_list)
+    |> iterator.from_list
+    |> iterator.flatten
+    |> iterator.to_list
+    |> should.equal(list.flatten(lists))
+  }
+
+  test([[], []])
+  test([[1], [2]])
+  test([[1, 2], [3, 4]])
+}
+
 // a |> from_list |> filter(f) |> to_list == a |> list.filter(_, f)
 pub fn filter_test() {
   let test = fn(subject, f) {