iterator.fold_until

2 files changed, 71 insertions, 0 deletions

diff --git a/src/gleam/iterator.gleam b/src/gleam/iterator.gleam
index 25cc441..414956d 100644
--- a/src/gleam/iterator.gleam
+++ b/src/gleam/iterator.gleam
@@ -1036,3 +1036,48 @@ if erlang {
     |> Iterator
   }
 }
+
+fn do_fold_until(
+  continuation: fn() -> Action(e),
+  f: fn(e, acc) -> list.ContinueOrStop(acc),
+  accumulator: acc,
+) -> acc {
+  case continuation() {
+    Stop -> accumulator
+    Continue(elem, next) ->
+      case f(elem, accumulator) {
+        list.Continue(accumulator) -> do_fold_until(next, f, accumulator)
+        list.Stop(accumulator) -> accumulator
+      }
+  }
+}
+
+/// Like `fold`, `fold_until` reduces an iterator of elements into a single value by calling a given
+/// function on each element in turn, but uses a `list.ContinueOrStop` to determine 
+/// whether or not to keep iterating.
+///
+/// If called on an iterator of infinite length then this function will only ever
+/// return if the give function returns list.Stop.
+///
+///
+/// ## Examples
+///    > let f = fn(e, acc) {
+///    >   case e {
+///    >     _ if e < 4 -> list.Continue(e + acc)
+///    >     _ -> list.Stop(acc)
+///    >   }
+///    > }
+///    >
+///    > [1, 2, 3, 4]
+///    > |> from_list
+///    > |> iterator.fold_until(from: acc, with: f) 
+///    6
+///
+pub fn fold_until(
+  over iterator: Iterator(e),
+  from initial: acc,
+  with f: fn(e, acc) -> list.ContinueOrStop(acc),
+) -> acc {
+  iterator.continuation
+  |> do_fold_until(f, initial)
+}
diff --git a/test/gleam/iterator_test.gleam b/test/gleam/iterator_test.gleam
index 0c0176f..c8e2c4d 100644
--- a/test/gleam/iterator_test.gleam
+++ b/test/gleam/iterator_test.gleam
@@ -421,3 +421,29 @@ if erlang {
     |> should.equal([1, 100, 2, 3, 4])
   }
 }
+
+// a |> from_list |> fold_until(a, f) == a |> list.fold_until(_, a, f)
+pub fn fold_until_test() {
+  let test = fn(subject, acc, f) {
+    subject
+    |> iterator.from_list()
+    |> iterator.fold_until(acc, f)
+    |> should.equal(list.fold_until(subject, acc, f))
+  }
+
+  let f = fn(e, acc) {
+    case e {
+      _ if e < 6 -> list.Continue([e, ..acc])
+      _ -> list.Stop(acc)
+    }
+  }
+  test([], [], f)
+  test([1], [], f)
+  test([1, 2, 3], [], f)
+  test([1, 2, 3, 4, 5, 6, 7, 8], [], f)
+
+  [1, 2, 3, 4, 5, 6, 7, 8]
+  |> iterator.from_list()
+  |> iterator.fold_until([], f)
+  |> should.equal([5, 4, 3, 2, 1])
+}