Support Stream.drop/2 with negative count

Author: José Valim

lib/elixir/lib/stream.ex

@@ -154,18 +154,51 @@ defmodule Stream do
   @doc """
   Lazily drops the next `n` items from the enumerable.
 
+  If a negative `n` is given, it will drop the last `n` items from
+  the collection. Note that the mechanism by which this is implemented
+  will delay the emission of any item until `n` additional items have
+  been emitted by the enum.
+
   ## Examples
 
       iex> stream = Stream.drop(1..10, 5)
       iex> Enum.to_list(stream)
       [6,7,8,9,10]
 
+      iex> stream = Stream.drop(1..10, -5)
+      iex> Enum.to_list(stream)
+      [1,2,3,4,5]
+
   """
   @spec drop(Enumerable.t, non_neg_integer) :: Enumerable.t
   def drop(enum, n) when n >= 0 do
     lazy enum, n, fn(f1) -> R.drop(f1) end
   end
 
+  def drop(enum, n) when n < 0 do
+    n = abs(n)
+
+    lazy enum, { 0, [], [] }, fn(f1) ->
+      fn
+        entry, [h, { count, buf1, [] } | t] ->
+          do_drop(:cont, n, entry, h, count, buf1, [], t)
+        entry, [h, { count, buf1, [next|buf2] } | t] ->
+          { reason, [h|t] } = f1.(next, [h|t])
+          do_drop(reason, n, entry, h, count, buf1, buf2, t)
+      end
+    end
+  end
+
+  defp do_drop(reason, n, entry, h, count, buf1, buf2, t) do
+    buf1  = [entry|buf1]
+    count = count + 1
+    if count == n do
+      { reason, [h, { 0, [], :lists.reverse(buf1) }|t] }
+    else
+      { reason, [h, { count, buf1, buf2 }|t] }
+    end
+  end
+
   @doc """
   Lazily drops elements of the enumerable while the given
   function returns true.
@@ -364,7 +397,7 @@ defmodule Stream do
     &do_take(enum, abs(n), &1, &2)
   end
 
-  def do_take(enum, n, acc, f) do
+  defp do_take(enum, n, acc, f) do
     { _, { _count, buf1, buf2 } } =
       Enumerable.reduce(enum, { :cont, { 0, [], [] } }, fn
         entry, { count, buf1, buf2 } ->
@@ -373,7 +406,7 @@ defmodule Stream do
           if count == n do
             { :cont, { 0, [], buf1 } }
           else
-            { :cont, { count, buf1, buf2} }
+            { :cont, { count, buf1, buf2 } }
           end
       end)
 

lib/elixir/test/elixir/stream_test.exs

@@ -93,6 +93,36 @@ defmodule StreamTest do
     assert Stream.drop(nats, 2) |> Enum.take(5) == [3,4,5,6,7]
   end
 
+  test "drop with negative count" do
+    stream = Stream.drop(1..10, -5)
+    assert is_lazy(stream)
+    assert Enum.to_list(stream) == [1,2,3,4,5]
+
+    stream = Stream.drop(1..10, -5)
+    list   = Enum.to_list(stream)
+    assert Enum.zip(list, list) == Enum.zip(stream, stream)
+  end
+
+  test "drop with negative count stream entries" do
+    par = self
+    pid = spawn_link fn ->
+      Enum.each Stream.drop(&inbox_stream/2, -3),
+                fn x -> par <- { :stream, x } end
+    end
+
+    pid <- { :stream, 1 }
+    pid <- { :stream, 2 }
+    pid <- { :stream, 3 }
+    refute_receive { :stream, 1 }
+
+    pid <- { :stream, 4 }
+    assert_receive { :stream, 1 }
+
+    pid <- { :stream, 5 }
+    assert_receive { :stream, 2 }
+    refute_receive { :stream, 3 }
+  end
+
   test "drop_while" do
     stream = Stream.drop_while(1..10, &(&1 <= 5))
     assert is_lazy(stream)
@@ -316,4 +346,19 @@ defmodule StreamTest do
   defp is_lazy(stream) do
     is_record(stream, Stream.Lazy) or is_function(stream, 2)
   end
+
+  defp inbox_stream({ :suspend, acc }, f) do
+    { :suspended, acc, &inbox_stream(&1, f) }
+  end
+
+  defp inbox_stream({ :halt, acc }, _f) do
+    { :halted, acc }
+  end
+
+  defp inbox_stream({ :cont, acc }, f) do
+    receive do
+      { :stream, item } ->
+        inbox_stream(f.(item, acc), f)
+    end
+  end
 end