Merge pull request #1521 from ericmj/stream-fixes

Fix Stream bugs and add more composing tests

Author: José Valim

lib/elixir/lib/stream.ex

@@ -27,7 +27,7 @@ defmodule Stream do
   We say the functions in `Stream` are *lazy* and the functions in `Enum`
   are *eager*.
 
-  Due to their laziness, streams are useful when working with large 
+  Due to their laziness, streams are useful when working with large
   (or even infinite) collections. When chaining many operations with `Enum`,
   intermediate lists are created, while `Stream` creates a recipe of
   computations that are executed at a later moment. Let's see another
@@ -298,7 +298,7 @@ defmodule Stream do
 
   """
   @spec repeatedly((() -> element)) :: t
-  def repeatedly(generator_fun) 
+  def repeatedly(generator_fun)
   when is_function(generator_fun, 0) do
     do_repeatedly(generator_fun, &1, &2)
   end
@@ -323,14 +323,14 @@ defmodule Stream do
 
   """
   @spec take(Enumerable.t, non_neg_integer) :: t
-  def take(enumerable, n) when n >= 0 do
+  def take(_enumerable, 0), do: Lazy[enumerable: [], fun: & &1]
+
+  def take(enumerable, n) when n > 0 do
     Lazy[enumerable: enumerable,
          fun: fn(f1) ->
-           fn
-             (entry, { acc, n }) when n > 0 ->
-               { f1.(entry, acc), n-1 }
-             (_entry, acc) ->
-               throw { :stream_lazy, acc }
+           fn(entry, { acc, n }) ->
+             acc = { f1.(entry, acc), n-1 }
+             if n > 1, do: acc, else: throw { :stream_lazy, acc }
            end
          end,
          acc: n]
@@ -351,11 +351,12 @@ defmodule Stream do
   def take_while(enumerable, f) do
     Lazy[enumerable: enumerable,
          fun: fn(f1) ->
-           fn
-             entry, { acc, true } ->
-               if f.(entry), do: { f1.(entry, acc), true }, else: { acc, false }
-             _entry, acc ->
-               acc
+           fn(entry, { acc, true }) ->
+             if f.(entry) do
+               { f1.(entry, acc), true }
+             else
+               throw { :stream_lazy, { acc, false } }
+             end
            end
          end,
          acc: true]

lib/elixir/test/elixir/stream_test.exs

@@ -40,6 +40,9 @@ defmodule StreamTest do
 
     assert Enum.to_list(Stream.drop(1..5, 0)) == [1,2,3,4,5]
     assert Enum.to_list(Stream.drop(1..3, 5)) == []
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Stream.drop(nats, 2) |> Enum.take(5) == [3,4,5,6,7]
   end
 
   test :drop_while do
@@ -49,12 +52,18 @@ defmodule StreamTest do
 
     assert Enum.to_list(Stream.drop_while(1..5, &1 <= 0)) == [1,2,3,4,5]
     assert Enum.to_list(Stream.drop_while(1..3, &1 <= 5)) == []
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Stream.drop_while(nats, &1 <= 5) |> Enum.take(5) == [6,7,8,9,10]
   end
 
   test :filter do
     stream = Stream.filter([1,2,3], fn(x) -> rem(x, 2) == 0 end)
     assert is_lazy(stream)
     assert Enum.to_list(stream) == [2]
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Stream.filter(nats, &(rem(&1, 2) == 0)) |> Enum.take(5) == [2,4,6,8,10]
   end
 
   test :iterate do
@@ -72,15 +81,21 @@ defmodule StreamTest do
     stream = Stream.map([1,2,3], &1 * 2)
     assert is_lazy(stream)
     assert Enum.to_list(stream) == [2,4,6]
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Stream.map(nats, &(&1 * 2)) |> Enum.take(5) == [2,4,6,8,10]
   end
 
   test :reject do
     stream = Stream.reject([1,2,3], fn(x) -> rem(x, 2) == 0 end)
     assert is_lazy(stream)
     assert Enum.to_list(stream) == [1,3]
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Stream.reject(nats, &(rem(&1, 2) == 0)) |> Enum.take(5) == [1,3,5,7,9]
   end
 
-  test :repeatedly do 
+  test :repeatedly do
     stream = Stream.repeatedly(fn -> 1 end)
     assert Enum.take(stream, 5) == [1,1,1,1,1]
     stream = Stream.repeatedly(&:random.uniform/0)
@@ -95,6 +110,9 @@ defmodule StreamTest do
 
     assert Enum.to_list(Stream.take(1..1000, 0)) == []
     assert Enum.to_list(Stream.take(1..3, 5)) == [1,2,3]
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Enum.to_list(Stream.take(nats, 5)) == [1,2,3,4,5]
   end
 
   test :take_while do
@@ -104,12 +122,18 @@ defmodule StreamTest do
 
     assert Enum.to_list(Stream.take_while(1..1000, &1 <= 0)) == []
     assert Enum.to_list(Stream.take_while(1..3, &1 <= 5)) == [1,2,3]
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Enum.to_list(Stream.take_while(nats, &(&1 <= 5))) == [1,2,3,4,5]
   end
 
   test :with_index do
     stream = Stream.with_index([1,2,3])
     assert is_lazy(stream)
     assert Enum.to_list(stream) == [{1,0},{2,1},{3,2}]
+
+    nats = Stream.iterate(1, &(&1 + 1))
+    assert Stream.with_index(nats) |> Enum.take(3) == [{1,0},{2,1},{3,2}]
   end
 
   defp is_lazy(stream) do