Add Enum.slide/3 (#11349)

Author: s3cur3

lib/elixir/lib/enum.ex

@@ -2532,6 +2532,159 @@ defmodule Enum do
     end
   end
 
+  @doc """
+  Slides a single or multiple elements given by `range_or_single_index` from `enumerable`
+  to `insertion_index`.
+
+  The semantics of the range to be moved match the semantics of `Enum.slice/2`.
+  Specifically, that means:
+
+   * Indices are normalized, meaning that negative indexes will be counted from the end
+      (for example, -1 means the last element of the enumerable). This will result in *two*
+      traversals of your enumerable on types like lists that don't provide a constant-time count.
+
+    * If the normalized index range's `last` is out of bounds, the range is truncated to the last element.
+
+    * If the normalized index range's `first` is out of bounds, the selected range for sliding
+      will be empty, so you'll get back your input list.
+
+    * Decreasing ranges (such as `5..0//1`) also select an empty range to be moved,
+      so you'll get back your input list.
+
+    * Ranges with any step but 1 will raise an error.
+
+  ## Examples
+
+      # Slide a single element
+      iex> Enum.slide([:a, :b, :c, :d, :e, :f, :g], 5, 1)
+      [:a, :f, :b, :c, :d, :e, :g]
+
+      # Slide a range of elements backward
+      iex> Enum.slide([:a, :b, :c, :d, :e, :f, :g], 3..5, 1)
+      [:a, :d, :e, :f, :b, :c, :g]
+
+      # Slide a range of elements forward
+      iex> Enum.slide([:a, :b, :c, :d, :e, :f, :g], 1..3, 5)
+      [:a, :e, :f, :b, :c, :d, :g]
+
+      # Slide with negative indices (counting from the end)
+      iex> Enum.slide([:a, :b, :c, :d, :e, :f, :g], 3..-1//1, 2)
+      [:a, :b, :d, :e, :f, :g, :c]
+      iex> Enum.slide([:a, :b, :c, :d, :e, :f, :g], -4..-2, 1)
+      [:a, :d, :e, :f, :b, :c, :g]
+
+  """
+  def slide(enumerable, range_or_single_index, insertion_index)
+
+  def slide(enumerable, single_index, insertion_index) when is_integer(single_index) do
+    slide(enumerable, single_index..single_index, insertion_index)
+  end
+
+  # This matches the behavior of Enum.slice/2
+  def slide(_, _.._//step = index_range, _insertion_index) when step != 1 do
+    raise ArgumentError,
+          "Enum.slide/3 does not accept ranges with custom steps, got: #{inspect(index_range)}"
+  end
+
+  # Normalize negative input ranges like Enum.slice/2
+  def slide(enumerable, first..last, insertion_index) when first < 0 or last < 0 do
+    count = Enum.count(enumerable)
+    normalized_first = if first >= 0, do: first, else: first + count
+    normalized_last = if last >= 0, do: last, else: last + count
+
+    if normalized_first >= 0 and normalized_first < count and normalized_first != insertion_index do
+      normalized_range = normalized_first..normalized_last//1
+      slide(enumerable, normalized_range, insertion_index)
+    else
+      Enum.to_list(enumerable)
+    end
+  end
+
+  def slide(enumerable, insertion_index.._, insertion_index) do
+    Enum.to_list(enumerable)
+  end
+
+  def slide(_, first..last, insertion_index)
+      when insertion_index > first and insertion_index < last do
+    raise "Insertion index for slide must be outside the range being moved " <>
+            "(tried to insert #{first}..#{last} at #{insertion_index})"
+  end
+
+  # Guarantees at this point: step size == 1 and first <= last and (insertion_index < first or insertion_index > last)
+  def slide(enumerable, first..last, insertion_index) do
+    impl = if is_list(enumerable), do: &slide_list_start/4, else: &slide_any/4
+
+    cond do
+      insertion_index <= first -> impl.(enumerable, insertion_index, first, last)
+      insertion_index > last -> impl.(enumerable, first, last + 1, insertion_index)
+    end
+  end
+
+  # Takes the range from middle..last and moves it to be in front of index start
+  defp slide_any(enumerable, start, middle, last) do
+    # We're going to deal with 4 "chunks" of the enumerable:
+    # 0. "Head," before the start index
+    # 1. "Slide back," between start (inclusive) and middle (exclusive)
+    # 2. "Slide front," between middle (inclusive) and last (inclusive)
+    # 3. "Tail," after last
+    #
+    # But, we're going to accumulate these into only two lists: pre and post.
+    # We'll reverse-accumulate the head into our pre list, then "slide back" into post,
+    # then "slide front" into pre, then "tail" into post.
+    #
+    # Then at the end, we're going to reassemble and reverse them, and end up with the
+    # chunks in the correct order.
+    {_size, pre, post} =
+      Enum.reduce(enumerable, {0, [], []}, fn item, {index, pre, post} ->
+        {pre, post} =
+          cond do
+            index < start -> {[item | pre], post}
+            index >= start and index < middle -> {pre, [item | post]}
+            index >= middle and index <= last -> {[item | pre], post}
+            true -> {pre, [item | post]}
+          end
+
+        {index + 1, pre, post}
+      end)
+
+    :lists.reverse(pre, :lists.reverse(post))
+  end
+
+  # Like slide_any/4 above, this optimized implementation of slide for lists depends
+  # on the indices being sorted such that we're moving middle..last to be in front of start.
+  defp slide_list_start([h | t], start, middle, last)
+       when start > 0 and start <= middle and middle <= last do
+    [h | slide_list_start(t, start - 1, middle - 1, last - 1)]
+  end
+
+  defp slide_list_start(list, 0, middle, last), do: slide_list_middle(list, middle, last, [])
+
+  defp slide_list_middle([h | t], middle, last, acc) when middle > 0 do
+    slide_list_middle(t, middle - 1, last - 1, [h | acc])
+  end
+
+  defp slide_list_middle(list, 0, last, start_to_middle) do
+    {slid_range, tail} = slide_list_last(list, last + 1, [])
+    slid_range ++ :lists.reverse(start_to_middle, tail)
+  end
+
+  # You asked for a middle index off the end of the list... you get what we've got
+  defp slide_list_middle([], _, _, acc) do
+    :lists.reverse(acc)
+  end
+
+  defp slide_list_last([h | t], last, acc) when last > 0 do
+    slide_list_last(t, last - 1, [h | acc])
+  end
+
+  defp slide_list_last(rest, 0, acc) do
+    {:lists.reverse(acc), rest}
+  end
+
+  defp slide_list_last([], _, acc) do
+    {:lists.reverse(acc), []}
+  end
+
   @doc """
   Applies the given function to each element in the `enumerable`,
   storing the result in a list and passing it as the accumulator

lib/elixir/test/elixir/enum_test.exs

@@ -803,6 +803,151 @@ defmodule EnumTest do
     assert Enum.reverse_slice([1, 2, 3], 10, 10) == [1, 2, 3]
   end
 
+  describe "slide/3" do
+    test "on an empty enum produces an empty list" do
+      for enum <- [[], %{}, 0..-1//1, MapSet.new()] do
+        assert Enum.slide(enum, 0..0, 0) == []
+      end
+    end
+
+    test "on a single-element enumerable is the same as transforming to list" do
+      for enum <- [["foo"], [1], [%{foo: "bar"}], %{foo: :bar}, MapSet.new(["foo"]), 1..1] do
+        assert Enum.slide(enum, 0..0, 0) == Enum.to_list(enum)
+      end
+    end
+
+    test "moves a single element" do
+      for zero_to_20 <- [0..20, Enum.to_list(0..20)] do
+        expected_numbers = Enum.flat_map([0..7, [14], 8..13, 15..20], &Enum.to_list/1)
+        assert Enum.slide(zero_to_20, 14..14, 8) == expected_numbers
+      end
+
+      assert Enum.slide([:a, :b, :c, :d, :e, :f], 3..3, 2) == [:a, :b, :d, :c, :e, :f]
+    end
+
+    test "on a subsection of a list reorders the range correctly" do
+      for zero_to_20 <- [0..20, Enum.to_list(0..20)] do
+        expected_numbers = Enum.flat_map([0..7, 14..18, 8..13, 19..20], &Enum.to_list/1)
+        assert Enum.slide(zero_to_20, 14..18, 8) == expected_numbers
+      end
+
+      assert Enum.slide([:a, :b, :c, :d, :e, :f], 3..4, 2) == [:a, :b, :d, :e, :c, :f]
+    end
+
+    test "handles negative indices" do
+      make_negative_range = fn first..last, length ->
+        (first - length)..(last - length)//1
+      end
+
+      test_specs = [
+        {[], 0..0, 0},
+        {[1], 0..0, 0},
+        {[-2, 1], 1..1, 1},
+        {[4, -3, 2, -1], 3..3, 2},
+        {[-5, -3, 4, 4, 5], 0..2, 3},
+        {[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 4..7, 9},
+        {[0, 1, 2, 3, 4, 5, 6, 7, 8, 9], 4..7, 0}
+      ]
+
+      for {list, range, insertion_point} <- test_specs do
+        negative_range = make_negative_range.(range, length(list))
+
+        assert Enum.slide(list, negative_range, insertion_point) ==
+                 Enum.slide(list, range, insertion_point)
+      end
+    end
+
+    test "handles mixed positive and negative indices" do
+      for zero_to_20 <- [0..20, Enum.to_list(0..20)] do
+        assert Enum.slide(zero_to_20, -6..-1, 8) ==
+                 Enum.slide(zero_to_20, 15..20, 8)
+
+        assert Enum.slide(zero_to_20, 15..-1//1, 8) ==
+                 Enum.slide(zero_to_20, 15..20, 8)
+
+        assert Enum.slide(zero_to_20, -6..20, 8) ==
+                 Enum.slide(zero_to_20, 15..20, 8)
+      end
+    end
+
+    test "raises an error when the step is not exactly 1" do
+      slide_ranges_that_should_fail = [2..10//2, 8..-1, 10..2//-1, 10..4//-2, -1..-8//-1]
+
+      for zero_to_20 <- [0..20, Enum.to_list(0..20)],
+          range_that_should_fail <- slide_ranges_that_should_fail do
+        assert_raise(ArgumentError, fn ->
+          Enum.slide(zero_to_20, range_that_should_fail, 1)
+        end)
+      end
+    end
+
+    test "doesn't change the order when the first and middle indices match" do
+      for zero_to_20 <- [0..20, Enum.to_list(0..20)] do
+        assert Enum.slide(zero_to_20, 8..18, 8) == Enum.to_list(0..20)
+      end
+
+      assert Enum.slide([:a, :b, :c, :d, :e, :f], 1..3, 1) == [:a, :b, :c, :d, :e, :f]
+    end
+
+    test "on the whole of an enumerable reorders it correctly" do
+      for zero_to_20 <- [0..20, Enum.to_list(0..20)] do
+        expected_numbers = Enum.flat_map([10..20, 0..9], &Enum.to_list/1)
+        assert Enum.slide(zero_to_20, 10..20, 0) == expected_numbers
+      end
+
+      assert Enum.slide([:a, :b, :c, :d, :e, :f], 4..5, 0) == [:e, :f, :a, :b, :c, :d]
+    end
+
+    test "raises when the insertion point is inside the range" do
+      for zero_to_20 <- [0..20, Enum.to_list(0..20)] do
+        assert_raise RuntimeError, fn ->
+          Enum.slide(zero_to_20, 10..18, 14)
+        end
+      end
+    end
+
+    test "accepts range starts that are off the end of the enum, returning the input list" do
+      assert Enum.slide([], 1..5, 0) == []
+
+      for zero_to_20 <- [0..20, Enum.to_list(0..20)] do
+        assert Enum.slide(zero_to_20, 21..25, 3) == Enum.to_list(0..20)
+      end
+    end
+
+    test "accepts range ends that are off the end of the enum, truncating the moved range" do
+      for zero_to_10 <- [0..10, Enum.to_list(0..10)] do
+        assert Enum.slide(zero_to_10, 8..15, 4) == Enum.slide(zero_to_10, 8..10, 4)
+      end
+    end
+
+    test "matches behavior for lists vs. ranges" do
+      range = 0..20
+      list = Enum.to_list(range)
+      # Below 32 elements, the map implementation currently sticks values in order.
+      # If ever the MapSet implementation changes, this will fail (not affecting the correctness
+      # of slide). I figured it'd be worth testing this for the time being just to have
+      # another enumerable (aside from range) testing the generic implementation.
+      set = MapSet.new(list)
+
+      test_specs = [
+        {0..0, 0},
+        {0..0, 20},
+        {11..11, 14},
+        {11..11, 3},
+        {4..8, 19},
+        {4..8, 0},
+        {4..8, 2},
+        {10..20, 0}
+      ]
+
+      for {slide_range, insertion_point} <- test_specs do
+        slide = &Enum.slide(&1, slide_range, insertion_point)
+        assert slide.(list) == slide.(set)
+        assert slide.(list) == slide.(range)
+      end
+    end
+  end
+
   test "scan/2" do
     assert Enum.scan([1, 2, 3, 4, 5], &(&1 + &2)) == [1, 3, 6, 10, 15]
     assert Enum.scan([], &(&1 + &2)) == []