From 5d5a9320ebda0217e2a5092d5068722aa4fa377c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Valim?= <jose.valim@dashbit.co>
Date: Sat, 12 Sep 2020 14:03:01 +0200
Subject: [PATCH 1/2] Handle subtypes in map keys

---
 lib/elixir/lib/module/types.ex                |  27 +++-
 lib/elixir/lib/module/types/infer.ex          | 125 ++++++++++++------
 lib/elixir/lib/module/types/of.ex             |  57 ++++++--
 .../test/elixir/module/types/infer_test.exs   |  25 ++++
 .../test/elixir/module/types/map_test.exs     |  78 ++++++++++-
 .../test/elixir/module/types/pattern_test.exs |   4 +-
 6 files changed, 254 insertions(+), 62 deletions(-)

diff --git a/lib/elixir/lib/module/types.ex b/lib/elixir/lib/module/types.ex
index efbd0502f8a..c77215994f6 100644
--- a/lib/elixir/lib/module/types.ex
+++ b/lib/elixir/lib/module/types.ex
@@ -1,6 +1,10 @@
 defmodule Module.Types do
   @moduledoc false
 
+  defmodule Error do
+    defexception [:message]
+  end
+
   import Module.Types.Helpers
   alias Module.Types.{Expr, Pattern}
 
@@ -8,12 +12,31 @@ defmodule Module.Types do
   def warnings(module, file, defs, no_warn_undefined, cache) do
     stack = stack()
 
-    Enum.flat_map(defs, fn {{fun, _arity} = function, kind, meta, clauses} ->
+    Enum.flat_map(defs, fn {{fun, arity} = function, kind, meta, clauses} ->
       context = context(with_file_meta(meta, file), module, function, no_warn_undefined, cache)
 
       Enum.flat_map(clauses, fn {_meta, args, guards, body} ->
         def_expr = {kind, meta, [guards_to_expr(guards, {fun, [], args})]}
-        warnings_from_clause(args, guards, body, def_expr, stack, context)
+
+        try do
+          warnings_from_clause(args, guards, body, def_expr, stack, context)
+        rescue
+          e ->
+            def_expr = {kind, meta, [guards_to_expr(guards, {fun, [], args}), [do: body]]}
+
+            error =
+              Error.exception("""
+              found error while checking types for #{Exception.format_mfa(module, fun, arity)}
+
+              #{Macro.to_string(def_expr)}
+
+              Please report this bug: https://github.com/elixir-lang/elixir/issues
+
+              #{Exception.format_banner(:error, e, __STACKTRACE__)}\
+              """)
+
+            reraise error, __STACKTRACE__
+        end
       end)
     end)
   end
diff --git a/lib/elixir/lib/module/types/infer.ex b/lib/elixir/lib/module/types/infer.ex
index 546c06e519c..6c0c32ebab6 100644
--- a/lib/elixir/lib/module/types/infer.ex
+++ b/lib/elixir/lib/module/types/infer.ex
@@ -7,6 +7,7 @@ defmodule Module.Types.Infer do
   Unifies two types and returns the unified type and an updated typing context
   or an error in case of a typing conflict.
   """
+  # TODO: handle unions
   def unify(source, target, stack, context) do
     case do_unify(source, target, stack, context) do
       {:ok, type, context} ->
@@ -32,21 +33,21 @@ defmodule Module.Types.Infer do
   end
 
   defp do_unify(type, {:var, var}, stack, context) do
-    case Map.fetch!(context.types, var) do
-      {:var, var_type} ->
+    case context.types do
+      %{^var => {:var, var_type}} ->
         do_unify(type, {:var, var_type}, stack, context)
 
-      _other ->
+      %{} ->
         unify_var(var, type, stack, context, _var_source = false)
     end
   end
 
   defp do_unify({:var, var}, type, stack, context) do
-    case Map.fetch!(context.types, var) do
-      {:var, var_type} ->
+    case context.types do
+      %{^var => {:var, var_type}} ->
         do_unify({:var, var_type}, type, stack, context)
 
-      _other ->
+      %{} ->
         unify_var(var, type, stack, context, _var_source = true)
     end
   end
@@ -96,8 +97,8 @@ defmodule Module.Types.Infer do
   end
 
   defp unify_var(var, type, stack, context, var_source?) do
-    case Map.fetch!(context.types, var) do
-      :unbound ->
+    case context.types do
+      %{^var => :unbound} ->
         context = refine_var(var, type, stack, context)
         stack = push_unify_stack(var, stack)
 
@@ -111,7 +112,7 @@ defmodule Module.Types.Infer do
           {:ok, {:var, var}, context}
         end
 
-      var_type ->
+      %{^var => var_type} ->
         # Only add trace if the variable wasn't already "expanded"
         context =
           if variable_expanded?(var, stack, context) do
@@ -268,13 +269,15 @@ defmodule Module.Types.Infer do
   If the variable has already been added, return the existing type variable.
   """
   def new_var(var, context) do
-    case Map.fetch(context.vars, var_name(var)) do
-      {:ok, type} ->
+    var_name = var_name(var)
+
+    case context.vars do
+      %{^var_name => type} ->
         {type, context}
 
-      :error ->
+      %{} ->
         type = {:var, context.counter}
-        vars = Map.put(context.vars, var_name(var), type)
+        vars = Map.put(context.vars, var_name, type)
         types_to_vars = Map.put(context.types_to_vars, context.counter, var)
         types = Map.put(context.types, context.counter, :unbound)
         traces = Map.put(context.traces, context.counter, [])
@@ -312,7 +315,16 @@ defmodule Module.Types.Infer do
     {type, context}
   end
 
-  def resolve_var({:var, var}, context), do: resolve_var(Map.fetch!(context.types, var), context)
+  @doc """
+  Resolves a variable raising if it is unbound.
+  """
+  def resolve_var({:var, var}, context) do
+    case context.types do
+      %{^var => :unbound} -> raise "cannot resolve unbound var"
+      %{^var => type} -> resolve_var(type, context)
+    end
+  end
+
   def resolve_var(other, _context), do: other
 
   # Check unify stack to see if variable was already expanded
@@ -321,15 +333,15 @@ defmodule Module.Types.Infer do
   end
 
   defp variable_same?(left, right, context) do
-    case Map.fetch(context.types, left) do
-      {:ok, {:var, new_left}} ->
+    case context.types do
+      %{^left => {:var, new_left}} ->
         variable_same?(new_left, right, context)
 
-      _ ->
-        case Map.fetch(context.types, right) do
-          {:ok, {:var, new_right}} -> variable_same?(left, new_right, context)
-          _ -> false
-        end
+      %{^right => {:var, new_right}} ->
+        variable_same?(left, new_right, context)
+
+      %{} ->
+        false
     end
   end
 
@@ -370,11 +382,11 @@ defmodule Module.Types.Infer do
   # Bad: `{var} = var`
   # Good: `x = y; y = z; z = x`
   defp recursive_type?({:var, var} = parent, parents, context) do
-    case Map.fetch!(context.types, var) do
-      :unbound ->
+    case context.types do
+      %{^var => :unbound} ->
         false
 
-      type ->
+      %{^var => type} ->
         if type in parents do
           not Enum.all?(parents, &match?({:var, _}, &1))
         else
@@ -402,12 +414,53 @@ defmodule Module.Types.Infer do
     false
   end
 
+  @doc """
+  Checks if the type has a type var.
+  """
+  def has_unbound_var?({:var, var}, context) do
+    case context.types do
+      %{^var => :unbound} -> true
+      %{^var => type} -> has_unbound_var?(type, context)
+    end
+  end
+
+  def has_unbound_var?({:tuple, args}, context),
+    do: Enum.any?(args, &has_unbound_var?(&1, context))
+
+  def has_unbound_var?({:union, args}, context),
+    do: Enum.any?(args, &has_unbound_var?(&1, context))
+
+  def has_unbound_var?(_type, _context), do: false
+
   @doc """
   Checks if the first argument is a subtype of the second argument.
-  Only checks for simple and concrete types.
+
+  This function assumes that:
+
+    * dynamic is not considered a subtype of all other types but the top type
+    * unbound variables are not subtype of anything
+
   """
   # TODO: boolean <: false | true
   # TODO: number <: float | integer
+  # TODO: implement subtype for maps
+  def subtype?(type, type, _context), do: true
+
+  def subtype?({:var, var}, other, context) do
+    case context.types do
+      %{^var => :unbound} -> false
+      %{^var => type} -> subtype?(type, other, context)
+    end
+  end
+
+  def subtype?(other, {:var, var}, context) do
+    case context.types do
+      %{^var => :unbound} -> false
+      %{^var => type} -> subtype?(other, type, context)
+    end
+  end
+
+  def subtype?(_, :dynamic, _context), do: true
   def subtype?({:atom, boolean}, :boolean, _context) when is_boolean(boolean), do: true
   def subtype?({:atom, atom}, :atom, _context) when is_atom(atom), do: true
   def subtype?(:boolean, :atom, _context), do: true
@@ -415,18 +468,15 @@ defmodule Module.Types.Infer do
   def subtype?(:integer, :number, _context), do: true
   def subtype?({:tuple, _}, :tuple, _context), do: true
 
-  # TODO: Lift unions to unify/3?
-  def subtype?({:union, left_types}, {:union, right_types} = right_union, context) do
-    # Since we can't unify unions we give up when encountering variables
-    Enum.any?(left_types ++ right_types, &match?({:var, _}, &1)) or
-      Enum.all?(left_types, &subtype?(&1, right_union, context))
+  def subtype?({:union, left_types}, {:union, _} = right_union, context) do
+    Enum.all?(left_types, &subtype?(&1, right_union, context))
   end
 
   def subtype?(left, {:union, right_types}, context) do
     Enum.any?(right_types, &subtype?(left, &1, context))
   end
 
-  def subtype?(left, right, _context), do: left == right
+  def subtype?(_left, _right, _context), do: false
 
   @doc """
   Returns a "simplified" union using `subtype?/3` to remove redundant types.
@@ -436,17 +486,14 @@ defmodule Module.Types.Infer do
   `{boolean()} | {atom()}` will not be merged or types with variables that
   are distinct but equivalent such as `a | b when a ~ b`.
   """
-  # TODO: Translate union of all top types to dynamic()
+  def to_union([type], _context), do: type
+
   def to_union(types, context) when types != [] do
     flat_types = flatten_union(types)
 
-    if :dynamic in flat_types do
-      :dynamic
-    else
-      case unique_super_types(flat_types, context) do
-        [type] -> type
-        types -> {:union, types}
-      end
+    case unique_super_types(flat_types, context) do
+      [type] -> type
+      types -> {:union, types}
     end
   end
 
diff --git a/lib/elixir/lib/module/types/of.ex b/lib/elixir/lib/module/types/of.ex
index 0daba1f477f..d7a1ed6ba9b 100644
--- a/lib/elixir/lib/module/types/of.ex
+++ b/lib/elixir/lib/module/types/of.ex
@@ -15,7 +15,7 @@ defmodule Module.Types.Of do
   Handles open maps (with dynamic => dynamic).
   """
   def open_map(args, stack, context, fun) do
-    with {:ok, pairs, context} <- of_pairs(args, stack, context, fun) do
+    with {:ok, pairs, context} <- map_pairs(args, stack, context, fun) do
       {:ok, {:map, pairs_to_unions(pairs, context) ++ [{:optional, :dynamic, :dynamic}]}, context}
     end
   end
@@ -24,34 +24,63 @@ defmodule Module.Types.Of do
   Handles closed maps (without dynamic => dynamic).
   """
   def closed_map(args, stack, context, fun) do
-    with {:ok, pairs, context} <- of_pairs(args, stack, context, fun) do
+    with {:ok, pairs, context} <- map_pairs(args, stack, context, fun) do
       {:ok, {:map, pairs_to_unions(pairs, context)}, context}
     end
   end
 
-  defp of_pairs(pairs, stack, context, fun) do
+  defp map_pairs(pairs, stack, context, fun) do
     map_reduce_ok(pairs, context, fn {key, value}, context ->
       with {:ok, key_type, context} <- fun.(key, stack, context),
            {:ok, value_type, context} <- fun.(value, stack, context),
-           do: {:ok, {:required, key_type, value_type}, context}
+           do: {:ok, {key_type, value_type}, context}
     end)
   end
 
+  defp pairs_to_unions([{key, value}], _context), do: [{:required, key, value}]
+
   defp pairs_to_unions(pairs, context) do
-    # We are currently creating overlapping key types
+    case Enum.split_with(pairs, fn {key, _value} -> Infer.has_unbound_var?(key, context) end) do
+      {[], pairs} -> pairs_to_unions(pairs, [], context)
+      {[_ | _], pairs} -> pairs_to_unions([{:dynamic, :dynamic} | pairs], [], context)
+    end
+  end
 
-    Enum.reduce(pairs, [], fn {kind_left, key, value_left}, pairs ->
-      case List.keyfind(pairs, key, 1) do
-        {:required, ^key, value_right} ->
-          value = Infer.to_union([value_left, value_right], context)
-          List.keystore(pairs, key, 1, {:required, key, value})
+  defp pairs_to_unions([{key, value} | ahead], behind, context) do
+    {matched_ahead, values} = find_matching_values(ahead, key, [], [])
 
-        nil ->
-          [{kind_left, key, value_left} | pairs]
-      end
-    end)
+    # In case nothing matches, use the original ahead
+    ahead = matched_ahead || ahead
+
+    all_values =
+      [value | values] ++
+        find_subtype_values(ahead, key, context) ++
+        find_subtype_values(behind, key, context)
+
+    pairs_to_unions(ahead, [{key, Infer.to_union(all_values, context)} | behind], context)
+  end
+
+  defp pairs_to_unions([], acc, context) do
+    acc
+    |> Enum.sort(&Infer.subtype?(elem(&1, 0), elem(&2, 0), context))
+    |> Enum.map(fn {key, value} -> {:required, key, value} end)
+  end
+
+  defp find_subtype_values(pairs, key, context) do
+    for {pair_key, pair_value} <- pairs, Infer.subtype?(pair_key, key, context), do: pair_value
+  end
+
+  defp find_matching_values([{key, value} | ahead], key, acc, values) do
+    find_matching_values(ahead, key, acc, [value | values])
+  end
+
+  defp find_matching_values([{_, _} = pair | ahead], key, acc, values) do
+    find_matching_values(ahead, key, [pair | acc], values)
   end
 
+  defp find_matching_values([], _key, acc, [_ | _] = values), do: {Enum.reverse(acc), values}
+  defp find_matching_values([], _key, _acc, []), do: {nil, []}
+
   @doc """
   Handles structs.
   """
diff --git a/lib/elixir/test/elixir/module/types/infer_test.exs b/lib/elixir/test/elixir/module/types/infer_test.exs
index 8d3c5beb1d9..175f6946691 100644
--- a/lib/elixir/test/elixir/module/types/infer_test.exs
+++ b/lib/elixir/test/elixir/module/types/infer_test.exs
@@ -386,6 +386,31 @@ defmodule Module.Types.InferTest do
     end
   end
 
+  describe "has_unbound_var?/2" do
+    setup do
+      context = new_context()
+      {unbound_var, context} = add_var(context)
+      {bound_var, context} = add_var(context)
+      {:ok, _, context} = unify(bound_var, :integer, context)
+      %{context: context, unbound_var: unbound_var, bound_var: bound_var}
+    end
+
+    test "returns true when there are unbound vars",
+         %{context: context, unbound_var: unbound_var} do
+      assert has_unbound_var?(unbound_var, context)
+      assert has_unbound_var?({:union, [unbound_var]}, context)
+      assert has_unbound_var?({:tuple, [unbound_var]}, context)
+    end
+
+    test "returns false when there are no unbound vars",
+         %{context: context, bound_var: bound_var} do
+      refute has_unbound_var?(bound_var, context)
+      refute has_unbound_var?({:union, [bound_var]}, context)
+      refute has_unbound_var?({:tuple, [bound_var]}, context)
+      refute has_unbound_var?(:integer, context)
+    end
+  end
+
   test "subtype?/3" do
     assert subtype?({:atom, :foo}, :atom, new_context())
     assert subtype?({:atom, true}, :boolean, new_context())
diff --git a/lib/elixir/test/elixir/module/types/map_test.exs b/lib/elixir/test/elixir/module/types/map_test.exs
index 66c04dca79b..8e1893858be 100644
--- a/lib/elixir/test/elixir/module/types/map_test.exs
+++ b/lib/elixir/test/elixir/module/types/map_test.exs
@@ -16,7 +16,7 @@ defmodule Module.Types.MapTest do
     assert quoted_expr([a], %{123 => a}) == {:ok, {:map, [{:required, :integer, {:var, 0}}]}}
 
     assert quoted_expr(%{123 => :foo, 456 => :bar}) ==
-             {:ok, {:map, [{:required, :integer, {:union, [{:atom, :bar}, {:atom, :foo}]}}]}}
+             {:ok, {:map, [{:required, :integer, {:union, [{:atom, :foo}, {:atom, :bar}]}}]}}
   end
 
   test "struct" do
@@ -24,9 +24,9 @@ defmodule Module.Types.MapTest do
              {:ok,
               {:map,
                [
-                 {:required, {:atom, :foo}, {:atom, :atom}},
-                 {:required, {:atom, :baz}, {:map, []}},
                  {:required, {:atom, :bar}, :integer},
+                 {:required, {:atom, :baz}, {:map, []}},
+                 {:required, {:atom, :foo}, {:atom, :atom}},
                  {:required, {:atom, :__struct__}, {:atom, Module.Types.MapTest.Struct}}
                ]}}
 
@@ -34,9 +34,9 @@ defmodule Module.Types.MapTest do
              {:ok,
               {:map,
                [
-                 {:required, {:atom, :bar}, {:atom, :atom}},
-                 {:required, {:atom, :foo}, :integer},
                  {:required, {:atom, :baz}, {:map, []}},
+                 {:required, {:atom, :foo}, :integer},
+                 {:required, {:atom, :bar}, {:atom, :atom}},
                  {:required, {:atom, :__struct__}, {:atom, Module.Types.MapTest.Struct}}
                ]}}
   end
@@ -296,4 +296,72 @@ defmodule Module.Types.MapTest do
       assert quoted_expr([var], [map_size(var) == 1], var.foo) == {:ok, {:var, 0}}
     end
   end
+
+  test "with bound var keys" do
+    assert quoted_expr(
+             [atom, bool, true = var],
+             [is_atom(atom) and is_boolean(bool)],
+             %{atom => :atom, bool => :bool, var => true}
+           ) ==
+             {:ok,
+              {:map,
+               [
+                 {:required, {:atom, true}, {:atom, true}},
+                 {:required, :boolean, {:union, [{:atom, :bool}, {:atom, true}]}},
+                 {:required, :atom, {:union, [{:atom, :atom}, {:atom, :bool}, {:atom, true}]}}
+               ]}}
+
+    assert quoted_expr(
+             [atom, bool, true = var],
+             [is_atom(atom) and is_boolean(bool)],
+             %{var => true, bool => :bool, atom => :atom}
+           ) ==
+             {:ok,
+              {:map,
+               [
+                 {:required, {:atom, true}, {:atom, true}},
+                 {:required, :boolean, {:union, [{:atom, :bool}, {:atom, true}]}},
+                 {:required, :atom, {:union, [{:atom, :atom}, {:atom, :bool}, {:atom, true}]}}
+               ]}}
+
+    assert quoted_expr(
+             [atom, bool, true = var],
+             [is_atom(atom) and is_boolean(bool)],
+             %{var => true, atom => :atom, bool => :bool}
+           ) ==
+             {:ok,
+              {:map,
+               [
+                 {:required, {:atom, true}, {:atom, true}},
+                 {:required, :boolean, {:union, [{:atom, :bool}, {:atom, true}]}},
+                 {:required, :atom, {:union, [{:atom, :atom}, {:atom, :bool}, {:atom, true}]}}
+               ]}}
+  end
+
+  test "with unbound var keys" do
+    assert quoted_expr(
+             [var, struct],
+             (
+               map = %{var => :foo}
+               %^var{} = struct
+               map
+             )
+           ) == {:ok, {:map, [{:required, :atom, {:atom, :foo}}]}}
+
+    # If we have multiple keys, the unbound key must become required(dynamic) => dynamic
+    assert quoted_expr(
+             [var, struct],
+             (
+               map = %{var => :foo, :foo => :bar}
+               %^var{} = struct
+               map
+             )
+           ) ==
+             {:ok,
+              {:map,
+               [
+                 {:required, {:atom, :foo}, {:atom, :bar}},
+                 {:required, :dynamic, :dynamic}
+               ]}}
+  end
 end
diff --git a/lib/elixir/test/elixir/module/types/pattern_test.exs b/lib/elixir/test/elixir/module/types/pattern_test.exs
index 0f5573d43a7..c73e9d59ed9 100644
--- a/lib/elixir/test/elixir/module/types/pattern_test.exs
+++ b/lib/elixir/test/elixir/module/types/pattern_test.exs
@@ -124,7 +124,7 @@ defmodule Module.Types.PatternTest do
                {:ok,
                 {:map,
                  [
-                   {:required, :integer, {:union, [{:atom, :bar}, {:atom, :foo}]}},
+                   {:required, :integer, {:union, [{:atom, :foo}, {:atom, :bar}]}},
                    {:optional, :dynamic, :dynamic}
                  ]}}
 
@@ -147,8 +147,8 @@ defmodule Module.Types.PatternTest do
                {:ok,
                 {:map,
                  [
-                   {:required, {:atom, :bar}, {:atom, :atom}},
                    {:required, {:atom, :foo}, :integer},
+                   {:required, {:atom, :bar}, {:atom, :atom}},
                    {:required, {:atom, :__struct__}, {:atom, Module.Types.PatternTest.Struct}},
                    {:required, {:atom, :baz}, :dynamic}
                  ]}}

From 3ea1f0733f74114d76abb584622f1b1371606288 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jos=C3=A9=20Valim?= <jose.valim@dashbit.co>
Date: Sat, 12 Sep 2020 14:42:25 +0200
Subject: [PATCH 2/2] Fixes

---
 lib/elixir/lib/module/types/infer.ex | 16 +++++++++++-----
 1 file changed, 11 insertions(+), 5 deletions(-)

diff --git a/lib/elixir/lib/module/types/infer.ex b/lib/elixir/lib/module/types/infer.ex
index 6c0c32ebab6..10ccea76233 100644
--- a/lib/elixir/lib/module/types/infer.ex
+++ b/lib/elixir/lib/module/types/infer.ex
@@ -7,7 +7,6 @@ defmodule Module.Types.Infer do
   Unifies two types and returns the unified type and an updated typing context
   or an error in case of a typing conflict.
   """
-  # TODO: handle unions
   def unify(source, target, stack, context) do
     case do_unify(source, target, stack, context) do
       {:ok, type, context} ->
@@ -85,10 +84,17 @@ defmodule Module.Types.Infer do
   end
 
   defp do_unify(source, target, stack, context) do
-    if subtype?(source, target, context) do
-      {:ok, source, context}
-    else
-      error(:unable_unify, {source, target, stack}, context)
+    cond do
+      # This condition exists to handle unions with unbound vars.
+      # TODO: handle unions properly.
+      has_unbound_var?(source, context) or has_unbound_var?(target, context) ->
+        {:ok, source, context}
+
+      subtype?(source, target, context) ->
+        {:ok, source, context}
+
+      true ->
+        error(:unable_unify, {source, target, stack}, context)
     end
   end