Add Access.slice/1 (#12008)

Author: vrcca

lib/elixir/lib/access.ex

@@ -829,4 +829,118 @@ defmodule Access do
   defp get_and_update_filter([], _func, _next, updates, gets) do
     {:lists.reverse(gets), :lists.reverse(updates)}
   end
+
+  @doc ~S"""
+  Returns a function that accesses all items of a list that are within the provided range.
+
+  The range will be normalized following the same rules from `Enum.slice/2`.
+
+  The returned function is typically passed as an accessor to `Kernel.get_in/2`,
+  `Kernel.get_and_update_in/3`, and friends.
+
+  ## Examples
+
+      iex> list = [%{name: "john", salary: 10}, %{name: "francine", salary: 30}, %{name: "vitor", salary: 25}]
+      iex> get_in(list, [Access.slice(1..2), :name])
+      ["francine", "vitor"]
+      iex> get_and_update_in(list, [Access.slice(1..3//2), :name], fn prev ->
+      ...>   {prev, String.upcase(prev)}
+      ...> end)
+      {["francine"], [%{name: "john", salary: 10}, %{name: "FRANCINE", salary: 30}, %{name: "vitor", salary: 25}]}
+
+  `slice/1` can also be used to pop elements out of a list or
+  a key inside of a list:
+
+      iex> list = [%{name: "john", salary: 10}, %{name: "francine", salary: 30}, %{name: "vitor", salary: 25}]
+      iex> pop_in(list, [Access.slice(-2..-1)])
+      {[%{name: "francine", salary: 30}, %{name: "vitor", salary: 25}], [%{name: "john", salary: 10}]}
+      iex> pop_in(list, [Access.slice(-2..-1), :name])
+      {["francine", "vitor"], [%{name: "john", salary: 10}, %{salary: 30}, %{salary: 25}]}
+
+  When no match is found, an empty list is returned and the update function is never called
+
+      iex> list = [%{name: "john", salary: 10}, %{name: "francine", salary: 30}, %{name: "vitor", salary: 25}]
+      iex> get_in(list, [Access.slice(5..10//2), :name])
+      []
+      iex> get_and_update_in(list, [Access.slice(5..10//2), :name], fn prev ->
+      ...>   {prev, String.upcase(prev)}
+      ...> end)
+      {[], [%{name: "john", salary: 10}, %{name: "francine", salary: 30}, %{name: "vitor", salary: 25}]}
+
+  An error is raised if the accessed structure is not a list:
+
+      iex> get_in(%{}, [Access.slice(2..10//3)])
+      ** (ArgumentError) Access.slice/1 expected a list, got: %{}
+
+  An error is raised if the step of the range is negative:
+
+      iex> get_in([], [Access.slice(2..10//-1)])
+      ** (ArgumentError) Access.slice/1 does not accept ranges with negative steps, got: 2..10//-1
+
+  """
+  @doc since: "1.14"
+  @spec slice(Range.t()) :: access_fun(data :: list, current_value :: list)
+  def slice(%Range{} = range) do
+    if range.step > 0 do
+      fn op, data, next -> slice(op, data, range, next) end
+    else
+      raise ArgumentError,
+            "Access.slice/1 does not accept ranges with negative steps, got: #{inspect(range)}"
+    end
+  end
+
+  defp slice(:get, data, %Range{} = range, next) when is_list(data) do
+    data
+    |> Enum.slice(range)
+    |> Enum.map(next)
+  end
+
+  defp slice(:get_and_update, data, range, next) when is_list(data) do
+    range = normalize_range(range, data)
+
+    if range.first > range.last do
+      {[], data}
+    else
+      get_and_update_slice(data, range, next, [], [], 0)
+    end
+  end
+
+  defp slice(_op, data, _range, _next) do
+    raise ArgumentError, "Access.slice/1 expected a list, got: #{inspect(data)}"
+  end
+
+  defp normalize_range(%Range{first: first, last: last, step: step}, list)
+       when first < 0 or last < 0 do
+    count = length(list)
+    first = if first >= 0, do: first, else: Kernel.max(first + count, 0)
+    last = if last >= 0, do: last, else: last + count
+    Range.new(first, last, step)
+  end
+
+  defp normalize_range(range, _list), do: range
+
+  defp get_and_update_slice([head | rest], range, next, updates, gets, index) do
+    if index in range do
+      case next.(head) do
+        :pop ->
+          get_and_update_slice(rest, range, next, updates, [head | gets], index + 1)
+
+        {get, update} ->
+          get_and_update_slice(
+            rest,
+            range,
+            next,
+            [update | updates],
+            [get | gets],
+            index + 1
+          )
+      end
+    else
+      get_and_update_slice(rest, range, next, [head | updates], gets, index + 1)
+    end
+  end
+
+  defp get_and_update_slice([], _range, _next, updates, gets, _index) do
+    {:lists.reverse(gets), :lists.reverse(updates)}
+  end
 end

lib/elixir/test/elixir/access_test.exs

@@ -135,6 +135,86 @@ defmodule AccessTest do
     end
   end
 
+  describe "slice/1" do
+    @test_list [1, 2, 3, 4, 5, 6, 7]
+
+    test "retrieves a range from the start of the list" do
+      assert [2, 3] == get_in(@test_list, [Access.slice(1..2)])
+    end
+
+    test "retrieves a range from the end of the list" do
+      assert [6, 7] == get_in(@test_list, [Access.slice(-2..-1)])
+    end
+
+    test "retrieves a range from positive first and negative last" do
+      assert [2, 3, 4, 5, 6] == get_in(@test_list, [Access.slice(1..-2//1)])
+    end
+
+    test "retrieves a range from negative first and positive last" do
+      assert [6, 7] == get_in(@test_list, [Access.slice(-2..7//1)])
+    end
+
+    test "retrieves a range with steps" do
+      assert [1, 3] == get_in(@test_list, [Access.slice(0..2//2)])
+      assert [2, 5] == get_in(@test_list, [Access.slice(1..4//3)])
+      assert [2] == get_in(@test_list, [Access.slice(1..2//3)])
+      assert [1, 3, 5, 7] == get_in(@test_list, [Access.slice(0..6//2)])
+    end
+
+    test "pops a range from the start of the list" do
+      assert {[2, 3], [1, 4, 5, 6, 7]} == pop_in(@test_list, [Access.slice(1..2)])
+    end
+
+    test "pops a range from the end of the list" do
+      assert {[6, 7], [1, 2, 3, 4, 5]} == pop_in(@test_list, [Access.slice(-2..-1)])
+    end
+
+    test "pops a range from positive first and negative last" do
+      assert {[2, 3, 4, 5, 6], [1, 7]} == pop_in(@test_list, [Access.slice(1..-2//1)])
+    end
+
+    test "pops a range from negative first and positive last" do
+      assert {[6, 7], [1, 2, 3, 4, 5]} == pop_in(@test_list, [Access.slice(-2..7//1)])
+    end
+
+    test "pops a range with steps" do
+      assert {[1, 3, 5], [2, 4, 6, 7]} == pop_in(@test_list, [Access.slice(0..4//2)])
+      assert {[2], [1, 3, 4, 5, 6, 7]} == pop_in(@test_list, [Access.slice(1..2//2)])
+      assert {[1, 4], [1, 2, 5, 6, 7]} == pop_in([1, 2, 1, 4, 5, 6, 7], [Access.slice(2..3)])
+    end
+
+    test "updates range from the start of the list" do
+      assert [-1, 2, 3, 4, 5, 6, 7] == update_in(@test_list, [Access.slice(0..0)], &(&1 * -1))
+
+      assert [1, -2, -3, 4, 5, 6, 7] == update_in(@test_list, [Access.slice(1..2)], &(&1 * -1))
+    end
+
+    test "updates range from the end of the list" do
+      assert [1, 2, 3, 4, 5, -6, -7] == update_in(@test_list, [Access.slice(-2..-1)], &(&1 * -1))
+
+      assert [-1, -2, 3, 4, 5, 6, 7] == update_in(@test_list, [Access.slice(-7..-6)], &(&1 * -1))
+    end
+
+    test "updates a range from positive first and negative last" do
+      assert [1, -2, -3, -4, -5, -6, 7] ==
+               update_in(@test_list, [Access.slice(1..-2//1)], &(&1 * -1))
+    end
+
+    test "updates a range from negative first and positive last" do
+      assert [1, 2, 3, 4, 5, -6, -7] ==
+               update_in(@test_list, [Access.slice(-2..7//1)], &(&1 * -1))
+    end
+
+    test "updates a range with steps" do
+      assert [-1, 2, -3, 4, -5, 6, 7] ==
+               update_in(@test_list, [Access.slice(0..4//2)], &(&1 * -1))
+    end
+
+    test "returns empty when the start of the range is greater than the end" do
+      assert [] == get_in(@test_list, [Access.slice(2..1//1)])
+    end
+  end
+
   describe "at/1" do
     @test_list [1, 2, 3, 4, 5, 6]