Share chunking logic between stream and enum

Signed-off-by: José Valim <jose.valim@plataformatec.com.br>

Author: José Valim

lib/elixir/lib/enum.ex

@@ -360,17 +360,7 @@ defmodule Enum do
   @spec chunk(t, pos_integer, pos_integer, t | nil) :: [list]
   def chunk(enumerable, count, step, leftover \\ nil)
       when is_integer(count) and count > 0 and is_integer(step) and step > 0 do
-    limit = :erlang.max(count, step)
-
-    {acc, {buffer, i}} =
-      reduce(enumerable, {[], {[], 0}}, R.chunk(count, step, limit))
-
-    if is_nil(leftover) || i == 0 do
-      :lists.reverse(acc)
-    else
-      buffer = :lists.reverse(buffer, take(leftover, count - i))
-      :lists.reverse([buffer | acc])
-    end
+    R.chunk(&chunk_by/4, enumerable, count, step, leftover)
   end
 
   @doc """
@@ -429,18 +419,7 @@ defmodule Enum do
   """
   @spec chunk_by(t, (element -> any)) :: [list]
   def chunk_by(enumerable, fun) do
-    chunk_by(enumerable, nil, fn
-      entry, nil ->
-        {:cont, {[entry], fun.(entry)}}
-      entry, {acc, value} ->
-        case fun.(entry) do
-          ^value -> {:cont, {[entry | acc], value}}
-          new_value -> {:cont, :lists.reverse(acc), {[entry], new_value}}
-        end
-    end, fn
-      nil -> {:cont, :done}
-      {acc, _value} -> {:cont, :lists.reverse(acc), :done}
-    end)
+    R.chunk_by(&chunk_by/4, enumerable, fun)
   end
 
   @doc """

lib/elixir/lib/stream.ex

@@ -161,23 +161,7 @@ defmodule Stream do
   @spec chunk(Enumerable.t, pos_integer, pos_integer, Enumerable.t | nil) :: Enumerable.t
   def chunk(enum, n, step, leftover \\ nil)
       when is_integer(n) and n > 0 and is_integer(step) and step > 0 do
-    limit = :erlang.max(n, step)
-    if is_nil(leftover) do
-      lazy enum, {[], 0}, fn(f1) -> R.chunk(n, step, limit, f1) end
-    else
-      lazy enum, {[], 0},
-           fn(f1) -> R.chunk(n, step, limit, f1) end,
-           &do_chunk(&1, n, leftover, &2)
-     end
-  end
-
-  defp do_chunk(acc(_, {_, 0}, _) = acc, _, _, _) do
-    {:cont, acc}
-  end
-
-  defp do_chunk(acc(h, {buffer, count} = old, t), n, leftover, f1) do
-    buffer = :lists.reverse(buffer, Enum.take(leftover, n - count))
-    next_with_acc(f1, buffer, h, old, t)
+    R.chunk(&chunk_by/4, enum, n, step, leftover)
   end
 
   @doc """
@@ -194,18 +178,7 @@ defmodule Stream do
   """
   @spec chunk_by(Enumerable.t, (element -> any)) :: Enumerable.t
   def chunk_by(enum, fun) do
-    chunk_by(enum, nil, fn
-      entry, nil ->
-        {:cont, {[entry], fun.(entry)}}
-      entry, {acc, value} ->
-        case fun.(entry) do
-          ^value -> {:cont, {[entry | acc], value}}
-          new_value -> {:cont, :lists.reverse(acc), {[entry], new_value}}
-        end
-    end, fn
-      nil -> {:cont, :done}
-      {acc, _value} -> {:cont, :lists.reverse(acc), :done}
-    end)
+    R.chunk_by(&chunk_by/4, enum, fun)
   end
 
   @doc """

lib/elixir/lib/stream/reducers.ex

@@ -1,28 +1,48 @@
 defmodule Stream.Reducers do
-  # Collection of reducers shared by Enum and Stream.
+  # Collection of reducers and utilities shared by Enum and Stream.
   @moduledoc false
 
-  defmacro chunk(amount, step, limit, fun \\ nil) do
-    quote do
-      fn entry, acc(head, {buffer, count}, tail) ->
-        buffer = [entry | buffer]
-        count = count + 1
-
-        new_state =
-          if count >= unquote(limit) do
-            left = count - unquote(step)
-            {Enum.take(buffer, left), left}
-          else
-            {buffer, count}
-          end
+  def chunk(chunk_by, enumerable, count, step, leftover) do
+    limit = :erlang.max(count, step)
+    chunk_by.(enumerable, {[], 0}, fn entry, {acc_buffer, acc_count} ->
+      acc_buffer = [entry | acc_buffer]
+      acc_count = acc_count + 1
 
-        if count == unquote(amount) do
-          next_with_acc(unquote(fun), :lists.reverse(buffer), head, new_state, tail)
+      new_state =
+        if acc_count >= limit do
+          remaining = acc_count - step
+          {Enum.take(acc_buffer, remaining), remaining}
         else
-          skip(acc(head, new_state, tail))
+          {acc_buffer, acc_count}
         end
+
+      if acc_count == count do
+        {:cont, :lists.reverse(acc_buffer), new_state}
+      else
+        {:cont, new_state}
       end
-    end
+    end, fn {acc_buffer, acc_count} ->
+      if is_nil(leftover) || acc_count == 0 do
+        {:cont, []}
+      else
+        {:cont, :lists.reverse(acc_buffer, Enum.take(leftover, count - acc_count)), []}
+      end
+    end)
+  end
+
+  def chunk_by(chunk_by, enumerable, fun) do
+    chunk_by.(enumerable, nil, fn
+      entry, nil ->
+        {:cont, {[entry], fun.(entry)}}
+      entry, {acc, value} ->
+        case fun.(entry) do
+          ^value -> {:cont, {[entry | acc], value}}
+          new_value -> {:cont, :lists.reverse(acc), {[entry], new_value}}
+        end
+    end, fn
+      nil -> {:cont, :done}
+      {acc, _value} -> {:cont, :lists.reverse(acc), :done}
+    end)
   end
 
   defmacro chunk_by(callback, fun \\ nil) do
@@ -38,7 +58,7 @@ defmodule Stream.Reducers do
 
   defmacro dedup(callback, fun \\ nil) do
     quote do
-      fn(entry, acc(head, prev, tail) = acc) ->
+      fn entry, acc(head, prev, tail) = acc ->
         value = unquote(callback).(entry)
         case prev do
           {:value, ^value} -> skip(acc)
@@ -84,7 +104,7 @@ defmodule Stream.Reducers do
 
   defmacro filter(callback, fun \\ nil) do
     quote do
-      fn(entry, acc) ->
+      fn entry, acc ->
         if unquote(callback).(entry) do
           next(unquote(fun), entry, acc)
         else
@@ -96,7 +116,7 @@ defmodule Stream.Reducers do
 
   defmacro filter_map(filter, mapper, fun \\ nil) do
     quote do
-      fn(entry, acc) ->
+      fn entry, acc ->
         if unquote(filter).(entry) do
           next(unquote(fun), unquote(mapper).(entry), acc)
         else
@@ -108,7 +128,7 @@ defmodule Stream.Reducers do
 
   defmacro map(callback, fun \\ nil) do
     quote do
-      fn(entry, acc) ->
+      fn entry, acc ->
         next(unquote(fun), unquote(callback).(entry), acc)
       end
     end
@@ -127,7 +147,7 @@ defmodule Stream.Reducers do
 
   defmacro reject(callback, fun \\ nil) do
     quote do
-      fn(entry, acc) ->
+      fn entry, acc ->
         unless unquote(callback).(entry) do
           next(unquote(fun), entry, acc)
         else
@@ -151,7 +171,7 @@ defmodule Stream.Reducers do
 
   defmacro scan3(callback, fun \\ nil) do
     quote do
-      fn(entry, acc(head, acc, tail)) ->
+      fn entry, acc(head, acc, tail) ->
         value = unquote(callback).(entry, acc)
         next_with_acc(unquote(fun), value, head, value, tail)
       end
@@ -160,7 +180,7 @@ defmodule Stream.Reducers do
 
   defmacro take(fun \\ nil) do
     quote do
-      fn(entry, acc(head, curr, tail) = original) ->
+      fn entry, acc(head, curr, tail) = original ->
         case curr do
           0 ->
             {:halt, original}
@@ -187,7 +207,7 @@ defmodule Stream.Reducers do
 
   defmacro take_while(callback, fun \\ nil) do
     quote do
-      fn(entry, acc) ->
+      fn entry, acc ->
         if unquote(callback).(entry) do
           next(unquote(fun), entry, acc)
         else
@@ -199,7 +219,7 @@ defmodule Stream.Reducers do
 
   defmacro uniq_by(callback, fun \\ nil) do
     quote do
-      fn(entry, acc(head, prev, tail) = original) ->
+      fn entry, acc(head, prev, tail) = original ->
         value = unquote(callback).(entry)
         if Map.has_key?(prev, value) do
           skip(original)
@@ -212,7 +232,7 @@ defmodule Stream.Reducers do
 
   defmacro with_index(fun \\ nil) do
     quote do
-      fn(entry, acc(head, counter, tail)) ->
+      fn entry, acc(head, counter, tail) ->
         next_with_acc(unquote(fun), {entry, counter}, head, counter + 1, tail)
       end
     end