Merge pull request #1652 from vanstee/stream-flat-map

Implement `Stream.flat_map/2`

Author: José Valim

lib/elixir/lib/stream.ex

@@ -266,6 +266,28 @@ defmodule Stream do
          end]
   end
 
+  @doc """
+  Creates a stream that will apply the given function on enumeration and
+  flatten the result.
+
+  ## Examples
+
+      iex> stream = Stream.flat_map([1, 2, 3], fn(x) -> [x, x * 2] end)
+      iex> Enum.to_list(stream)
+      [1, 2, 2, 4, 3, 6]
+
+  """
+
+  @spec flat_map(Enumerable.t, (element -> any)) :: t
+  def flat_map(enumerable, f) do
+    Lazy[enumerable: enumerable,
+         fun: fn(f1) ->
+           fn(entry, acc) ->
+             Enumerable.reduce(f.(entry), acc, f1)
+           end
+         end]
+  end
+
   @doc """
   Creates a stream that will reject elements according to
   the given function on enumeration.

lib/elixir/test/elixir/stream_test.exs

@@ -86,6 +86,15 @@ defmodule StreamTest do
     assert Stream.map(nats, &(&1 * 2)) |> Enum.take(5) == [2,4,6,8,10]
   end
 
+  test :flat_map do
+    stream = Stream.flat_map([1, 2, 3], &[&1, &1 * 2])
+    assert is_lazy(stream)
+    assert Enum.to_list(stream) == [1, 2, 2, 4, 3, 6]
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Stream.flat_map(nats, &[&1, &1 * 2]) |> Enum.take(6) == [1, 2, 2, 4, 3, 6]
+  end
+
   test :reject do
     stream = Stream.reject([1,2,3], fn(x) -> rem(x, 2) == 0 end)
     assert is_lazy(stream)