Fuse maps and tuples for printing (#14079)

Author: gldubc

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

@@ -1704,6 +1704,49 @@ defmodule Module.Types.Descr do
 
       {tag, fields, negs}
     end)
+    |> map_fusion()
+  end
+
+  # Given a dnf, fuse maps when possible
+  # e.g. %{a: integer(), b: atom()} or %{a: float(), b: atom()} into %{a: number(), b: atom()}
+  defp map_fusion(dnf) do
+    # Steps:
+    # 1. Group maps by tags and keys
+    # 2. Try fusions for each group until no fusion is found
+    # 3. Merge the groups back into a dnf
+    dnf
+    |> Enum.group_by(fn {tag, fields, _} -> {tag, Map.keys(fields)} end)
+    |> Enum.flat_map(fn {_, maps} -> fuse_maps(maps) end)
+  end
+
+  defp fuse_maps(maps) do
+    Enum.reduce(maps, [], fn map, acc ->
+      case Enum.split_while(acc, &fusible_maps?(map, &1)) do
+        {_, []} ->
+          [map | acc]
+
+        {others, [match | rest]} ->
+          fused = fuse_map_pair(map, match)
+          others ++ [fused | rest]
+      end
+    end)
+  end
+
+  # Two maps are fusible if they have no negations and differ in at most one element.
+  defp fusible_maps?({_, fields1, negs1}, {_, fields2, negs2}) do
+    negs1 != [] or negs2 != [] or
+      Map.keys(fields1)
+      |> Enum.count(fn key -> Map.get(fields1, key) != Map.get(fields2, key) end) > 1
+  end
+
+  defp fuse_map_pair({tag, fields1, []}, {_, fields2, []}) do
+    fused_fields =
+      Map.new(fields1, fn {key, type1} ->
+        type2 = Map.get(fields2, key)
+        {key, if(type1 != type2, do: union(type1, type2), else: type1)}
+      end)
+
+    {tag, fused_fields, []}
   end
 
   # If all fields are the same except one, we can optimize map difference.
@@ -1925,24 +1968,66 @@ defmodule Module.Types.Descr do
   defp tuple_to_quoted(dnf) do
     dnf
     |> tuple_simplify()
+    |> tuple_fusion()
     |> Enum.map(&tuple_each_to_quoted/1)
     |> case do
       [] -> []
       dnf -> Enum.reduce(dnf, &{:or, [], [&2, &1]}) |> List.wrap()
     end
   end
 
-  defp tuple_each_to_quoted({tag, positive_map, negative_maps}) do
-    case negative_maps do
+  # Given a dnf of tuples, fuses the tuple unions when possible,
+  # e.g. {integer(), atom()} or {float(), atom()} into {number(), atom()}
+  # The negations of two fused tuples are just concatenated.
+  defp tuple_fusion(dnf) do
+    # Steps:
+    # 1. Consider tuples without negations apart from those with
+    # 2. Group tuples by size and tag
+    # 3. Try fusions for each group until no fusion is found
+    # 4. Merge the groups back into a dnf
+    dnf
+    |> Enum.group_by(fn {tag, elems, _} -> {tag, length(elems)} end)
+    |> Enum.flat_map(fn {_, tuples} -> fuse_tuples(tuples) end)
+  end
+
+  defp fuse_tuples(tuples) do
+    Enum.reduce(tuples, [], fn tuple, acc ->
+      case Enum.split_while(acc, &fusible_tuples?(tuple, &1)) do
+        {_, []} ->
+          [tuple | acc]
+
+        {others, [match | rest]} ->
+          fused = fuse_tuple_pair(tuple, match)
+          others ++ [fused | rest]
+      end
+    end)
+  end
+
+  # Two tuples are fusible if they have no negations and differ in at most one element.
+  defp fusible_tuples?({_, elems1, negs1}, {_, elems2, negs2}) do
+    negs1 != [] or negs2 != [] or
+      Enum.zip(elems1, elems2) |> Enum.count(fn {a, b} -> a != b end) > 1
+  end
+
+  defp fuse_tuple_pair({tag, elems1, []}, {_, elems2, []}) do
+    fused_elements =
+      Enum.zip(elems1, elems2)
+      |> Enum.map(fn {a, b} -> if a != b, do: union(a, b), else: a end)
+
+    {tag, fused_elements, []}
+  end
+
+  defp tuple_each_to_quoted({tag, positive_tuple, negative_tuples}) do
+    case negative_tuples do
       [] ->
-        tuple_literal_to_quoted({tag, positive_map})
+        tuple_literal_to_quoted({tag, positive_tuple})
 
       _ ->
-        negative_maps
+        negative_tuples
         |> Enum.map(&tuple_literal_to_quoted/1)
         |> Enum.reduce(&{:or, [], [&2, &1]})
         |> Kernel.then(
-          &{:and, [], [tuple_literal_to_quoted({tag, positive_map}), {:not, [], [&1]}]}
+          &{:and, [], [tuple_literal_to_quoted({tag, positive_tuple}), {:not, [], [&1]}]}
         )
     end
   end

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

@@ -1268,6 +1268,91 @@ defmodule Module.Types.DescrTest do
       # assert difference(tuple([number(), term()]), tuple([integer(), atom()]))
       #        |> to_quoted_string() ==
       #          "{float(), term()} or {number(), term() and not atom()}"
+
+      assert union(tuple([integer(), atom()]), tuple([integer(), atom()])) |> to_quoted_string() ==
+               "{integer(), atom()}"
+
+      assert union(tuple([integer(), atom()]), tuple([float(), atom()])) |> to_quoted_string() ==
+               "{float() or integer(), atom()}"
+
+      assert union(tuple([integer(), atom()]), tuple([float(), atom()]))
+             |> union(tuple([pid(), pid(), port()]))
+             |> union(tuple([pid(), pid(), atom()]))
+             |> to_quoted_string() ==
+               "{float() or integer(), atom()} or {pid(), pid(), atom() or port()}"
+
+      assert union(open_tuple([integer()]), open_tuple([float()])) |> to_quoted_string() ==
+               "{float() or integer(), ...}"
+
+      # {:ok, {term(), integer()}} or {:ok, {term(), float()}} or {:exit, :kill} or {:exit, :timeout}
+      assert tuple([atom([:ok]), tuple([term(), empty_list()])])
+             |> union(tuple([atom([:ok]), tuple([term(), open_map()])]))
+             |> union(tuple([atom([:exit]), atom([:kill])]))
+             |> union(tuple([atom([:exit]), atom([:timeout])]))
+             |> to_quoted_string() ==
+               "{:exit, :kill or :timeout} or {:ok, {term(), %{...} or empty_list()}}"
+
+      # Detection of duplicates
+      assert tuple([atom([:ok]), term()])
+             |> union(tuple([atom([:ok]), term()]))
+             |> to_quoted_string() == "{:ok, term()}"
+
+      assert tuple([closed_map(a: integer(), b: atom()), open_map()])
+             |> union(tuple([closed_map(a: integer(), b: atom()), open_map()]))
+             |> to_quoted_string() ==
+               "{%{a: integer(), b: atom()}, %{...}}"
+
+      # Nested fusion
+      assert tuple([closed_map(a: integer(), b: atom()), open_map()])
+             |> union(tuple([closed_map(a: float(), b: atom()), open_map()]))
+             |> to_quoted_string() ==
+               "{%{a: float() or integer(), b: atom()}, %{...}}"
+
+      # Complex simplification of map/tuple combinations. Initial type is:
+      # ```
+      #   dynamic(
+      #     :error or
+      #     ({%Decimal{coef: :inf, exp: integer(), sign: integer()}, binary()} or
+      #       {%Decimal{coef: :NaN, exp: integer(), sign: integer()}, binary()} or
+      #       {%Decimal{coef: integer(), exp: integer(), sign: integer()}, term()} or
+      #       {%Decimal{coef: :inf, exp: integer(), sign: integer()} or
+      #           %Decimal{coef: :NaN, exp: integer(), sign: integer()} or
+      #           %Decimal{coef: integer(), exp: integer(), sign: integer()}, term()})
+      #   )
+      # ```
+      decimal_inf =
+        closed_map(
+          __struct__: atom([Decimal]),
+          coef: atom([:inf]),
+          exp: integer(),
+          sign: integer()
+        )
+
+      decimal_nan =
+        closed_map(
+          __struct__: atom([Decimal]),
+          coef: atom([:NaN]),
+          exp: integer(),
+          sign: integer()
+        )
+
+      decimal_int =
+        closed_map(__struct__: atom([Decimal]), coef: integer(), exp: integer(), sign: integer())
+
+      assert atom([:error])
+             |> union(
+               tuple([decimal_inf, binary()])
+               |> union(
+                 tuple([decimal_nan, binary()])
+                 |> union(
+                   tuple([decimal_int, term()])
+                   |> union(tuple([union(decimal_inf, union(decimal_nan, decimal_int)), term()]))
+                 )
+               )
+             )
+             |> dynamic()
+             |> to_quoted_string() ==
+               "dynamic(\n  :error or\n    ({%Decimal{coef: integer() or (:NaN or :inf), exp: integer(), sign: integer()}, term()} or\n       {%Decimal{coef: :NaN or :inf, exp: integer(), sign: integer()}, binary()})\n)"
     end
 
     test "map" do
@@ -1311,6 +1396,50 @@ defmodule Module.Types.DescrTest do
       assert difference(open_map(a: number(), b: atom()), open_map(a: integer()))
              |> to_quoted_string() == "%{..., a: float(), b: atom()}"
 
+      # Basic map fusion
+      assert union(closed_map(a: integer()), closed_map(a: integer())) |> to_quoted_string() ==
+               "%{a: integer()}"
+
+      assert union(closed_map(a: integer()), closed_map(a: float())) |> to_quoted_string() ==
+               "%{a: float() or integer()}"
+
+      # Nested fusion
+      assert union(closed_map(a: integer(), b: atom()), closed_map(a: float(), b: atom()))
+             |> union(closed_map(x: pid(), y: pid(), z: port()))
+             |> union(closed_map(x: pid(), y: pid(), z: atom()))
+             |> to_quoted_string() ==
+               "%{a: float() or integer(), b: atom()} or %{x: pid(), y: pid(), z: atom() or port()}"
+
+      # Open map fusion
+      assert union(open_map(a: integer()), open_map(a: float())) |> to_quoted_string() ==
+               "%{..., a: float() or integer()}"
+
+      # Fusing complex nested maps with unions
+      assert closed_map(status: atom([:ok]), data: closed_map(value: term(), count: empty_list()))
+             |> union(
+               closed_map(status: atom([:ok]), data: closed_map(value: term(), count: open_map()))
+             )
+             |> union(closed_map(status: atom([:error]), reason: atom([:timeout])))
+             |> union(closed_map(status: atom([:error]), reason: atom([:crash])))
+             |> to_quoted_string() ==
+               "%{data: %{count: %{...} or empty_list(), value: term()}, status: :ok} or\n  %{reason: :crash or :timeout, status: :error}"
+
+      # Difference and union tests
+      assert closed_map(status: atom([:ok]), value: term())
+             |> difference(closed_map(status: atom([:ok]), value: float()))
+             |> union(
+               closed_map(status: atom([:ok]), value: term())
+               |> difference(closed_map(status: atom([:ok]), value: integer()))
+             )
+             |> to_quoted_string() ==
+               "%{status: :ok, value: term()}"
+
+      # Nested map fusion
+      assert closed_map(data: closed_map(x: integer(), y: atom()), meta: open_map())
+             |> union(closed_map(data: closed_map(x: float(), y: atom()), meta: open_map()))
+             |> to_quoted_string() ==
+               "%{data: %{x: float() or integer(), y: atom()}, meta: %{...}}"
+
       # Test complex combinations
       assert intersection(open_map(a: number(), b: atom()), open_map(a: integer(), c: boolean()))
              |> union(difference(open_map(x: atom()), open_map(x: boolean())))