Support Stream.take/2 with negative count

Author: José Valim

lib/elixir/lib/enum.ex

@@ -1433,10 +1433,10 @@ defmodule Enum do
   @doc """
   Takes the first `count` items from the collection.
 
-  If a negative value `count` is given, the last `count`
-  values will be taken. The collection is enumerated
-  once to retrieve the proper index. The remaining
-  calculation is performed from the end.
+  If a negative `count` is given, the last `count` values will
+  be taken. For such, the collection is fully enumerated keeping up
+  to `2 * count` elements in memory. Once the end of the collection is
+  reached, the last `count` elements are returned.
 
   ## Examples
 
@@ -1473,7 +1473,8 @@ defmodule Enum do
   end
 
   def take(collection, count) when count < 0 do
-    do_take_reverse(reverse(collection), abs(count), [])
+    Stream.take(collection, count).({ :cont, [] }, &{ :cont, [&1|&2] })
+    |> elem(1) |> :lists.reverse
   end
 
   @doc """
@@ -1913,18 +1914,6 @@ defmodule Enum do
     []
   end
 
-  defp do_take_reverse([h|t], counter, acc) when counter > 0 do
-    do_take_reverse(t, counter - 1, [h|acc])
-  end
-
-  defp do_take_reverse(_list, 0, acc) do
-    acc
-  end
-
-  defp do_take_reverse([], _, acc) do
-    acc
-  end
-
   ## take_while
 
   defp do_take_while([h|t], fun) do

lib/elixir/lib/stream.ex

@@ -231,9 +231,9 @@ defmodule Stream do
 
   ## Examples
 
-      iex> stream = Stream.filter_map([1, 2, 3], fn(x) -> rem(x, 2) == 0 end)
+      iex> stream = Stream.filter_map(1..6, fn(x) -> rem(x, 2) == 0 end, &(&1 * 2))
       iex> Enum.to_list(stream)
-      [2]
+      [4,8,12]
 
   """
   @spec filter_map(Enumerable.t, (element -> as_boolean(term)), (element -> any)) :: Enumerable.t
@@ -332,23 +332,62 @@ defmodule Stream do
   Lazily takes the next `n` items from the enumerable and stops
   enumeration.
 
+  If a negative `n` is given, the last `n` values will be taken.
+  For such, the collection is fully enumerated keeping up to `2 * n`
+  elements in memory. Once the end of the collection is reached,
+  the last `count` elements will be executed. Therefore, using
+  a negative `n` in an infinite collection will never return.
+
   ## Examples
 
       iex> stream = Stream.take(1..100, 5)
       iex> Enum.to_list(stream)
       [1,2,3,4,5]
 
+      iex> stream = Stream.take(1..100, -5)
+      iex> Enum.to_list(stream)
+      [96,97,98,99,100]
+
       iex> stream = Stream.cycle([1, 2, 3]) |> Stream.take(5)
       iex> Enum.to_list(stream)
       [1,2,3,1,2]
 
   """
   @spec take(Enumerable.t, non_neg_integer) :: Enumerable.t
+  def take(_enum, 0), do: Lazy[enum: []]
+
   def take(enum, n) when n > 0 do
     lazy enum, n, fn(f1) -> R.take(f1) end
   end
 
-  def take(_enum, 0), do: Lazy[enum: []]
+  def take(enum, n) when n < 0 do
+    &do_take(enum, abs(n), &1, &2)
+  end
+
+  def do_take(enum, n, acc, f) do
+    { _, { _count, buf1, buf2 } } =
+      Enumerable.reduce(enum, { :cont, { 0, [], [] } }, fn
+        entry, { count, buf1, buf2 } ->
+          buf1  = [entry|buf1]
+          count = count + 1
+          if count == n do
+            { :cont, { 0, [], buf1 } }
+          else
+            { :cont, { count, buf1, buf2} }
+          end
+      end)
+
+    Enumerable.reduce(do_take_last(buf1, buf2, n, []), acc, f)
+  end
+
+  defp do_take_last(_buf1, _buf2, 0, acc),
+    do: acc
+  defp do_take_last([], [], _, acc),
+    do: acc
+  defp do_take_last([], [h|t], n, acc),
+    do: do_take_last([], t, n-1, [h|acc])
+  defp do_take_last([h|t], buf2, n, acc),
+    do: do_take_last(t, buf2, n-1, [h|acc])
 
   @doc """
   Creates a stream that takes every `n` item from the enumerable.

lib/elixir/test/elixir/stream_test.exs

@@ -232,6 +232,17 @@ defmodule StreamTest do
     assert Enum.to_list(stream) == [5,6,7,8,9,10,11,12,13,14,15]
   end
 
+  test "take with negative count" do
+    Process.put(:stream_each, [])
+
+    stream = Stream.take(1..100, -5)
+    assert is_lazy(stream)
+
+    stream = Stream.each(stream, &Process.put(:stream_each, [&1|Process.get(:stream_each)]))
+    assert Enum.to_list(stream) == [96,97,98,99,100]
+    assert Process.get(:stream_each) == [100,99,98,97,96]
+  end
+
   test "take is zippable" do
     stream = Stream.take(1..1000, 5)
     list   = Enum.to_list(stream)