diff --git a/src/gleam/iterator.gleam b/src/gleam/iterator.gleam index 336c0a10..216a3bee 100644 --- a/src/gleam/iterator.gleam +++ b/src/gleam/iterator.gleam @@ -1627,3 +1627,59 @@ pub fn each(over iterator: Iterator(a), with f: fn(a) -> b) -> Nil { pub fn yield(element: a, next: fn() -> Iterator(a)) -> Iterator(a) { Iterator(fn() { Continue(element, fn() { next().continuation() }) }) } + +/// Returns an iterator of sliding windows. +/// +/// ## Examples +/// +/// ```gleam +/// from_list([1,2,3,4,5]) +/// |> window(3) +/// |> map(to_list) +/// |> to_list +/// // -> [[1, 2, 3], [2, 3, 4], [3, 4, 5]] +/// ``` +/// +/// ```gleam +/// from_list([1,2]) +/// |> window(4) +/// |> map(to_list) +/// |> to_list +/// // -> [] +/// ``` +/// +pub fn window(i: Iterator(a), by n: Int) -> Iterator(Iterator(a)) { + let yield = fn(iterator) { + let window = take(iterator, n) + + case length(window) == n { + True -> Next(window, drop(iterator, 1)) + False -> Done + } + } + + case n > 0 { + True -> unfold(from: i, with: yield) + False -> empty() + } +} + +/// Returns an iterator of tuples containing two contiguous elements. +/// +/// ## Examples +/// +/// ```gleam +/// window_by_2([1,2,3,4]) +/// |> to_list +/// // -> [#(1, 2), #(2, 3), #(3, 4)] +/// ``` +/// +/// ```gleam +/// window_by_2([1]) +/// |> to_list +/// // -> [] +/// ``` +/// +pub fn window_by_2(i: Iterator(a)) -> Iterator(#(a, a)) { + zip(i, drop(i, 1)) +} diff --git a/test/gleam/iterator_test.gleam b/test/gleam/iterator_test.gleam index 6ae1d4c7..01415c6c 100644 --- a/test/gleam/iterator_test.gleam +++ b/test/gleam/iterator_test.gleam @@ -730,3 +730,34 @@ pub fn yield_computes_only_necessary_values_test() { |> iterator.to_list |> should.equal([1, 2, 3]) } + +pub fn window_test() { + let testcase = fn(subject, n, expected) { + subject + |> iterator.from_list + |> iterator.window(n) + |> iterator.map(iterator.to_list) + |> iterator.to_list + |> should.equal(expected) + } + + testcase([1, 2, 3], 2, [[1, 2], [2, 3]]) + testcase([1, 2, 3], 3, [[1, 2, 3]]) + testcase([1, 2, 3], 4, []) + testcase([1, 2, 3, 4, 5], 3, [[1, 2, 3], [2, 3, 4], [3, 4, 5]]) + testcase([1, 2, 3], 0, []) + testcase([1, 2, 3], -1, []) +} + +pub fn window_by_2_test() { + let testcase = fn(subject, expected) { + subject + |> iterator.from_list + |> iterator.window_by_2 + |> iterator.to_list + |> should.equal(expected) + } + + testcase([1, 2, 3, 4], [#(1, 2), #(2, 3), #(3, 4)]) + testcase([1], []) +}