Add Stream.chunks/2,3,4

Author: José Valim

lib/elixir/lib/enum.ex

@@ -311,15 +311,12 @@ defmodule Enum do
 
   """
   @spec chunks(t, non_neg_integer, non_neg_integer) :: [list]
-  @spec chunks(t, non_neg_integer, non_neg_integer, list | nil) :: [list]
+  @spec chunks(t, non_neg_integer, non_neg_integer, t | nil) :: [list]
   def chunks(coll, n, step, pad // nil) when n > 0 and step > 0 do
-    { acc, buffer, i } =
-      reduce(coll, { [], [], 0 }, fn
-        x, { acc, buffer, i } when i < n ->
-          chunks_n(acc, [x|buffer], i + 1, n, step)
-        x, { acc, buffer, i } when i < step ->
-          chunks_step(acc, [x|buffer], i + 1, step)
-      end)
+    limit = :erlang.max(n, step)
+
+    { _, { acc, { buffer, i } } } =
+      Enumerable.reduce(coll, { :cont, { [], { [], 0 } } }, R.chunks(n, step, limit))
 
     if nil?(pad) || i == 0 do
       :lists.reverse(acc)
@@ -1599,25 +1596,7 @@ defmodule Enum do
 
   ## Helpers
 
-  @compile { :inline, chunks_n: 5, chunks_step: 4, to_string: 2 }
-
-  defp chunks_n(acc, buffer, i, n, step) when i == n do
-    acc = [:lists.reverse(buffer)|acc]
-    chunks_step(acc, buffer, i, step)
-  end
-
-  defp chunks_n(acc, buffer, i, _n, _step) do
-    { acc, buffer, i }
-  end
-
-  defp chunks_step(acc, buffer, i, step) when i >= step do
-    left = i - step
-    { acc, take(buffer, left), left }
-  end
-
-  defp chunks_step(acc, buffer, i, _step) do
-    { acc, buffer, i }
-  end
+  @compile { :inline, to_string: 2 }
 
   defp enumerate_and_count(collection, count) when is_list(collection) do
     { collection, length(collection) - abs(count) }

lib/elixir/lib/stream.ex

@@ -163,7 +163,56 @@ defmodule Stream do
   ## Transformers
 
   @doc """
-  Chunk the `enum` by buffering elements for which `fun` returns
+  Shortcut to `chunks(coll, n, n)`.
+  """
+  @spec chunks(Enumerable.t, non_neg_integer) :: Enumerable.t
+  def chunks(coll, n), do: chunks(coll, n, n, nil)
+
+  @doc """
+  Streams the enumerable in chunks, containing `n` items each, where
+  each new chunk starts `step` elements into the enumerable.
+
+  `step` is optional and, if not passed, defaults to `n`, i.e.
+  chunks do not overlap. If the final chunk does not have `n`
+  elements to fill the chunk, elements are taken as necessary
+  from `pad` if it was passed. If `pad` is passed and does not
+  have enough elements to fill the chunk, then the chunk is
+  returned anyway with less than `n` elements. If `pad` is not
+  passed at all or is nil, then the partial chunk is discarded
+  from the result.
+
+  ## Examples
+
+      iex> Stream.chunks([1, 2, 3, 4, 5, 6], 2) |> Enum.to_list
+      [[1, 2], [3, 4], [5, 6]]
+      iex> Stream.chunks([1, 2, 3, 4, 5, 6], 3, 2) |> Enum.to_list
+      [[1, 2, 3], [3, 4, 5]]
+      iex> Stream.chunks([1, 2, 3, 4, 5, 6], 3, 2, [7]) |> Enum.to_list
+      [[1, 2, 3], [3, 4, 5], [5, 6, 7]]
+      iex> Stream.chunks([1, 2, 3, 4, 5, 6], 3, 3, []) |> Enum.to_list
+      [[1, 2, 3], [4, 5, 6]]
+
+  """
+  @spec chunks(Enumerable.t, non_neg_integer, non_neg_integer) :: Enumerable.t
+  @spec chunks(Enumerable.t, non_neg_integer, non_neg_integer, Enumerable.t | nil) :: Enumerable.t
+  def chunks(enum, n, step, pad // nil) when n > 0 and step > 0 do
+    limit = :erlang.max(n, step)
+    lazy enum, { [], 0 },
+         fn(f1) -> R.chunks(n, step, limit, f1) end,
+         fn(f1) -> &do_chunks(&1, n, pad, f1) end
+  end
+
+  defp do_chunks(acc(h, { buffer, count } = old, t) = acc, n, pad, f1) do
+    if nil?(pad) || count == 0 do
+      { :cont, acc }
+    else
+      buffer = :lists.reverse(buffer) ++ Enum.take(pad, n - count)
+      cont_with_acc(f1, buffer, h, old, t)
+    end
+  end
+
+  @doc """
+  Chunks the `enum` by buffering elements for which `fun` returns
   the same value and only emit them when `fun` returns a new value
   or the `enum` finishes,
 

lib/elixir/lib/stream/reducers.ex

@@ -2,6 +2,29 @@ defmodule Stream.Reducers do
   # Collection of reducers shared by Enum and Stream.
   @moduledoc false
 
+  defmacro chunks(n, step, limit, f // nil) do
+    quote do
+      fn entry, acc(h, { buffer, count }, t) ->
+        buffer = [entry|buffer]
+        count  = count + 1
+
+        new =
+          if count >= unquote(limit) do
+            left = count - unquote(step)
+            { Enum.take(buffer, left), left }
+          else
+            { buffer, count }
+          end
+
+        if count == unquote(n) do
+          cont_with_acc(unquote(f), :lists.reverse(buffer), h, new, t)
+        else
+          { :cont, acc(h, new, t) }
+        end
+      end
+    end
+  end
+
   defmacro chunks_by(fun, f // nil) do
     quote do
       fn

lib/elixir/test/elixir/stream_test.exs

@@ -25,6 +25,29 @@ defmodule StreamTest do
     assert Enum.to_list(stream) == [3,5,7]
   end
 
+  test "chunks" do
+    assert Stream.chunks([1, 2, 3, 4, 5], 2) |> Enum.to_list ==
+           [[1, 2], [3, 4]]
+    assert Stream.chunks([1, 2, 3, 4, 5], 2, 2, [6]) |> Enum.to_list ==
+           [[1, 2], [3, 4], [5, 6]]
+    assert Stream.chunks([1, 2, 3, 4, 5, 6], 3, 2) |> Enum.to_list ==
+           [[1, 2, 3], [3, 4, 5]]
+    assert Stream.chunks([1, 2, 3, 4, 5, 6], 2, 3) |> Enum.to_list ==
+           [[1, 2], [4, 5]]
+    assert Stream.chunks([1, 2, 3, 4, 5, 6], 3, 2, []) |> Enum.to_list ==
+           [[1, 2, 3], [3, 4, 5], [5, 6]]
+    assert Stream.chunks([1, 2, 3, 4, 5, 6], 3, 3, []) |> Enum.to_list ==
+           [[1, 2, 3], [4, 5, 6]]
+    assert Stream.chunks([1, 2, 3, 4, 5], 4, 4, 6..10) |> Enum.to_list ==
+           [[1, 2, 3, 4], [5, 6, 7, 8]]
+  end
+
+  test "chunks is zippable" do
+    stream = Stream.chunks([1, 2, 3, 4, 5, 6], 3, 2, [])
+    list   = Enum.to_list(stream)
+    assert Enum.zip(list, list) == Enum.zip(stream, stream)
+  end
+
   test "chunks_by" do
     stream = Stream.chunks_by([1, 2, 2, 3, 4, 4, 6, 7, 7], &(rem(&1, 2) == 1))