Add iterator.scan (#173)

2 files changed, 43 insertions, 0 deletions

diff --git a/src/gleam/iterator.gleam b/src/gleam/iterator.gleam
index 010d82a..96fd848 100644
--- a/src/gleam/iterator.gleam
+++ b/src/gleam/iterator.gleam
@@ -535,6 +535,22 @@ pub fn iterate(
   unfold(initial, fn(element) { Next(element, f(element)) })
 }
 
+fn do_scan(
+  continuation: fn() -> Action(element),
+  accumulator: acc,
+  f: fn(element, acc) -> acc,
+) -> fn() -> Action(acc) {
+  fn() {
+    case continuation() {
+      Continue(el, next) -> {
+        let accumulated = f(el, accumulator)
+        Continue(accumulated, do_scan(next, accumulated, f))
+      }
+      Stop -> Stop
+    }
+  }
+}
+
 fn do_zip(
   left: fn() -> Action(a),
   right: fn() -> Action(b),
@@ -552,6 +568,26 @@ fn do_zip(
   }
 }
 
+/// Creates an iterator from an existing iterator and a stateful function.
+///
+/// Specifically, this behaves like `fold`, but yields intermediate results.
+///
+/// ## Examples
+///
+///    Generate a sequence of partial sums:
+///    > from_list([1, 2, 3, 4, 5]) |> scan(from: 0, with: fn(el, acc) { acc + el }) |> to_list
+///    [1, 3, 6, 10, 15]
+///
+pub fn scan(
+  over iterator: Iterator(element),
+  from initial: acc,
+  with f: fn(element, acc) -> acc,
+) -> Iterator(acc) {
+  iterator.continuation
+  |> do_scan(initial, f)
+  |> Iterator
+}
+
 /// Zips two iterators together, emitting values from both
 /// until the shorter one runs out.
 ///
diff --git a/test/gleam/iterator_test.gleam b/test/gleam/iterator_test.gleam
index ffe57e4..f1ee880 100644
--- a/test/gleam/iterator_test.gleam
+++ b/test/gleam/iterator_test.gleam
@@ -265,6 +265,13 @@ pub fn iterate_test() {
   |> should.equal([1, 3, 9, 27, 81])
 }
 
+pub fn scan_test() {
+  iterator.from_list([1, 2, 3, 4, 5])
+  |> iterator.scan(from: 0, with: fn(el, acc) { acc + el })
+  |> iterator.to_list
+  |> should.equal([1, 3, 6, 10, 15])
+}
+
 pub fn zip_test() {
   iterator.from_list(["a", "b", "c"])
   |> iterator.zip(iterator.range(20, 30))