Make the reducer function part of the suspended continuation

Author: José Valim

lib/elixir/lib/enum.ex

@@ -80,10 +80,9 @@ defprotocol Enumerable do
 
   The continuation is the closure returned as result when
   the enumeration is suspended. When invoked, it expects
-  an accumulator and a reducer function as argument, it returns
-  the result.
+  a new accumulator and it returns the result.
   """
-  @type continuation :: (acc, reducer -> result)
+  @type continuation :: (acc -> result)
 
   @doc """
   The reduce function.
@@ -94,10 +93,10 @@ defprotocol Enumerable do
 
   As an example, here is the implementation of `reduce` for lists:
 
-      def reduce(_,     { :halt, acc }, _fun),    do: { :halted, acc }
-      def reduce(list,  { :suspend, acc }, _fun), do: { :suspended, acc, &reduce(list, &1, &2) }
-      def reduce([],    { :cont, acc }, _fun),    do: { :done, acc }
-      def reduce([h|t], { :cont, acc }, fun),     do: reduce(t, fun.(h, acc), fun)
+      def reduce(_,     { :halt, acc }, _fun),   do: { :halted, acc }
+      def reduce(list,  { :suspend, acc }, fun), do: { :suspended, acc, &reduce(list, &1, fun) }
+      def reduce([],    { :cont, acc }, _fun),   do: { :done, acc }
+      def reduce([h|t], { :cont, acc }, fun),    do: reduce(t, fun.(h, acc), fun)
 
   """
   @spec reduce(t, acc, reducer) :: result
@@ -905,17 +904,16 @@ defmodule Enum do
 
   def reduce(collection, fun) do
     result =
-      Enumerable.reduce(collection, { :cont, nil }, fn x, _ ->
-        { :suspend, x }
-      end)
+      Enumerable.reduce(collection, { :cont, :first }, fn
+        x, :first ->
+          { :cont, { :acc, x } }
+        x, { :acc, acc } ->
+          { :cont, { :acc, fun.(x, acc) } }
+      end) |> elem(1)
 
     case result do
-      { :suspended, x, continuation } ->
-        continuation.({ :cont, x }, fn x, acc ->
-          { :cont, fun.(x, acc) }
-        end) |> elem(1)
-      { _, _ } ->
-        raise Enum.EmptyError
+      :first        -> raise Enum.EmptyError
+      { :acc, acc } -> acc
     end
   end
 
@@ -1957,10 +1955,10 @@ defmodule Enum do
 end
 
 defimpl Enumerable, for: List do
-  def reduce(_,     { :halt, acc }, _fun),    do: { :halted, acc }
-  def reduce(list,  { :suspend, acc }, _fun), do: { :suspended, acc, &reduce(list, &1, &2) }
-  def reduce([],    { :cont, acc }, _fun),    do: { :done, acc }
-  def reduce([h|t], { :cont, acc }, fun),     do: reduce(t, fun.(h, acc), fun)
+  def reduce(_,     { :halt, acc }, _fun),   do: { :halted, acc }
+  def reduce(list,  { :suspend, acc }, fun), do: { :suspended, acc, &reduce(list, &1, fun) }
+  def reduce([],    { :cont, acc }, _fun),   do: { :done, acc }
+  def reduce([h|t], { :cont, acc }, fun),    do: reduce(t, fun.(h, acc), fun)
 
   def member?([], _),       do: false
   def member?(list, value), do: :lists.member(value, list)

lib/elixir/lib/hash_dict.ex

@@ -330,7 +330,7 @@ defmodule HashDict do
 
   # Reduces the bucket
   defp bucket_reduce(_, { :halt, acc }, _fun),        do: { :halted, acc }
-  defp bucket_reduce(list, { :suspend, acc }, _fun),  do: { :suspended, acc, &bucket_reduce(list, &1, &2) }
+  defp bucket_reduce(list, { :suspend, acc }, fun),   do: { :suspended, acc, &bucket_reduce(list, &1, fun) }
   defp bucket_reduce([[k|v]|t], { :cont, acc }, fun), do: bucket_reduce(t, fun.({ k, v }, acc), fun)
   defp bucket_reduce([], { :cont, acc }, _fun),       do: { :done, acc }
 
@@ -401,8 +401,8 @@ defmodule HashDict do
     { :halted, acc }
   end
 
-  defp node_reduce(list, -1, { :suspend, acc }, _fun, count, next) do
-    { :suspended, acc, &node_reduce(list, -1, &1, &2, count, next) }
+  defp node_reduce(list, -1, { :suspend, acc }, fun, count, next) do
+    { :suspend, acc, &node_reduce(list, -1, &1, fun, count, next) }
   end
 
   defp node_reduce([[k|v]|t], -1, { :cont, acc }, fun, _count, next) do

lib/elixir/lib/hash_set.ex

@@ -304,10 +304,10 @@ defmodule HashSet do
     { [], 0 }
   end
 
-  defp bucket_reduce(_, { :halt, acc }, _fun),       do: { :halted, acc }
-  defp bucket_reduce(list, { :suspend, acc }, _fun), do: { :suspended, acc, &bucket_reduce(list, &1, &2) }
-  defp bucket_reduce([h|t], { :cont, acc }, fun),    do: bucket_reduce(t, fun.(h, acc), fun)
-  defp bucket_reduce([], { :cont, acc }, _fun),      do: { :done, acc }
+  defp bucket_reduce(_, { :halt, acc }, _fun),      do: { :halted, acc }
+  defp bucket_reduce(list, { :suspend, acc }, fun), do: { :suspended, acc, &bucket_reduce(list, &1, fun) }
+  defp bucket_reduce([h|t], { :cont, acc }, fun),   do: bucket_reduce(t, fun.(h, acc), fun)
+  defp bucket_reduce([], { :cont, acc }, _fun),     do: { :done, acc }
 
   defp bucket_fold(bucket, acc, fun) do
     :lists.foldl(fun, acc, bucket)
@@ -399,8 +399,8 @@ defmodule HashSet do
     { :halted, acc }
   end
 
-  defp node_reduce(list, -1, { :suspend, acc }, _fun, count, next) do
-    { :suspended, acc, &node_reduce(list, -1, &1, &2, count, next) }
+  defp node_reduce(list, -1, { :suspend, acc }, fun, count, next) do
+    { :suspended, acc, &node_reduce(list, -1, &1, fun, count, next) }
   end
 
   defp node_reduce([h|t], -1, { :cont, acc }, fun, _count, next) do

lib/elixir/lib/range.ex

@@ -25,8 +25,8 @@ defimpl Enumerable, for: Range do
     { :halted, acc }
   end
 
-  defp reduce(x, y, { :suspend, acc }, _fun, next, up) do
-    { :suspended, acc, &reduce(x, y, &1, &2, next, up) }
+  defp reduce(x, y, { :suspend, acc }, fun, next, up) do
+    { :suspended, acc, &reduce(x, y, &1, fun, next, up) }
   end
 
   defp reduce(x, y, { :cont, acc }, fun, next, true) when x <= y do

lib/elixir/lib/stream.ex

@@ -106,25 +106,24 @@ defmodule Stream do
     end
 
     defp do_reduce(Lazy[enum: enum, funs: funs, accs: accs], acc, fun) do
-      do_reduce(&Enumerable.reduce(enum, &1, &2), :lists.reverse(accs), funs, acc, fun)
+      composed = :lists.foldl(fn fun, acc -> fun.(acc) end, fun, funs)
+      do_each(&Enumerable.reduce(enum, &1, composed), :lists.reverse(accs), acc)
     end
 
-    defp do_reduce(_reduce, _accs, _funs, { :halt, acc }, _fun) do
+    defp do_each(_reduce, _accs, { :halt, acc }) do
       { :halted, acc }
     end
 
-    defp do_reduce(reduce, accs, funs, { :suspend, acc }, _fun) do
-      { :suspended, acc, &do_reduce(reduce, accs, funs, &1, &2) }
+    defp do_each(reduce, accs, { :suspend, acc }) do
+      { :suspended, acc, &do_each(reduce, accs, &1) }
     end
 
-    defp do_reduce(reduce, accs, funs, { :cont, acc }, fun) do
-      composed = :lists.foldl(fn fun, acc -> fun.(acc) end, fun, funs)
-
-      case reduce.({ :cont, [acc|accs] }, composed) do
+    defp do_each(reduce, accs, { :cont, acc }) do
+      case reduce.({ :cont, [acc|accs] }) do
         { reason, [acc|_] } ->
           { reason, acc }
         { :suspended, [acc|accs], continuation } ->
-          { :suspended, acc, &do_reduce(continuation, accs, funs, &1, &2) }
+          { :suspended, acc, &do_each(continuation, accs, &1) }
       end
     end
   end
@@ -241,20 +240,21 @@ defmodule Stream do
     lazy enum, fn(f1) ->
       fn(entry, acc) ->
         enum = f.(entry)
-        do_flat_map(&Enumerable.reduce(enum, &1, &2), { :cont, acc }, f1)
+        fun  = &do_flat_map_each(f1, &1, &2)
+        do_flat_map(&Enumerable.reduce(enum, &1, fun), { :cont, acc })
       end
     end
   end
 
-  defp do_flat_map(reduce, acc, f1) do
+  defp do_flat_map(reduce, acc) do
     try do
-      reduce.(acc, &do_flat_map_each(f1, &1, &2))
+      reduce.(acc)
     catch
       { :stream_flat_map, acc } -> acc
     else
       { :done, acc }         -> { :cont, acc }
       { :halted, acc }       -> { :cont, acc }
-      { :suspended, acc, c } -> { :suspend, acc, &do_flat_map(c, &1, &2) }
+      { :suspended, acc, c } -> { :suspend, acc, &do_flat_map(c, &1) }
     end
   end
 
@@ -400,20 +400,21 @@ defmodule Stream do
   end
 
   defp do_concat(enumerables, acc, fun) do
+    fun = &do_concat_each(fun, &1, &2)
     Enumerable.reduce(enumerables, acc, fn x, acc ->
-      do_concat_reduce(&Enumerable.reduce(x, &1, &2), { :cont, acc }, fun)
+      do_concat_reduce(&Enumerable.reduce(x, &1, fun), { :cont, acc })
     end)
   end
 
-  defp do_concat_reduce(reduce, acc, fun) do
+  defp do_concat_reduce(reduce, acc) do
     try do
-      reduce.(acc, &do_concat_each(fun, &1, &2))
+      reduce.(acc)
     catch
       { :stream_concat, acc } -> acc
     else
       { :done, acc }         -> { :cont, acc }
       { :halted, acc }       -> { :cont, acc }
-      { :suspended, acc, c } -> { :suspend, acc, &do_concat_reduce(c, &1, &2) }
+      { :suspended, acc, c } -> { :suspend, acc, &do_concat_reduce(c, &1) }
     end
   end
 
@@ -439,26 +440,28 @@ defmodule Stream do
   """
   @spec cycle(Enumerable.t) :: t
   def cycle(enumerable) do
-    reduce = &Enumerable.reduce(enumerable, &1, &2)
-    &do_cycle(reduce, reduce, &1, &2)
+    fn acc, fun ->
+      reduce = &Enumerable.reduce(enumerable, &1, fun)
+      do_cycle(reduce, reduce, acc)
+    end
   end
 
-  defp do_cycle(_reduce, _cycle, { :halt, acc }, _fun) do
+  defp do_cycle(_reduce, _cycle, { :halt, acc }) do
     { :halted, acc }
   end
 
-  defp do_cycle(reduce, cycle, { :suspend, acc }, _fun) do
-    { :suspended, acc, &do_cycle(reduce, cycle, &1, &2) }
+  defp do_cycle(reduce, cycle, { :suspend, acc }) do
+    { :suspended, acc, &do_cycle(reduce, cycle, &1) }
   end
 
-  defp do_cycle(reduce, cycle, acc, fun) do
-    case reduce.(acc, fun) do
+  defp do_cycle(reduce, cycle, acc) do
+    case reduce.(acc) do
       { :done, acc } ->
-        do_cycle(cycle, cycle, { :cont, acc }, fun)
+        do_cycle(cycle, cycle, { :cont, acc })
       { :halted, acc } ->
         { :halted, acc }
       { :suspended, acc, continuation } ->
-        { :suspended, acc, &do_cycle(continuation, cycle, &1, &2) }
+        { :suspended, acc, &do_cycle(continuation, cycle, &1) }
     end
   end
 
@@ -494,7 +497,19 @@ defmodule Stream do
   """
   @spec repeatedly((() -> element)) :: t
   def repeatedly(generator_fun) when is_function(generator_fun, 0) do
-    unfold(:repeat, fn :repeat -> { generator_fun.(), :repeat } end)
+    &do_repeatedly(generator_fun, &1, &2)
+  end
+
+  defp do_repeatedly(generator_fun, { :suspend, acc }, fun) do
+    { :suspended, acc, &do_repeatedly(generator_fun, &1, fun) }
+  end
+
+  defp do_repeatedly(_generator_fun, { :halt, acc }, _fun) do
+    { :halted, acc }
+  end
+
+  defp do_repeatedly(generator_fun, { :cont, acc }, fun) do
+    do_repeatedly(generator_fun, fun.(generator_fun.(), acc), fun)
   end
 
   @doc """
@@ -553,8 +568,8 @@ defmodule Stream do
     &do_unfold(next_acc, next_fun, &1, &2)
   end
 
-  defp do_unfold(next_acc, next_fun, { :suspend, acc }, _fun) do
-    { :suspended, acc, &do_unfold(next_acc, next_fun, &1, &2) }
+  defp do_unfold(next_acc, next_fun, { :suspend, acc }, fun) do
+    { :suspended, acc, &do_unfold(next_acc, next_fun, &1, fun) }
   end
 
   defp do_unfold(_next_acc, _next_fun, { :halt, acc }, _fun) do