Merge pull request #1675 from ericmj/enum-at-fetch-neg

Support negative indicies in Enum.at/fetch/fetch!

Author: José Valim

lib/elixir/lib/enum.ex

@@ -226,9 +226,9 @@ defmodule Enum do
       :none
 
   """
-  @spec at(t, index) :: element | nil
-  @spec at(t, index, default) :: element | default
-  def at(collection, n, default // nil) when n >= 0 do
+  @spec at(t, integer) :: element | nil
+  @spec at(t, integer, default) :: element | default
+  def at(collection, n, default // nil) do
     case fetch(collection, n) do
       { :ok, h } -> h
       :error     -> default
@@ -398,7 +398,7 @@ defmodule Enum do
       :error
 
   """
-  @spec fetch(t, index) :: { :ok, element } | :error
+  @spec fetch(t, integer) :: { :ok, element } | :error
   def fetch(collection, n) when is_list(collection) and n >= 0 do
     do_fetch(collection, n)
   end
@@ -417,6 +417,11 @@ defmodule Enum do
     { :enum_fetch, entry } -> { :ok, entry }
   end
 
+  def fetch(collection, n) when n < 0 do
+    { list, count } = iterate_and_count_oob(collection, n)
+    if count >= 0, do: fetch(list, count), else: :error
+  end
+
   @doc """
   Finds the element at the given index (zero-based).
   Raises `OutOfBoundsError` if the given position
@@ -434,8 +439,8 @@ defmodule Enum do
       ** (Enum.OutOfBoundsError) out of bounds error
 
   """
-  @spec fetch!(t, index) :: element | no_return
-  def fetch!(collection, n) when n >= 0 do
+  @spec fetch!(t, integer) :: element | no_return
+  def fetch!(collection, n) do
     case fetch(collection, n) do
       { :ok, h } -> h
       :error     -> raise Enum.OutOfBoundsError

lib/elixir/test/elixir/enum_test.exs

@@ -52,6 +52,8 @@ defmodule EnumTest.List do
     assert Enum.at([2, 4, 6], 2) == 6
     assert Enum.at([2, 4, 6], 4) == nil
     assert Enum.at([2, 4, 6], 4, :none) == :none
+    assert Enum.at([2, 4, 6], -2) == 4
+    assert Enum.at([2, 4, 6], -4) == nil
   end
 
   test :concat_1 do
@@ -78,9 +80,15 @@ defmodule EnumTest.List do
   test :fetch! do
     assert Enum.fetch!([2, 4, 6], 0) == 2
     assert Enum.fetch!([2, 4, 6], 2) == 6
+    assert Enum.fetch!([2, 4, 6], -2) == 4
+
     assert_raise Enum.OutOfBoundsError, fn ->
       Enum.fetch!([2, 4, 6], 4)
     end
+
+    assert_raise Enum.OutOfBoundsError, fn ->
+      Enum.fetch!([2, 4, 6], -4)
+    end
   end
 
   test :drop do
@@ -129,6 +137,8 @@ defmodule EnumTest.List do
     assert Enum.fetch([2, 4, 6], 0) == { :ok, 2 }
     assert Enum.fetch([2, 4, 6], 2) == { :ok, 6 }
     assert Enum.fetch([2, 4, 6], 4) == :error
+    assert Enum.fetch([2, 4, 6], -2) == { :ok, 4}
+    assert Enum.fetch([2, 4, 6], -4) == :error
   end
 
   test :first do
@@ -379,6 +389,8 @@ defmodule EnumTest.Range do
   test :fetch! do
     assert Enum.fetch!(2..6, 0) == 2
     assert Enum.fetch!(2..6, 4) == 6
+    assert Enum.fetch!(2..6, -1) == 6
+    assert Enum.fetch!(2..6, -2) == 5
     assert Enum.fetch!(-2..-6, 0) == -2
     assert Enum.fetch!(-2..-6, 4) == -6
 
@@ -389,6 +401,10 @@ defmodule EnumTest.Range do
     assert_raise Enum.OutOfBoundsError, fn ->
       assert Enum.fetch!(-2..-6, 8)
     end
+
+    assert_raise Enum.OutOfBoundsError, fn ->
+      assert Enum.fetch!(2..6, -8)
+    end
   end
 
   test :count do