Add Stream.chunks_by/2

Author: José Valim

lib/elixir/lib/enum.ex

@@ -340,24 +340,14 @@ defmodule Enum do
   """
   @spec chunks_by(t, (element -> any)) :: [list]
   def chunks_by(coll, fun) do
-    res =
-      reduce(coll, nil, fn
-        x, {acc, buffer, value} ->
-          new_value = fun.(x)
-          if new_value == value do
-            { acc, [x | buffer], value }
-          else
-            { [:lists.reverse(buffer) | acc], [x], new_value }
-          end
-        x, nil ->
-          { [], [x], fun.(x) }
-      end)
+    { _, { acc, res } } =
+      Enumerable.reduce(coll, { :cont, { [], nil } }, R.chunks_by(fun))
 
-    if nil?(res) do
-      []
-    else
-      { acc, buffer, _ } = res
-      :lists.reverse([:lists.reverse(buffer) | acc])
+    case res do
+      { buffer, _ } ->
+        :lists.reverse([:lists.reverse(buffer) | acc])
+      nil ->
+        []
     end
   end
 

lib/elixir/lib/stream.ex

@@ -83,7 +83,7 @@ defmodule Stream do
   like `Stream.cycle/1`, `Stream.unfold/2`, `Stream.resource/3` and more.
   """
 
-  defrecord Lazy, enum: nil, funs: [], accs: []
+  defrecord Lazy, enum: nil, funs: [], accs: [], done: nil
 
   defimpl Enumerable, for: Lazy do
     @compile :inline_list_funs
@@ -103,25 +103,36 @@ defmodule Stream do
       { :error, __MODULE__ }
     end
 
-    defp do_reduce(Lazy[enum: enum, funs: funs, accs: accs], acc, fun) do
+    defp do_reduce(Lazy[enum: enum, funs: funs, accs: accs, done: done], acc, fun) do
       composed = :lists.foldl(fn fun, acc -> fun.(acc) end, fun, funs)
-      do_each(&Enumerable.reduce(enum, &1, composed), :lists.reverse(accs), acc)
+      do_each(&Enumerable.reduce(enum, &1, composed), done && { done, fun }, :lists.reverse(accs), acc)
     end
 
-    defp do_each(_reduce, _accs, { :halt, acc }) do
+    defp do_each(_reduce, _done, _accs, { :halt, acc }) do
       { :halted, acc }
     end
 
-    defp do_each(reduce, accs, { :suspend, acc }) do
-      { :suspended, acc, &do_each(reduce, accs, &1) }
+    defp do_each(reduce, done, accs, { :suspend, acc }) do
+      { :suspended, acc, &do_each(reduce, done, accs, &1) }
     end
 
-    defp do_each(reduce, accs, { :cont, acc }) do
+    defp do_each(reduce, done, accs, { :cont, acc }) do
       case reduce.({ :cont, [acc|accs] }) do
-        { reason, [acc|_] } ->
-          { reason, acc }
         { :suspended, [acc|accs], continuation } ->
-          { :suspended, acc, &do_each(continuation, accs, &1) }
+          { :suspended, acc, &do_each(continuation, done, accs, &1) }
+        { :halted, [acc|_] } ->
+          { :halted, acc }
+        { :done, [acc|_] = accs } ->
+          case done do
+            nil ->
+              { :done, acc }
+            { done, fun } ->
+              case done.(fun).(accs) do
+                { :cont, [acc|_] }    -> { :done, acc }
+                { :halt, [acc|_] }    -> { :halted, acc }
+                { :suspend, [acc|_] } -> { :suspended, acc, &({ :done, &1 |> elem(1) }) }
+              end
+          end
       end
     end
   end
@@ -151,6 +162,33 @@ defmodule Stream do
 
   ## Transformers
 
+  @doc """
+  Chunk 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,
+
+  ## Examples
+
+      iex> stream = Stream.chunks_by([1, 2, 2, 3, 4, 4, 6, 7, 7], &(rem(&1, 2) == 1))
+      iex> Enum.to_list(stream)
+      [[1], [2, 2], [3], [4, 4, 6], [7, 7]]
+
+  """
+  @spec chunks_by(Enumerable.t, (element -> any)) :: Enumerable.t
+  def chunks_by(enum, fun) do
+    lazy enum, nil,
+         fn(f1) -> R.chunks_by(fun, f1) end,
+         fn(f1) -> &do_chunks_by(&1, f1) end
+  end
+
+  defp do_chunks_by(acc(_, nil, _) = acc, _f1) do
+    { :cont, acc }
+  end
+
+  defp do_chunks_by(acc(h, { buffer, _ }, t), f1) do
+    cont_with_acc(f1, :lists.reverse(buffer), h, nil, t)
+  end
+
   @doc """
   Lazily drops the next `n` items from the enumerable.
 
@@ -720,23 +758,20 @@ defmodule Stream do
 
   ## Helpers
 
-  @compile { :inline, lazy: 2, lazy: 3 }
+  @compile { :inline, lazy: 2, lazy: 3, lazy: 4 }
 
-  defp lazy(enum, fun) do
-    case enum do
-      Lazy[funs: funs] = lazy ->
-        lazy.funs([fun|funs])
-      _ ->
-        Lazy[enum: enum, funs: [fun], accs: []]
-    end
-  end
+  defp lazy(Lazy[funs: funs] = lazy, fun),
+    do: lazy.funs([fun|funs])
+  defp lazy(enum, fun),
+    do: Lazy[enum: enum, funs: [fun]]
 
-  defp lazy(enum, acc, fun) do
-    case enum do
-      Lazy[funs: funs, accs: accs] = lazy ->
-        lazy.funs([fun|funs]).accs([acc|accs])
-      _ ->
-        Lazy[enum: enum, funs: [fun], accs: [acc]]
-    end
-  end
+  defp lazy(Lazy[funs: funs, accs: accs] = lazy, acc, fun),
+    do: lazy.funs([fun|funs]).accs([acc|accs])
+  defp lazy(enum, acc, fun),
+    do: Lazy[enum: enum, funs: [fun], accs: [acc]]
+
+  defp lazy(Lazy[done: nil, funs: funs, accs: accs] = lazy, acc, fun, done),
+    do: lazy.funs([fun|funs]).accs([acc|accs]).done(done)
+  defp lazy(enum, acc, fun, done),
+    do: Lazy[enum: enum, funs: [fun], accs: [acc], done: done]
 end

lib/elixir/lib/stream/reducers.ex

@@ -2,6 +2,22 @@ defmodule Stream.Reducers do
   # Collection of reducers shared by Enum and Stream.
   @moduledoc false
 
+  defmacro chunks_by(fun, f // nil) do
+    quote do
+      fn
+        entry, acc(h, { buffer, value }, t) ->
+          new_value = unquote(fun).(entry)
+          if new_value == value do
+            { :cont, acc(h, { [entry|buffer], value }, t) }
+          else
+            cont_with_acc(unquote(f), :lists.reverse(buffer), h, { [entry], new_value }, t)
+          end
+        entry, acc(h, nil, t) ->
+          { :cont, acc(h, { [entry], unquote(fun).(entry) }, t) }
+      end
+    end
+  end
+
   defmacro drop(f // nil) do
     quote do
       fn

lib/elixir/test/elixir/stream_test.exs

@@ -25,6 +25,22 @@ defmodule StreamTest do
     assert Enum.to_list(stream) == [3,5,7]
   end
 
+  test "chunks_by" do
+    stream = Stream.chunks_by([1, 2, 2, 3, 4, 4, 6, 7, 7], &(rem(&1, 2) == 1))
+
+    assert is_lazy(stream)
+    assert Enum.to_list(stream) ==
+           [[1], [2, 2], [3], [4, 4, 6], [7, 7]]
+    assert stream |> Stream.take(3) |> Enum.to_list ==
+           [[1], [2, 2], [3]]
+  end
+
+  test "chunks_by is zippable" do
+    stream = Stream.chunks_by([1, 2, 2, 3], &(rem(&1, 2) == 1))
+    list   = Enum.to_list(stream)
+    assert Enum.zip(list, list) == Enum.zip(stream, stream)
+  end
+
   test "concat/1" do
     stream = Stream.concat([1..3, [], [4, 5, 6], [], 7..9])
     assert is_function(stream)