Merge pull request #1693 from pminten/stream-generate

Add Stream.unfold

Author: José Valim

lib/elixir/lib/stream.ex

@@ -427,6 +427,31 @@ defmodule Stream do
          acc: true]
   end
 
+  @doc """
+  Emit a sequence of values and accumulators. Successive values are generated by
+  calling `next_fun` with the previous accumulator.
+
+  If the return value is nil iteration ends.
+
+  ## Examples
+
+      iex> Stream.unfold(5, fn 0 -> nil; n -> {n, n-1} end) |> Enum.to_list()
+      [5, 4, 3, 2, 1]
+  """
+  @spec unfold(acc, (acc -> { element, acc } | nil)) :: t
+  def unfold(acc, f) do
+    fn acc1, f1 ->
+      do_unfold(acc, f, acc1, f1)
+    end
+  end
+
+  defp do_unfold(gen_acc, gen_fun, acc, fun) do
+    case gen_fun.(gen_acc) do
+      nil                -> acc
+      { v, new_gen_acc } -> do_unfold(new_gen_acc, gen_fun, fun.(v, acc), fun)
+    end
+  end
+
   @doc """
   Creates a stream where each item in the enumerable will
   be accompanied by its index.

lib/elixir/test/elixir/stream_test.exs

@@ -175,6 +175,19 @@ defmodule StreamTest do
     stream = Stream.drop(1..100, 5)
     assert Stream.take_while(stream, &(&1 < 11)) |> Enum.to_list == [6,7,8,9,10]
   end
+  
+  test :unfold do
+    stream = Stream.unfold(10, fn x -> if x > 0, do: {x, x-1}, else: nil end)
+    assert Enum.take(stream, 5) == [10, 9, 8, 7, 6]
+    stream = Stream.unfold(5, fn x -> if x > 0, do: {x, x-1}, else: nil end)
+    assert Enum.to_list(stream) == [5, 4, 3, 2, 1]
+
+    # Only calculate values if needed
+    stream = Stream.unfold(1, fn x -> if x > 0, do: {x, x-1}, else: throw(:boom) end)
+    assert Enum.take(stream, 1) == [1]
+    stream = Stream.unfold(5, fn x -> if x > 0, do: {x, x-1}, else: nil end)
+    assert Enum.to_list(Stream.take(stream, 2)) == [5, 4]
+  end
 
   test :with_index do
     stream = Stream.with_index([1,2,3])