diff options
-rw-r--r-- | src/gleam/iterator.gleam | 62 | ||||
-rw-r--r-- | test/gleam/iterator_test.gleam | 14 |
2 files changed, 76 insertions, 0 deletions
diff --git a/src/gleam/iterator.gleam b/src/gleam/iterator.gleam index 96fd848..6db10e0 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_take_while( + continuation: fn() -> Action(element), + predicate: fn(element) -> Bool, +) -> fn() -> Action(element) { + fn() { + case continuation() { + Stop -> Stop + Continue(e, next) -> + case predicate(e) { + False -> Stop + True -> Continue(e, do_take_while(next, predicate)) + } + } + } +} + fn do_scan( continuation: fn() -> Action(element), accumulator: acc, @@ -568,6 +584,52 @@ fn do_zip( } } +/// Creates an iterator that yields elements while the predicate returns `True`. +/// +/// ## Examples +/// +/// > from_list([1, 2, 3, 2, 4]) |> take_while(satisfying: fn(x) { x < 3 }) |> to_list +/// [1, 2] +/// +pub fn take_while( + in iterator: Iterator(element), + satisfying predicate: fn(element) -> Bool, +) -> Iterator(element) { + iterator.continuation + |> do_take_while(predicate) + |> Iterator +} + +fn do_drop_while( + continuation: fn() -> Action(element), + predicate: fn(element) -> Bool, +) -> Action(element) { + case continuation() { + Stop -> Stop + Continue(e, next) -> + case predicate(e) { + False -> Continue(e, next) + True -> do_drop_while(next, predicate) + } + } +} + +/// Creates an iterator that drops elements while the predicate returns `True`, +/// and then yields the remaining elements. +/// +/// ## Examples +/// +/// > from_list([1, 2, 3, 4, 2, 5]) |> drop_while(satisfying: fn(x) { x < 4 }) |> to_list +/// [4, 2, 5] +/// +pub fn drop_while( + in iterator: Iterator(element), + satisfying predicate: fn(element) -> Bool, +) -> Iterator(element) { + fn() { do_drop_while(iterator.continuation, predicate) } + |> Iterator +} + /// Creates an iterator from an existing iterator and a stateful function. /// /// Specifically, this behaves like `fold`, but yields intermediate results. diff --git a/test/gleam/iterator_test.gleam b/test/gleam/iterator_test.gleam index f1ee880..b2ca8ac 100644 --- a/test/gleam/iterator_test.gleam +++ b/test/gleam/iterator_test.gleam @@ -265,6 +265,20 @@ pub fn iterate_test() { |> should.equal([1, 3, 9, 27, 81]) } +pub fn take_while_test() { + iterator.from_list([1, 2, 3, 2, 4]) + |> iterator.take_while(satisfying: fn(x) { x < 3 }) + |> iterator.to_list + |> should.equal([1, 2]) +} + +pub fn drop_while_test() { + iterator.from_list([1, 2, 3, 4, 2, 5]) + |> iterator.drop_while(satisfying: fn(x) { x < 4 }) + |> iterator.to_list + |> should.equal([4, 2, 5]) +} + pub fn scan_test() { iterator.from_list([1, 2, 3, 4, 5]) |> iterator.scan(from: 0, with: fn(el, acc) { acc + el }) |