diff options
| -rw-r--r-- | src/gleam/iterator.gleam | 36 | ||||
| -rw-r--r-- | test/gleam/iterator_test.gleam | 7 |
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)) |