Add iterator.sized_chunk and list.sized_chunk

4 files changed, 136 insertions, 5 deletions

diff --git a/src/gleam/iterator.gleam b/src/gleam/iterator.gleam
index 3a939d0..38602c1 100644
--- a/src/gleam/iterator.gleam
+++ b/src/gleam/iterator.gleam
@@ -6,11 +6,6 @@ type Action(element) {
   Continue(element, fn() -> Action(element))
 }
 
-// Shortcut for an empty iterator
-fn stop() -> Action(element) {
-  Stop
-}
-
 /// An iterator is a lazily evaluated sequence of element.
 ///
 /// Iterators are useful when working with collections that are too large to
@@ -32,6 +27,11 @@ pub type Step(element, accumulator) {
   Done
 }
 
+// Shortcut for an empty iterator.
+fn stop() -> Action(element) {
+  Stop
+}
+
 // Creating Iterators
 fn do_unfold(
   initial: acc,
@@ -661,6 +661,7 @@ pub fn zip(left: Iterator(a), right: Iterator(b)) -> Iterator(tuple(a, b)) {
   |> Iterator
 }
 
+// Result of collecting a single chunk by key
 type Chunk(element, key) {
   AnotherBy(List(element), key, element, fn() -> Action(element))
   LastBy(List(element))
@@ -717,3 +718,69 @@ pub fn chunk(
   }
   |> Iterator
 }
+
+// Result of collecting a single sized chunk
+type SizedChunk(element) {
+  Another(List(element), fn() -> Action(element))
+  Last(List(element))
+  None
+}
+
+fn next_sized_chunk(
+  continuation: fn() -> Action(element),
+  left: Int,
+  current_chunk: List(element),
+) -> SizedChunk(element) {
+  case continuation() {
+    Stop ->
+      case current_chunk {
+        [] -> None
+        remaining -> Last(list.reverse(remaining))
+      }
+    Continue(e, next) -> {
+      let chunk = [e, ..current_chunk]
+      case left > 1 {
+        False -> Another(list.reverse(chunk), next)
+        True -> next_sized_chunk(next, left - 1, chunk)
+      }
+    }
+  }
+}
+
+fn do_sized_chunk(
+  continuation: fn() -> Action(element),
+  count: Int,
+) -> fn() -> Action(List(element)) {
+  fn() {
+    case next_sized_chunk(continuation, count, []) {
+      None -> Stop
+      Last(chunk) -> Continue(chunk, stop)
+      Another(chunk, next_element) ->
+        Continue(chunk, do_sized_chunk(next_element, count))
+    }
+  }
+}
+
+/// Creates an iterator that emits chunks of given size.
+///
+/// If the last chunk does not have `count` elements, it is yielded
+/// as a partial chunk, with less than `count` elements.
+///
+/// For any `count` less than 1 this function behaves as if it was set to 1.
+///
+/// ## Examples
+///
+///    > from_list([1, 2, 3, 4, 5, 6]) |> chunk(into: 2) |> to_list
+///    [[1, 2], [3, 4], [5, 6]]
+///
+///    > from_list([1, 2, 3, 4, 5, 6, 7, 8]) |> chunk(into: 3) |> to_list
+///    [[1, 2, 3], [4, 5, 6], [7, 8]]
+///
+pub fn sized_chunk(
+  over iterator: Iterator(element),
+  into count: Int,
+) -> Iterator(List(element)) {
+  iterator.continuation
+  |> do_sized_chunk(count)
+  |> Iterator
+}
diff --git a/src/gleam/list.gleam b/src/gleam/list.gleam
index 066c584..422f26b 100644
--- a/src/gleam/list.gleam
+++ b/src/gleam/list.gleam
@@ -1303,3 +1303,45 @@ pub fn chunk(in list: List(a), by f: fn(a) -> key) -> List(List(a)) {
     [head, ..tail] -> do_chunk(tail, f, f(head), [head], [])
   }
 }
+
+fn do_sized_chunk(
+  list: List(a),
+  count: Int,
+  left: Int,
+  current_chunk: List(a),
+  acc: List(List(a)),
+) -> List(List(a)) {
+  case list {
+    [] ->
+      case current_chunk {
+        [] -> reverse(acc)
+        remaining -> reverse([reverse(remaining), ..acc])
+      }
+    [head, ..tail] -> {
+      let chunk = [head, ..current_chunk]
+      case left > 1 {
+        False -> do_sized_chunk(tail, count, count, [], [reverse(chunk), ..acc])
+        True -> do_sized_chunk(tail, count, left - 1, chunk, acc)
+      }
+    }
+  }
+}
+
+/// Returns a list of chunks containing `count` elements each.
+///
+/// If the last chunk does not have `count` elements, it is instead
+/// a partial chunk, with less than `count` elements.
+///
+/// For any `count` less than 1 this function behaves as if it was set to 1.
+///
+/// ## Examples
+///
+///    > [1, 2, 3, 4, 5, 6] |> chunk(into: 2)
+///    [[1, 2], [3, 4], [5, 6]]
+///
+///    > [1, 2, 3, 4, 5, 6, 7, 8] |> chunk(into: 3)
+///    [[1, 2, 3], [4, 5, 6], [7, 8]]
+///
+pub fn sized_chunk(in list: List(a), into count: Int) -> List(List(a)) {
+  do_sized_chunk(list, count, count, [], [])
+}
diff --git a/test/gleam/iterator_test.gleam b/test/gleam/iterator_test.gleam
index 49c65fe..0c85ede 100644
--- a/test/gleam/iterator_test.gleam
+++ b/test/gleam/iterator_test.gleam
@@ -305,3 +305,15 @@ pub fn chunk_test() {
   |> iterator.to_list
   |> should.equal([[1], [2, 2], [3], [4, 4, 6], [7, 7]])
 }
+
+pub fn sized_chunk_test() {
+  iterator.from_list([1, 2, 3, 4, 5, 6])
+  |> iterator.sized_chunk(into: 2)
+  |> iterator.to_list
+  |> should.equal([[1, 2], [3, 4], [5, 6]])
+
+  iterator.from_list([1, 2, 3, 4, 5, 6, 7, 8])
+  |> iterator.sized_chunk(into: 3)
+  |> iterator.to_list
+  |> should.equal([[1, 2, 3], [4, 5, 6], [7, 8]])
+}
diff --git a/test/gleam/list_test.gleam b/test/gleam/list_test.gleam
index 8c414a8..b561d33 100644
--- a/test/gleam/list_test.gleam
+++ b/test/gleam/list_test.gleam
@@ -596,3 +596,13 @@ pub fn chunk_test() {
   |> list.chunk(by: fn(n) { n % 2 })
   |> should.equal([[1], [2, 2], [3], [4, 4, 6], [7, 7]])
 }
+
+pub fn sized_chunk_test() {
+  [1, 2, 3, 4, 5, 6]
+  |> list.sized_chunk(into: 2)
+  |> should.equal([[1, 2], [3, 4], [5, 6]])
+
+  [1, 2, 3, 4, 5, 6, 7, 8]
+  |> list.sized_chunk(into: 3)
+  |> should.equal([[1, 2, 3], [4, 5, 6], [7, 8]])
+}