Refactor Enum.chunks to use Enumerable.reduce/3, avoiding starting lazy collections twice

lib/elixir/lib/enum.ex

@@ -1233,6 +1233,12 @@ defmodule Enum do
     end) |> elem(0)
   end
 
+  @doc """
+  Shortcut to `chunks(coll, n, n)`.
+  """
+  @spec chunks(t, non_neg_integer) :: [list]
+  def chunks(coll, n), do: chunks(coll, n, n, nil)
+
   @doc """
   Returns a collection of lists containing `n` items each, where
   each new chunk starts `step` elements into the collection.
@@ -1253,23 +1259,49 @@ defmodule Enum do
       [[1, 2, 3], [3, 4, 5]]
       iex> Enum.chunks([1, 2, 3, 4, 5, 6], 3, 2, [7])
       [[1, 2, 3], [3, 4, 5], [5, 6, 7]]
+
   """
-  @spec chunks(t, integer) :: [list]
-  @spec chunks(t, integer, integer) :: [list]
-  @spec chunks(t, integer, integer, list) :: [list]
-  def chunks(coll, n), do: chunks(coll, n, n, nil)
-  def chunks(coll, n, step, pad // nil)
-  def chunks([], _, _, _), do: []
-  def chunks(coll, n, step, pad) do
-    if segment = pad_segment(Enum.take(coll, n), pad, n) do
-      [segment | chunks(Enum.drop(coll, step), n, step, pad)]
+  @spec chunks(t, non_neg_integer, non_neg_integer) :: [list]
+  @spec chunks(t, non_neg_integer, non_neg_integer, list) :: [list]
+  def chunks(coll, n, step, pad // nil) when n > 0 and step > 0 do
+    { acc, buffer, i } =
+      Enumerable.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)
+
+    if nil?(pad) do
+      :lists.reverse(acc)
     else
-      []
+      buffer = :lists.reverse(buffer) ++ take(pad, n - i)
+      :lists.reverse([buffer|acc])
     end
   end
 
   ## Helpers
 
+  @compile { :inline, chunks_n: 5, chunks_step: 4 }
+
+  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
+
   defp iterate_and_count(collection, count) do
     { list, total_items } = do_iterate_and_count(collection)
     { list, Kernel.max(0, total_items - abs(count)) }
@@ -1284,13 +1316,6 @@ defmodule Enum do
     map_reduce(collection, 0, reducer)
   end
 
-  defp pad_segment(segment, pad, n) do
-    case count = Enum.count(segment) do
-      ^n -> segment
-      _  -> if pad, do: segment ++ Enum.take(pad, n - count)
-    end
-  end
-
   ## Implementations
 
   ## all?

lib/elixir/test/elixir/enum_test.exs

@@ -302,6 +302,7 @@ defmodule EnumTest.List do
     assert Enum.chunks([1, 2, 3, 4, 5, 6], 3, 2) == [[1, 2, 3], [3, 4, 5]]
     assert Enum.chunks([1, 2, 3, 4, 5, 6], 2, 3) == [[1, 2], [4, 5]]
     assert Enum.chunks([1, 2, 3, 4, 5, 6], 3, 2, []) == [[1, 2, 3], [3, 4, 5], [5, 6]]
+    assert Enum.chunks([1, 2, 3, 4, 5], 4, 4, 6..10) == [[1, 2, 3, 4], [5, 6, 7, 8]]
   end
 end
 
@@ -435,7 +436,6 @@ defmodule EnumTest.Range do
     after
       Process.delete(:enum_test_each)
     end
-
   end
 
   test :filter do
@@ -632,6 +632,15 @@ defmodule EnumTest.Range do
     assert Enum.min_by(1..1, fn(x) -> :math.pow(-2, x) end) == 1
     assert Enum.min_by(1..3, fn(x) -> :math.pow(-2, x) end) == 3
   end
+
+  test :chunks do
+    assert Enum.chunks(1..5, 2) == [[1, 2], [3, 4]]
+    assert Enum.chunks(1..5, 2, 2, [6]) == [[1, 2], [3, 4], [5, 6]]
+    assert Enum.chunks(1..6, 3, 2) == [[1, 2, 3], [3, 4, 5]]
+    assert Enum.chunks(1..6, 2, 3) == [[1, 2], [4, 5]]
+    assert Enum.chunks(1..6, 3, 2, []) == [[1, 2, 3], [3, 4, 5], [5, 6]]
+    assert Enum.chunks(1..5, 4, 4, 6..10) == [[1, 2, 3, 4], [5, 6, 7, 8]]
+  end
 end
 
 defmodule EnumTest.Others do