Add iterator.take_while_inclusive

CHANGELOG.md

@@ -5,6 +5,7 @@
 - The `float` module gains the `loosely_equals` function.
 - The `io` module gains `print_error` and `println_error` functions for
   printing to stderr.
+- The `iterator` module gains the `take_while_inclusive` function.
 - The `io.debug` function now prints to stderr instead of stdout when using
   the Erlang target or running in Node.js (but still uses `console.log`
   when running as JavaScript in a browser)

src/gleam/iterator.gleam

@@ -653,6 +653,45 @@ pub fn take_while(
   |> Iterator
 }
 
+fn do_take_while_inclusive(
+  continuation: fn() -> Action(element),
+  predicate: fn(element) -> Bool,
+  continue: Bool,
+) -> fn() -> Action(element) {
+  fn() {
+    case continue {
+      False -> Stop
+      True ->
+        case continuation() {
+          Stop -> Stop
+          Continue(e, next) ->
+            Continue(e, do_take_while_inclusive(next, predicate, predicate(e)))
+        }
+    }
+  }
+}
+
+/// Creates an iterator that yields elements while the predicate returns `True`,
+/// and also yields the element for which the predicate first returned `False`.
+///
+/// ## Examples
+///
+/// ```gleam
+/// > from_list([1, 2, 3, 2, 4])
+/// > |> take_while_inclusive(satisfying: fn(x) { x < 3 })
+/// > |> to_list
+/// [1, 2, 3]
+/// ```
+///
+pub fn take_while_inclusive(
+  in iterator: Iterator(element),
+  satisfying predicate: fn(element) -> Bool,
+) -> Iterator(element) {
+  iterator.continuation
+  |> do_take_while_inclusive(predicate, True)
+  |> Iterator
+}
+
 fn do_drop_while(
   continuation: fn() -> Action(element),
   predicate: fn(element) -> Bool,

test/gleam/iterator_test.gleam

@@ -284,6 +284,13 @@ pub fn take_while_test() {
   |> should.equal([1, 2])
 }
 
+pub fn take_while_inclusive_test() {
+  iterator.from_list([1, 2, 3, 2, 4])
+  |> iterator.take_while_inclusive(satisfying: fn(x) { x < 3 })
+  |> iterator.to_list
+  |> should.equal([1, 2, 3])
+}
+
 pub fn drop_while_test() {
   iterator.from_list([1, 2, 3, 4, 2, 5])
   |> iterator.drop_while(satisfying: fn(x) { x < 4 })