Add Enum.flat_map_reduce/3

Author: José Valim

CHANGELOG.md

@@ -1,6 +1,7 @@
 # v0.12.2-dev
 
 * Enhancements
+  * [Enum] Add `Enum.flat_map_reduce/3`
   * [ExUnit] Support @moduletag in ExUnit cases
   * [Kernel] Improve stacktraces to be relative to the compilation path and include the related application
   * [Stream] Add `Stream.transform/3`

lib/elixir/lib/enum.ex

@@ -766,7 +766,7 @@ defmodule Enum do
   Returns a new collection appending the result of invoking `fun`
   on each corresponding item of `collection`.
 
-  The given function should return a list.
+  The given function should return an enumerable.
 
   ## Examples
 
@@ -777,15 +777,48 @@ defmodule Enum do
       [1, 2, 3, 4, 5, 6]
 
   """
-  @spec flat_map(t, (element -> any)) :: list
-  def flat_map([], _fun), do: []
-
+  @spec flat_map(t, (element -> t)) :: list
   def flat_map(collection, fun) do
     reduce(collection, [], fn(entry, acc) ->
       reduce(fun.(entry), acc, &[&1|&2])
     end) |> :lists.reverse
   end
 
+  @doc """
+  Maps and reduces a collection, flattening the given results.
+
+  It expects an accumulator and a function that receives each stream item
+  and an accumulator, and must return a tuple containing a new stream
+  (often a list) with the new accumulator or a tuple with `:halt` as first
+  element and the accumulator as second.
+
+  ## Examples
+
+      iex> enum = 1..100
+      iex> n = 3
+      iex> Enum.flat_map_reduce(enum, 0, fn i, acc ->
+      ...>   if acc < n, do: { [i], acc + 1 }, else: { :halt, acc }
+      ...> end)
+      { [1,2,3], 3 }
+
+  """
+  @spec flat_map_reduce(t, acc, fun) :: list when
+        fun: (element, acc -> { t, acc } | { :halt, acc }),
+        acc: any
+  def flat_map_reduce(collection, acc, fun) do
+    { _, { list, acc } } =
+      Enumerable.reduce(collection, { :cont, { [], acc } }, fn(entry, { list, acc }) ->
+        case fun.(entry, acc) do
+          { :halt, acc } ->
+            { :halt, { list, acc } }
+          { entries, acc } ->
+            { :cont, { reduce(entries, list, &[&1|&2]), acc } }
+        end
+      end)
+
+    { :lists.reverse(list), acc }
+  end
+
   @doc """
   Intersperses `element` between each element of the enumeration.
 

lib/elixir/lib/stream.ex

@@ -602,7 +602,12 @@ defmodule Stream do
 
   It expects an accumulator and a function that receives each stream item
   and an accumulator, and must return a tuple containing a new stream
-  (often a list) with the new accumulator or simply return the atom `:halt`.
+  (often a list) with the new accumulator or a tuple with `:halt` as first
+  element and the accumulator as second.
+
+  Note: this function is similar to `Enum.flat_map_reduce/3` except the
+  latter returns both the flat list and accumulator, while this one returns
+  only the stream.
 
   ## Examples
 
@@ -613,14 +618,14 @@ defmodule Stream do
       iex> enum = 1..100
       iex> n = 3
       iex> stream = Stream.transform(enum, 0, fn i, acc ->
-      ...>   if acc < n, do: { [i], acc + 1 }, else: :halt
+      ...>   if acc < n, do: { [i], acc + 1 }, else: { :halt, acc }
       ...> end)
       iex> Enum.to_list(stream)
       [1,2,3]
 
   """
   @spec transform(Enumerable.t, acc, fun) :: Enumerable.t when
-        fun: (element, acc -> { Enumerable.t, acc } | :halt),
+        fun: (element, acc -> { Enumerable.t, acc } | { :halt, acc }),
         acc: any
   def transform(enum, acc, reducer) do
     &do_transform(enum, acc, reducer, &1, &2)
@@ -647,11 +652,11 @@ defmodule Stream do
             do_transform(user_acc, user, fun, next_acc, next, inner_acc, inner)
           { list, user_acc } when is_list(list) ->
             do_list_transform(user_acc, user, fun, next_acc, next, inner_acc, inner, &Enumerable.List.reduce(list, &1, fun))
-          { other, user_acc } ->
-            do_other_transform(user_acc, user, fun, next_acc, next, inner_acc, inner, &Enumerable.reduce(other, &1, inner))
-          :halt ->
+          { :halt, _user_acc } ->
             next.({ :halt, next_acc })
             { :halted, elem(inner_acc, 1) }
+          { other, user_acc } ->
+            do_other_transform(user_acc, user, fun, next_acc, next, inner_acc, inner, &Enumerable.reduce(other, &1, inner))
         end
       { reason, _ } ->
         { reason, elem(inner_acc, 1) }

lib/elixir/test/elixir/enum_test.exs

@@ -170,6 +170,15 @@ defmodule EnumTest.List do
     assert Enum.flat_map([1, 2, 3], fn(x) -> x..x+1 end) == [1, 2, 2, 3, 3, 4]
   end
 
+  test :flat_map_reduce do
+    assert Enum.flat_map_reduce([1, 2, 3], 0, &{ [&1, &2], &1 + &2 }) ==
+           { [1, 0, 2, 1, 3, 3], 6 }
+
+    assert Enum.flat_map_reduce(1..100, 0, fn i, acc ->
+      if acc < 3, do: { [i], acc + 1 }, else: { :halt, acc }
+    end) == { [1,2,3], 3 }
+  end
+
   test :reduce do
     assert Enum.reduce([], 1, fn(x, acc) -> x + acc end) == 1
     assert Enum.reduce([1, 2, 3], 1, fn(x, acc) -> x + acc end) == 7

lib/elixir/test/elixir/stream_test.exs

@@ -313,7 +313,7 @@ defmodule StreamTest do
     stream = Stream.resource(fn -> 1 end,
                              fn acc -> { acc, acc + 1 } end,
                              fn _ -> Process.put(:stream_transform, true) end)
-    stream = Stream.transform(stream, 0, fn i, acc -> if acc < 3, do: { [i], acc + 1 }, else: :halt end)
+    stream = Stream.transform(stream, 0, fn i, acc -> if acc < 3, do: { [i], acc + 1 }, else: { :halt, acc } end)
 
     Process.put(:stream_transform, false)
     assert Enum.to_list(stream) == [1,2,3]