Add Stream.scan/2 and Stream.scan/3

Author: José Valim

lib/elixir/lib/enum.ex

@@ -1011,7 +1011,6 @@ defmodule Enum do
     end
   end
 
-
   @doc """
   Returns the minimum value.
   Raises `EmptyError` if the collection is empty.
@@ -1205,6 +1204,42 @@ defmodule Enum do
     end)
   end
 
+  @doc """
+  Applies the given function to each element in the collection,
+  storing the result in a list and passing it as accumulator
+  for the next computation.
+
+  ## Examples
+
+      iex> Enum.scan(1..5, &(&1 + &2))
+      [1,3,6,10,15]
+
+  """
+  @spec scan(t, (element, any -> any)) :: list
+  def scan(enum, fun) do
+    { _, { res, _ } } =
+      Enumerable.reduce(enum, { :cont, { [], :first } }, R.scan_2(fun))
+    :lists.reverse(res)
+  end
+
+  @doc """
+  Applies the given function to each element in the collection,
+  storing the result in a list and passing it as accumulator
+  for the next computation. Uses the given `acc` as starting value.
+
+  ## Examples
+
+      iex> Enum.scan(1..5, 0, &(&1 + &2))
+      [1,3,6,10,15]
+
+  """
+  @spec scan(t, any, (element, any -> any)) :: list
+  def scan(enum, acc, fun) do
+    { _, { res, _ } } =
+      Enumerable.reduce(enum, { :cont, { [], acc } }, R.scan_3(fun))
+    :lists.reverse(res)
+  end
+
   @doc """
   Returns a list of collection elements shuffled.
 

lib/elixir/lib/stream.ex

@@ -377,6 +377,40 @@ defmodule Stream do
     lazy enum, fn(f1) -> R.map(fun, f1) end
   end
 
+  @doc """
+  Creates a stream that applies the given function to each
+  element, emits the result and uses the same result as accumulator
+  for the next computation.
+
+  ## Examples
+
+      iex> stream = Stream.scan(1..5, &(&1 + &2))
+      iex> Enum.to_list(stream)
+      [1,3,6,10,15]
+
+  """
+  @spec scan(Enumerable.t, (element, acc -> any)) :: Enumerable.t
+  def scan(enum, fun) do
+    lazy enum, :first, fn(f1) -> R.scan_2(fun, f1) end
+  end
+
+  @doc """
+  Creates a stream that applies the given function to each
+  element, emits the result and uses the same result as accumulator
+  for the next computation. Uses the given `acc` as starting value.
+
+  ## Examples
+
+      iex> stream = Stream.scan(1..5, 0, &(&1 + &2))
+      iex> Enum.to_list(stream)
+      [1,3,6,10,15]
+
+  """
+  @spec scan(Enumerable.t, acc, (element, acc -> any)) :: Enumerable.t
+  def scan(enum, acc, fun) do
+    lazy enum, acc, fn(f1) -> R.scan_3(fun, f1) end
+  end
+
   @doc """
   Creates a stream that will apply the given function on enumeration and
   flatten the result.

lib/elixir/lib/stream/reducers.ex

@@ -108,6 +108,27 @@ defmodule Stream.Reducers do
     end
   end
 
+  defmacro scan_2(callback, f // nil) do
+    quote do
+      fn
+        entry, acc(h, :first, t) ->
+          cont_with_acc(unquote(f), entry, h, { :ok, entry }, t)
+        entry, acc(h, { :ok, acc }, t) ->
+          value = unquote(callback).(entry, acc)
+          cont_with_acc(unquote(f), value, h, { :ok, value }, t)
+      end
+    end
+  end
+
+  defmacro scan_3(callback, f // nil) do
+    quote do
+      fn(entry, acc(h, acc, t)) ->
+        value = unquote(callback).(entry, acc)
+        cont_with_acc(unquote(f), value, h, value, t)
+      end
+    end
+  end
+
   defmacro take(f // nil) do
     quote do
       fn(entry, acc(h, n, t) = orig) ->

lib/elixir/test/elixir/enum_test.exs

@@ -231,6 +231,14 @@ defmodule EnumTest.List do
     assert Enum.reverse([1, 2, 3], [4, 5, 6]) == [3, 2, 1, 4, 5, 6]
   end
 
+  test :scan do
+    assert Enum.scan([1,2,3,4,5], &(&1 + &2)) == [1,3,6,10,15]
+    assert Enum.scan([], &(&1 + &2)) == []
+
+    assert Enum.scan([1,2,3,4,5], 0, &(&1 + &2)) == [1,3,6,10,15]
+    assert Enum.scan([], 0, &(&1 + &2)) == []
+  end
+
   test :shuffle do
     # set a fixed seed so the test can be deterministic
     :random.seed(1374, 347975, 449264)
@@ -636,6 +644,11 @@ defmodule EnumTest.Range do
     assert Enum.partition(range, fn(x) -> rem(x, 2) == 0 end) == { [2], [1, 3] }
   end
 
+  test :scan do
+    assert Enum.scan(1..5, &(&1 + &2)) == [1,3,6,10,15]
+    assert Enum.scan(1..5, 0, &(&1 + &2)) == [1,3,6,10,15]
+  end
+
   test :shuffle do
     # set a fixed seed so the test can be deterministic
     :random.seed(1374, 347975, 449264)

lib/elixir/test/elixir/stream_test.exs

@@ -280,6 +280,20 @@ defmodule StreamTest do
     assert r1 != r2
   end
 
+  test "scan/2" do
+    stream = Stream.scan(1..5, &(&1 + &2))
+    assert is_lazy(stream)
+    assert Enum.to_list(stream) == [1,3,6,10,15]
+    assert Stream.scan([], &(&1 + &2)) |> Enum.to_list == []
+  end
+
+  test "scan/3" do
+    stream = Stream.scan(1..5, 0, &(&1 + &2))
+    assert is_lazy(stream)
+    assert Enum.to_list(stream) == [1,3,6,10,15]
+    assert Stream.scan([], 0, &(&1 + &2)) |> Enum.to_list == []
+  end
+
   test "take" do
     stream = Stream.take(1..1000, 5)
     assert is_lazy(stream)