Improve map unificationemj/check-maps-2

Consider both sides of the unification and correctly match both required and optional keys against the other side.

6 files changed, 213 insertions, 141 deletions

diff --git a/lib/elixir/lib/map.ex b/lib/elixir/lib/map.ex
index 365a83dca..739261f6a 100644
--- a/lib/elixir/lib/map.ex
+++ b/lib/elixir/lib/map.ex
@@ -62,12 +62,12 @@ defmodule Map do
   keys on the left-hand side and their values match the ones on the left-hand
   side. This means that an empty map matches every map.
 
-      iex> %{} = %{foo: "bar"}
-      %{foo: "bar"}
-      iex> %{a: a} = %{:a => 1, "b" => 2, [:c, :e, :e] => 3}
-      iex> a
-      1
-      iex> %{:c => 3} = %{:a => 1, 2 => :b}
+      %{} = %{foo: "bar"}
+      #=> %{foo: "bar"}
+      %{a: a} = %{:a => 1, "b" => 2, [:c, :e, :e] => 3}
+      a
+      #=> 1
+      %{:c => 3} = %{:a => 1, 2 => :b}
       ** (MatchError) no match of right hand side value: %{2 => :b, :a => 1}
 
   Variables can be used as map keys both when writing map literals as well as
diff --git a/lib/elixir/lib/module/types.ex b/lib/elixir/lib/module/types.ex
index 1d782a2cb..f8ccd0696 100644
--- a/lib/elixir/lib/module/types.ex
+++ b/lib/elixir/lib/module/types.ex
@@ -266,6 +266,8 @@ defmodule Module.Types do
     end
   end
 
+  defp missing_field(_, _), do: :error
+
   defp format_traces([], _simplify?) do
     {[], []}
   end
diff --git a/lib/elixir/lib/module/types/infer.ex b/lib/elixir/lib/module/types/infer.ex
index 03d821388..8c1fc4e05 100644
--- a/lib/elixir/lib/module/types/infer.ex
+++ b/lib/elixir/lib/module/types/infer.ex
@@ -158,104 +158,104 @@ defmodule Module.Types.Infer do
     end
   end
 
-  defp unify_maps(source_pairs, target_pairs, %{context: :pattern} = stack, context) do
+  # * All required keys on each side need to match to the other side.
+  # * All optional keys on each side that do not match must be discarded.
+
+  defp unify_maps(source_pairs, target_pairs, stack, context) do
     source_pairs = expand_struct(source_pairs)
     target_pairs = expand_struct(target_pairs)
 
-    # Since maps in patterns only support literal keys (excluding maps)
-    # we can do exact type match without subtype checking
-
-    unique_right_pairs =
-      Enum.reject(target_pairs, fn {_kind, key, _value} ->
-        List.keyfind(source_pairs, key, 1)
-      end)
-
-    unique_pairs = source_pairs ++ unique_right_pairs
-
-    # Build union of all unique key-value pairs between the maps
-    result =
-      map_reduce_ok(unique_pairs, context, fn {source_kind, source_key, source_value}, context ->
-        case List.keyfind(target_pairs, source_key, 1) do
-          {target_kind, ^source_key, target_value} ->
-            case unify(source_value, target_value, stack, context) do
-              {:ok, value, context} ->
-                {:ok, {unify_kinds(source_kind, target_kind), source_key, value}, context}
-
-              {:error, reason} ->
-                {:error, reason}
-            end
-
-          nil ->
-            {:ok, {source_kind, source_key, source_value}, context}
+    {source_required, source_optional} = split_pairs(source_pairs)
+    {target_required, target_optional} = split_pairs(target_pairs)
+
+    with {:ok, source_required_pairs, context} <-
+           unify_source_required(source_required, target_pairs, source_pairs, stack, context),
+         {:ok, target_required_pairs, context} <-
+           unify_target_required(target_required, source_pairs, target_pairs, stack, context),
+         {:ok, source_optional_pairs, context} <-
+           unify_map_optional(source_optional, target_pairs, source_pairs, stack, context),
+         {:ok, target_optional_pairs, context} <-
+           unify_map_optional(target_optional, source_pairs, target_pairs, stack, context),
+         pairs =
+           [
+             source_required_pairs,
+             target_required_pairs,
+             source_optional_pairs,
+             target_optional_pairs
+           ]
+           |> Enum.concat()
+           # Remove duplicate pairs from matching in both left and right directions
+           |> Enum.uniq()
+           |> simplify_struct(),
+         do: {:ok, {:map, pairs}, context}
+  end
+
+  defp unify_source_required(source_required, target_pairs, source_pairs, stack, context) do
+    map_reduce_ok(source_required, context, fn {source_key, source_value}, context ->
+      Enum.find_value(target_pairs, fn {target_kind, target_key, target_value} ->
+        with {:ok, key, context} <- unify(source_key, target_key, stack, context) do
+          case unify(source_value, target_value, stack, context) do
+            {:ok, value, context} ->
+              {:ok, {:required, key, value}, context}
+
+            {:error, _reason} ->
+              source_map = {:map, [{:required, source_key, source_value}]}
+              target_map = {:map, [{target_kind, target_key, target_value}]}
+              error({:unable_unify, source_map, target_map}, stack, context)
+          end
+        else
+          {:error, _reason} -> nil
         end
-      end)
-
-    case result do
-      {:ok, pairs, context} -> {:ok, {:map, simplify_struct(pairs)}, context}
-      {:error, reason} -> {:error, reason}
-    end
+      end) || error({:unable_unify, {:map, source_pairs}, {:map, target_pairs}}, stack, context)
+    end)
   end
 
-  defp unify_maps(source_pairs, target_pairs, %{context: :expr} = stack, context) do
-    source_pairs = expand_struct(source_pairs)
-    target_pairs = expand_struct(target_pairs)
+  defp unify_target_required(target_required, source_pairs, target_pairs, stack, context) do
+    map_reduce_ok(target_required, context, fn {target_key, target_value}, context ->
+      Enum.find_value(source_pairs, fn {target_kind, source_key, source_value} ->
+        with {:ok, key, context} <- unify(source_key, target_key, stack, context) do
+          case unify(source_value, target_value, stack, context) do
+            {:ok, value, context} ->
+              {:ok, {:required, key, value}, context}
 
-    result =
-      flat_map_reduce_ok(source_pairs, context, fn {source_kind, _, _} = source_pair, context ->
-        # Currently we only match on exact and dynamic types
-        # since those are the only we get from map.key
-        with :error <- exact_map_match(source_pair, target_pairs, stack, context),
-             :error <- dynamic_map_match(source_pair, target_pairs, stack, context) do
-          if source_kind == :optional do
-            {:ok, [], context}
-          else
-            source_map = {:map, simplify_struct(source_pairs)}
-            target_map = {:map, simplify_struct(target_pairs)}
-            error({:unable_unify, source_map, target_map}, stack, context)
+            {:error, _reason} ->
+              source_map = {:map, [{:required, source_key, source_value}]}
+              target_map = {:map, [{target_kind, target_key, target_value}]}
+              error({:unable_unify, source_map, target_map}, stack, context)
           end
+        else
+          {:error, _reason} -> nil
         end
-      end)
-
-    case result do
-      {:ok, pairs, context} ->
-        pairs = Enum.uniq_by(pairs ++ target_pairs, fn {_kind, key, _value} -> key end)
-        {:ok, {:map, simplify_struct(pairs)}, context}
-
-      {:error, reason} ->
-        {:error, reason}
-    end
+      end) || error({:unable_unify, {:map, source_pairs}, {:map, target_pairs}}, stack, context)
+    end)
   end
 
-  defp exact_map_match({source_kind, source_key, source_value}, target_pairs, stack, context) do
-    case List.keyfind(target_pairs, source_key, 1) do
-      {target_kind, ^source_key, target_value} ->
-        case unify(source_value, target_value, stack, context) do
-          {:ok, value, context} ->
-            {:ok, [{unify_kinds(source_kind, target_kind), source_key, value}], context}
+  defp unify_map_optional(left_optional, right_pairs, left_pairs, stack, context) do
+    flat_map_reduce_ok(left_optional, context, fn {left_key, left_value}, context ->
+      Enum.find_value(right_pairs, fn {kind, right_key, right_value} ->
+        with :optional <- kind,
+             {:ok, key, context} <- unify(left_key, right_key, stack, context) do
+          case unify(left_value, right_value, stack, context) do
+            {:ok, value, context} ->
+              {:ok, [{:optional, key, value}], context}
 
-          {:error, reason} ->
-            {:error, reason}
+            {:error, _reason} ->
+              error({:unable_unify, {:map, left_pairs}, {:map, right_pairs}}, stack, context)
+          end
+        else
+          _ -> nil
         end
-
-      nil ->
-        :error
-    end
+      end) || {:ok, [], context}
+    end)
   end
 
-  defp dynamic_map_match({source_kind, source_key, source_value}, target_pairs, stack, context) do
-    case dynamic_pair(target_pairs) do
-      {:ok, {_target_kind, target_value}} ->
-        case unify(source_value, target_value, stack, context) do
-          {:ok, value, context} ->
-            {:ok, [{source_kind, source_key, value}], context}
+  defp split_pairs(pairs) do
+    {required, optional} =
+      Enum.split_with(pairs, fn {kind, _key, _value} -> kind == :required end)
 
-          {:error, reason} ->
-            {:error, reason}
-        end
-
-      :error ->
-        :error
-    end
+    required = Enum.map(required, fn {_kind, key, value} -> {key, value} end)
+    optional = Enum.map(optional, fn {_kind, key, value} -> {key, value} end)
+    {required, optional}
   end
 
   @doc """
@@ -467,6 +467,8 @@ defmodule Module.Types.Infer do
   def unify_kinds(:optional, :optional), do: :optional
 
   # Collect relevant information from context and traces to report error
+  # TODO: We should do this lazily since in some cases unification will error
+  #       but we continue attempting unifying other types
   defp error({:unable_unify, left, right}, stack, context) do
     {fun, arity} = context.function
     line = get_meta(stack.last_expr)[:line]
@@ -491,13 +493,11 @@ defmodule Module.Types.Infer do
       |> Enum.uniq()
 
     Enum.flat_map(stack, fn var_index ->
-      case Map.fetch(context.traces, var_index) do
-        {:ok, traces} ->
-          expr_var = Map.fetch!(context.types_to_vars, var_index)
-          Enum.map(traces, &{expr_var, &1})
-
-        _other ->
-          []
+      with {:ok, traces} <- Map.fetch(context.traces, var_index),
+           {:ok, expr_var} <- Map.fetch(context.types_to_vars, var_index) do
+        Enum.map(traces, &{expr_var, &1})
+      else
+        _other -> []
       end
     end)
   end
@@ -564,17 +564,11 @@ defmodule Module.Types.Infer do
     end
   end
 
+  # TODO: Resolve type variables if %{__struct__: var, ...}
   defp fetch_struct_pair(pairs) do
     case Enum.find(pairs, &match?({:required, {:atom, :__struct__}, {:atom, _}}, &1)) do
       {:required, {:atom, :__struct__}, {:atom, module}} -> {:ok, module}
       nil -> :error
     end
   end
-
-  defp dynamic_pair(pairs) do
-    case List.keyfind(pairs, :dynamic, 1) do
-      {kind, :dynamic, type} -> {:ok, {kind, type}}
-      nil -> :error
-    end
-  end
 end
diff --git a/lib/elixir/test/elixir/module/types/expr_test.exs b/lib/elixir/test/elixir/module/types/expr_test.exs
index 628da75a7..686eb55b8 100644
--- a/lib/elixir/test/elixir/module/types/expr_test.exs
+++ b/lib/elixir/test/elixir/module/types/expr_test.exs
@@ -6,7 +6,7 @@ defmodule Module.Types.ExprTest do
   alias Module.Types
   alias Module.Types.Infer
 
-  defmacrop quoted_expr(vars \\ [], body) do
+  defmacro quoted_expr(vars \\ [], body) do
     quote do
       {vars, body} = unquote(Macro.escape(expand_expr(vars, body)))
 
@@ -147,8 +147,8 @@ defmodule Module.Types.ExprTest do
               {:map,
                [
                  {:required, {:atom, :bar}, {:var, 0}},
-                 {:optional, :dynamic, :dynamic},
-                 {:required, {:atom, :foo}, {:var, 1}}
+                 {:required, {:atom, :foo}, {:var, 1}},
+                 {:optional, :dynamic, :dynamic}
                ]}}
 
     assert quoted_expr(
@@ -163,8 +163,8 @@ defmodule Module.Types.ExprTest do
               {:map,
                [
                  {:required, {:atom, :bar}, {:atom, :bar}},
-                 {:optional, :dynamic, :dynamic},
-                 {:required, {:atom, :foo}, {:atom, :foo}}
+                 {:required, {:atom, :foo}, {:atom, :foo}},
+                 {:optional, :dynamic, :dynamic}
                ]}}
 
     assert {:error,
@@ -180,39 +180,83 @@ defmodule Module.Types.ExprTest do
              )
   end
 
-  test "struct field" do
-    assert quoted_expr(%File.Stat{}.mtime) == {:ok, {:atom, nil}}
-    assert quoted_expr(%File.Stat{mtime: 123}.mtime) == {:ok, :integer}
+  defmodule :"Elixir.Module.Types.ExprTest.Struct2" do
+    defstruct [:field]
+  end
 
+  test "map fields and structs" do
     assert quoted_expr(
+             [map],
              (
-               map = %File.Stat{}
-               map.mtime
+               %Module.Types.ExprTest.Struct2{} = map
+               map.field
+               map
              )
-           ) == {:ok, {:atom, nil}}
+           ) ==
+             {:ok,
+              {:map,
+               [
+                 {:required, {:atom, :field}, {:var, 0}},
+                 {:required, {:atom, :__struct__}, {:atom, Module.Types.ExprTest.Struct2}}
+               ]}}
 
     assert quoted_expr(
+             [map],
              (
-               map = %File.Stat{mtime: 123}
-               map.mtime
+               _ = map.field
+               %Module.Types.ExprTest.Struct2{} = map
              )
-           ) == {:ok, :integer}
+           ) ==
+             {:ok,
+              {:map,
+               [
+                 {:required, {:atom, :field}, {:var, 0}},
+                 {:required, {:atom, :__struct__}, {:atom, Module.Types.ExprTest.Struct2}}
+               ]}}
+
+    assert {:error, {{:unable_unify, _, _, _}, _}} =
+             quoted_expr(
+               [map],
+               (
+                 %Module.Types.ExprTest.Struct2{} = map
+                 map.no_field
+               )
+             )
+
+    assert {:error, {{:unable_unify, _, _, _}, _}} =
+             quoted_expr(
+               [map],
+               (
+                 _ = map.no_field
+                 %Module.Types.ExprTest.Struct2{} = map
+               )
+             )
+  end
+
+  test "map pattern" do
+    assert quoted_expr(%{a: :b} = %{a: :b}) ==
+             {:ok, {:map, [{:required, {:atom, :a}, {:atom, :b}}]}}
+
+    assert quoted_expr(
+             (
+               a = :a
+               %{^a => :b} = %{:a => :b}
+             )
+           ) == {:ok, {:map, [{:required, {:atom, :a}, {:atom, :b}}]}}
 
     assert {:error,
-            {{:unable_unify,
-              {:map, [{:required, {:atom, :foo}, {:var, 0}}, {:optional, :dynamic, :dynamic}]},
-              {:map, [{:required, {:atom, :__struct__}, {:atom, File.Stat}} | _]}, _},
-             _}} = quoted_expr(%File.Stat{}.foo)
+            {{:unable_unify, {:map, [{:required, {:atom, :c}, {:atom, :d}}]},
+              {:map, [{:required, {:atom, :a}, {:atom, :b}}, {:optional, :dynamic, :dynamic}]},
+              _}, _}} = quoted_expr(%{a: :b} = %{c: :d})
 
     assert {:error,
-            {{:unable_unify,
-              {:map, [{:required, {:atom, :foo}, {:var, 1}}, {:optional, :dynamic, :dynamic}]},
-              {:map, [{:required, {:atom, :__struct__}, {:atom, File.Stat}} | _]}, _},
+            {{:unable_unify, {:map, [{:required, {:atom, :b}, {:atom, :error}}]},
+              {:map, [{:required, {:var, 0}, {:atom, :ok}}, {:optional, :dynamic, :dynamic}]}, _},
              _}} =
              quoted_expr(
                (
-                 map = %File.Stat{}
-                 map.foo
+                 a = :a
+                 %{^a => :ok} = %{:b => :error}
                )
              )
   end
diff --git a/lib/elixir/test/elixir/module/types/infer_test.exs b/lib/elixir/test/elixir/module/types/infer_test.exs
index e021195e9..aeec23cbf 100644
--- a/lib/elixir/test/elixir/module/types/infer_test.exs
+++ b/lib/elixir/test/elixir/module/types/infer_test.exs
@@ -100,23 +100,35 @@ defmodule Module.Types.InferTest do
     test "map" do
       assert unify_lift({:map, []}, {:map, []}) == {:ok, {:map, []}}
 
-      assert unify_lift({:map, [{:required, :integer, :atom}]}, {:map, []}) ==
-               {:ok, {:map, [{:required, :integer, :atom}]}}
-
-      assert unify_lift({:map, []}, {:map, [{:required, :integer, :atom}]}) ==
+      assert unify_lift(
+               {:map, [{:required, :integer, :atom}]},
+               {:map, [{:optional, :dynamic, :dynamic}]}
+             ) ==
                {:ok, {:map, [{:required, :integer, :atom}]}}
 
       assert unify_lift(
-               {:map, [{:required, :integer, :atom}]},
+               {:map, [{:optional, :dynamic, :dynamic}]},
                {:map, [{:required, :integer, :atom}]}
              ) ==
                {:ok, {:map, [{:required, :integer, :atom}]}}
 
       assert unify_lift(
+               {:map, [{:optional, :dynamic, :dynamic}]},
+               {:map, [{:required, :integer, :atom}, {:optional, :dynamic, :dynamic}]}
+             ) ==
+               {:ok, {:map, [{:required, :integer, :atom}, {:optional, :dynamic, :dynamic}]}}
+
+      assert unify_lift(
+               {:map, [{:required, :integer, :atom}, {:optional, :dynamic, :dynamic}]},
+               {:map, [{:optional, :dynamic, :dynamic}]}
+             ) ==
+               {:ok, {:map, [{:required, :integer, :atom}, {:optional, :dynamic, :dynamic}]}}
+
+      assert unify_lift(
                {:map, [{:required, :integer, :atom}]},
-               {:map, [{:required, :atom, :integer}]}
+               {:map, [{:required, :integer, :atom}]}
              ) ==
-               {:ok, {:map, [{:required, :integer, :atom}, {:required, :atom, :integer}]}}
+               {:ok, {:map, [{:required, :integer, :atom}]}}
 
       assert unify_lift(
                {:map, [{:required, {:atom, :foo}, :boolean}]},
@@ -124,7 +136,25 @@ defmodule Module.Types.InferTest do
              ) ==
                {:ok, {:map, [{:required, {:atom, :foo}, :boolean}]}}
 
-      assert {:error, {{:unable_unify, :integer, :atom, _}, _}} =
+      assert {:error,
+              {{:unable_unify, {:map, [{:required, :integer, :atom}]},
+                {:map, [{:required, :atom, :integer}]}, _},
+               _}} =
+               unify_lift(
+                 {:map, [{:required, :integer, :atom}]},
+                 {:map, [{:required, :atom, :integer}]}
+               )
+
+      assert {:error, {{:unable_unify, {:map, [{:required, :integer, :atom}]}, {:map, []}, _}, _}} =
+               unify_lift({:map, [{:required, :integer, :atom}]}, {:map, []})
+
+      assert {:error, {{:unable_unify, {:map, []}, {:map, [{:required, :integer, :atom}]}, _}, _}} =
+               unify_lift({:map, []}, {:map, [{:required, :integer, :atom}]})
+
+      assert {:error,
+              {{:unable_unify, {:map, [{:required, {:atom, :foo}, :integer}]},
+                {:map, [{:required, {:atom, :foo}, :atom}]}, _},
+               _}} =
                unify_lift(
                  {:map, [{:required, {:atom, :foo}, :integer}]},
                  {:map, [{:required, {:atom, :foo}, :atom}]}
diff --git a/lib/elixir/test/elixir/module/types_test.exs b/lib/elixir/test/elixir/module/types_test.exs
index 63e070aab..25c934861 100644
--- a/lib/elixir/test/elixir/module/types_test.exs
+++ b/lib/elixir/test/elixir/module/types_test.exs
@@ -229,9 +229,9 @@ defmodule Module.TypesTest do
                {:ok,
                 [
                   {:map,
-                   [{:optional, :dynamic, :dynamic}, {:required, {:atom, true}, {:atom, false}}]},
+                   [{:required, {:atom, true}, {:atom, false}}, {:optional, :dynamic, :dynamic}]},
                   {:map,
-                   [{:optional, :dynamic, :dynamic}, {:required, {:atom, true}, {:atom, false}}]}
+                   [{:required, {:atom, true}, {:atom, false}}, {:optional, :dynamic, :dynamic}]}
                 ]}
 
       assert quoted_head([%{true: bool}], [is_boolean(bool)]) ==
@@ -246,19 +246,21 @@ defmodule Module.TypesTest do
                   {:map,
                    [
                      {:required, {:atom, false}, {:atom, false}},
-                     {:optional, :dynamic, :dynamic},
-                     {:required, {:atom, true}, {:atom, true}}
+                     {:required, {:atom, true}, {:atom, true}},
+                     {:optional, :dynamic, :dynamic}
                    ]},
                   {:map,
                    [
                      {:required, {:atom, false}, {:atom, false}},
-                     {:optional, :dynamic, :dynamic},
-                     {:required, {:atom, true}, {:atom, true}}
+                     {:required, {:atom, true}, {:atom, true}},
+                     {:optional, :dynamic, :dynamic}
                    ]}
                 ]}
 
-      assert {:error, {{:unable_unify, {:atom, true}, {:atom, false}, _}, _}} =
-               quoted_head([%{true: false} = foo, %{true: true} = foo])
+      assert {:error,
+              {{:unable_unify, {:map, [{:required, {:atom, true}, {:atom, true}}]},
+                {:map, [{:required, {:atom, true}, {:atom, false}}]}, _},
+               _}} = quoted_head([%{true: false} = foo, %{true: true} = foo])
     end
 
     test "struct var guard" do