Implement optimization of tuple unions (#14228)

Author: sabiwara

lib/elixir/lib/module/types/descr.ex

@@ -343,7 +343,7 @@ defmodule Module.Types.Descr do
   def empty?(:term), do: false
 
   def empty?(%{} = descr) do
-    case Map.get(descr, :dynamic, descr) do
+    case :maps.get(:dynamic, descr, _default = descr) do
       :term ->
         false
 
@@ -1897,18 +1897,18 @@ defmodule Module.Types.Descr do
 
   defp map_non_negated_fuse(maps) do
     Enum.reduce(maps, [], fn map, acc ->
-      fuse_with_first_fusible(map, acc)
+      map_fuse_with_first_fusible(map, acc)
     end)
   end
 
-  defp fuse_with_first_fusible(map, []), do: [map]
+  defp map_fuse_with_first_fusible(map, []), do: [map]
 
-  defp fuse_with_first_fusible(map, [candidate | rest]) do
+  defp map_fuse_with_first_fusible(map, [candidate | rest]) do
     if fused = maybe_optimize_map_union(map, candidate) do
       # we found a fusible candidate, we're done
       [fused | rest]
     else
-      [candidate | fuse_with_first_fusible(map, rest)]
+      [candidate | map_fuse_with_first_fusible(map, rest)]
     end
   end
 
@@ -2157,9 +2157,106 @@ defmodule Module.Types.Descr do
     end
   end
 
-  # Removes duplicates in union, which should trickle to other operations.
-  # This is a cheap optimization that relies on structural equality.
-  defp tuple_union(left, right), do: left ++ (right -- left)
+  defp tuple_union(dnf1, dnf2) do
+    # Union is just concatenation, but we rely on some optimization strategies to
+    # avoid the list to grow when possible
+
+    # first pass trying to identify patterns where two maps can be fused as one
+    with [tuple1] <- dnf1,
+         [tuple2] <- dnf2,
+         optimized when optimized != nil <- maybe_optimize_tuple_union(tuple1, tuple2) do
+      [optimized]
+    else
+      # otherwise we just concatenate and remove structural duplicates
+      _ -> dnf1 ++ (dnf2 -- dnf1)
+    end
+  end
+
+  defp maybe_optimize_tuple_union({tag1, pos1, []} = tuple1, {tag2, pos2, []} = tuple2) do
+    case tuple_union_optimization_strategy(tag1, pos1, tag2, pos2) do
+      :all_equal ->
+        tuple1
+
+      {:one_index_difference, index, v1, v2} ->
+        new_pos = List.replace_at(pos1, index, union(v1, v2))
+        {tag1, new_pos, []}
+
+      :left_subtype_of_right ->
+        tuple2
+
+      :right_subtype_of_left ->
+        tuple1
+
+      nil ->
+        nil
+    end
+  end
+
+  defp maybe_optimize_tuple_union(_, _), do: nil
+
+  defp tuple_union_optimization_strategy(tag1, pos1, tag2, pos2)
+  defp tuple_union_optimization_strategy(tag, pos, tag, pos), do: :all_equal
+
+  # might be one extra loop but cheap and avoids doing deep subtype comparisons
+  defp tuple_union_optimization_strategy(:closed, pos1, :closed, pos2)
+       when length(pos1) != length(pos2),
+       do: nil
+
+  defp tuple_union_optimization_strategy(tag1, pos1, tag2, pos2) do
+    status =
+      case {tag1, tag2} do
+        {:open, :closed} -> :right_subtype_of_left
+        {:closed, :open} -> :left_subtype_of_right
+        {same, same} -> :all_equal
+      end
+
+    do_tuple_union_optimization_strategy(tag1, pos1, tag2, pos2, 0, status)
+  end
+
+  defp do_tuple_union_optimization_strategy(_tag1, [], _tag2, [], _i, status), do: status
+
+  defp do_tuple_union_optimization_strategy(:open, [], _tag2, _pos2, _i, status)
+       when status in [:all_equal, :right_subtype_of_left],
+       do: :right_subtype_of_left
+
+  defp do_tuple_union_optimization_strategy(_tag1, _pos1, :open, [], _i, status)
+       when status in [:all_equal, :left_subtype_of_right],
+       do: :left_subtype_of_right
+
+  defp do_tuple_union_optimization_strategy(tag1, [v1 | pos1], tag2, [v2 | pos2], i, status) do
+    if next_status = tuple_union_next_strategy(i, v1, v2, status) do
+      do_tuple_union_optimization_strategy(tag1, pos1, tag2, pos2, i + 1, next_status)
+    end
+  end
+
+  defp do_tuple_union_optimization_strategy(_tag1, _pos1, _tag2, _pos2, _i, _status), do: nil
+
+  defp tuple_union_next_strategy(index, v1, v2, status)
+
+  # structurally equal values do not impact the ongoing strategy
+  defp tuple_union_next_strategy(_index, same, same, status), do: status
+
+  defp tuple_union_next_strategy(index, v1, v2, :all_equal) do
+    {:one_index_difference, index, v1, v2}
+  end
+
+  defp tuple_union_next_strategy(_index, v1, v2, {:one_index_difference, _, d1, d2}) do
+    # we have at least two differences now, we switch strategy
+    # if both are subtypes in one direction, keep checking
+    cond do
+      subtype?(d1, d2) and subtype?(v1, v2) -> :left_subtype_of_right
+      subtype?(d2, d1) and subtype?(v2, v1) -> :right_subtype_of_left
+      true -> nil
+    end
+  end
+
+  defp tuple_union_next_strategy(_index, v1, v2, :left_subtype_of_right) do
+    if subtype?(v1, v2), do: :left_subtype_of_right
+  end
+
+  defp tuple_union_next_strategy(_index, v1, v2, :right_subtype_of_left) do
+    if subtype?(v2, v1), do: :right_subtype_of_left
+  end
 
   defp tuple_to_quoted(dnf, opts) do
     dnf
@@ -2189,27 +2286,19 @@ defmodule Module.Types.Descr do
 
   defp tuple_non_negated_fuse(tuples) do
     Enum.reduce(tuples, [], fn tuple, acc ->
-      case Enum.split_while(acc, &non_fusible_tuples?(tuple, &1)) do
-        {_, []} ->
-          [tuple | acc]
-
-        {others, [match | rest]} ->
-          fused = tuple_non_negated_fuse_pair(tuple, match)
-          others ++ [fused | rest]
-      end
+      tuple_fuse_with_first_fusible(tuple, acc)
     end)
   end
 
-  # Two tuples are fusible if they have no negations and differ in at most one element.
-  defp non_fusible_tuples?({_, elems1, []}, {_, elems2, []}) do
-    Enum.zip(elems1, elems2) |> Enum.count_until(fn {a, b} -> a != b end, 2) > 1
-  end
-
-  defp tuple_non_negated_fuse_pair({tag, elems1, []}, {_, elems2, []}) do
-    fused_elements =
-      Enum.zip_with(elems1, elems2, fn a, b -> if a == b, do: a, else: union(a, b) end)
+  defp tuple_fuse_with_first_fusible(tuple, []), do: [tuple]
 
-    {tag, fused_elements, []}
+  defp tuple_fuse_with_first_fusible(tuple, [candidate | rest]) do
+    if fused = maybe_optimize_tuple_union(tuple, candidate) do
+      # we found a fusible candidate, we're done
+      [fused | rest]
+    else
+      [candidate | tuple_fuse_with_first_fusible(tuple, rest)]
+    end
   end
 
   defp tuple_each_to_quoted({tag, positive_tuple, negative_tuples}, opts) do

lib/elixir/test/elixir/module/types/descr_test.exs

@@ -106,7 +106,7 @@ defmodule Module.Types.DescrTest do
              |> equal?(list(term()))
     end
 
-    test "optimizations" do
+    test "optimizations (maps)" do
       # The tests are checking the actual implementation, not the semantics.
       # This is why we are using structural comparisons.
       # It's fine to remove these if the implementation changes, but breaking
@@ -123,7 +123,7 @@ defmodule Module.Types.DescrTest do
                closed_map(a: integer(), b: atom())
              ) == closed_map(a: union(float(), integer()), b: atom())
 
-      # Optimization two: we can tell that one map is a trivial subtype of the other:
+      # Optimization two: we can tell that one map is a subtype of the other:
 
       assert union(
                closed_map(a: term(), b: term()),
@@ -145,6 +145,36 @@ defmodule Module.Types.DescrTest do
                closed_map(a: float(), b: tuple([atom(), binary()]))
              ) == closed_map(a: term(), b: tuple([term(), term()]))
     end
+
+    test "optimizations (tuples)" do
+      # Optimization one: same tags, all but one key are structurally equal
+      assert union(
+               open_tuple([float(), atom()]),
+               open_tuple([integer(), atom()])
+             ) == open_tuple([union(float(), integer()), atom()])
+
+      assert union(
+               tuple([float(), atom()]),
+               tuple([integer(), atom()])
+             ) == tuple([union(float(), integer()), atom()])
+
+      # Optimization two: we can tell that one tuple is a subtype of the other:
+
+      assert union(
+               tuple([term(), term()]),
+               tuple([float(), binary()])
+             ) == tuple([term(), term()])
+
+      assert union(
+               open_tuple([term()]),
+               tuple([float(), binary()])
+             ) == open_tuple([term()])
+
+      assert union(
+               tuple([float(), binary()]),
+               open_tuple([term()])
+             ) == open_tuple([term()])
+    end
   end
 
   describe "intersection" do
@@ -1402,8 +1432,7 @@ defmodule Module.Types.DescrTest do
              |> to_quoted_string() ==
                """
                dynamic(
-                 :error or {%Decimal{sign: integer(), coef: :NaN or :inf, exp: integer()}, binary()} or
-                   {%Decimal{sign: integer(), coef: :NaN or :inf or integer(), exp: integer()}, term()}
+                 :error or {%Decimal{sign: integer(), coef: :NaN or :inf or integer(), exp: integer()}, term()}
                )\
                """
     end