Support struct patterns

3 files changed, 115 insertions, 12 deletions

diff --git a/lib/elixir/lib/module/types/infer.ex b/lib/elixir/lib/module/types/infer.ex
index ade413a35..957efc828 100644
--- a/lib/elixir/lib/module/types/infer.ex
+++ b/lib/elixir/lib/module/types/infer.ex
@@ -12,6 +12,8 @@ defmodule Module.Types.Infer do
     end
   end
 
+  ## PATTERNS
+
   # :atom
   def of_pattern(atom, context) when is_atom(atom) do
     {:ok, {:literal, atom}, context}
@@ -102,26 +104,70 @@ defmodule Module.Types.Infer do
 
   # %{...}
   def of_pattern({:%{}, _meta, args} = expr, context) do
-    result =
-      expr_stack(expr, context, fn context ->
-        map_reduce_ok(args, context, fn {key, value}, context ->
-          with {:ok, key_type, context} <- of_pattern(key, context),
-               {:ok, value_type, context} <- of_pattern(value, context),
-               do: {:ok, {key_type, value_type}, context}
-        end)
-      end)
+    # TODO: Create unions when types for keys overlap, the pattern:
+    #       `%{1 => :foo, 2 => :bar}`
+    #       should create the type:
+    #       `%{integer() => :foo | :bar}`
 
-    case result do
+    case expr_stack(expr, context, &of_pairs(args, &1)) do
       {:ok, pairs, context} -> {:ok, {:map, pairs}, context}
       {:error, reason} -> {:error, reason}
     end
   end
 
-  # TODO: structs
-  def of_pattern({:%, _meta, _args}, context) do
-    {:ok, {:map, []}, context}
+  # %var{...}
+  def of_pattern({:%, _meta1, [var, {:%{}, _meta2, args}]} = expr, context) when is_var(var) do
+    # TODO: Create unions when types for keys overlap, see above
+
+    expr_stack(expr, context, fn context ->
+      with {:ok, pairs, context} <- of_pairs(args, context),
+           {var_type, context} = new_var(var, context),
+           {:ok, _, context} <- unify(var_type, :atom, context) do
+        pairs = [{{:literal, :__struct__}, var_type} | pairs]
+        {:ok, {:map, pairs}, context}
+      end
+    end)
   end
 
+  # %Struct{...}
+  def of_pattern({:%, _meta1, [module, {:%{}, _meta2, args}]} = expr, context)
+      when is_atom(module) do
+    # TODO: Create unions when types for keys overlap, see above
+
+    struct_pairs =
+      Enum.map(:maps.remove(:__struct__, module.__struct__()), fn {key, value} ->
+        {term_to_type(key), term_to_type(value)}
+      end)
+
+    case expr_stack(expr, context, &of_pairs(args, &1)) do
+      {:ok, pattern_pairs, context} -> {:ok, {:map, pattern_pairs ++ struct_pairs}, context}
+      {:error, reason} -> {:error, reason}
+    end
+  end
+
+  defp of_pairs(pairs, context) do
+    map_reduce_ok(pairs, context, fn {key, value}, context ->
+      with {:ok, key_type, context} <- of_pattern(key, context),
+           {:ok, value_type, context} <- of_pattern(value, context),
+           do: {:ok, {key_type, value_type}, context}
+    end)
+  end
+
+  def term_to_type(term) when is_atom(term), do: {:literal, term}
+  def term_to_type(term) when is_bitstring(term), do: :binary
+  def term_to_type(term) when is_float(term), do: :float
+  def term_to_type(term) when is_function(term), do: :fun
+  def term_to_type(term) when is_integer(term), do: :integer
+  def term_to_type(term) when is_pid(term), do: :pid
+  def term_to_type(term) when is_port(term), do: :port
+  def term_to_type(term) when is_reference(term), do: :reference
+  # NOTE: We may want to go into more detail for higher order types
+  def term_to_type(term) when is_list(term), do: {:cons, :dynamic, :dynamic}
+  def term_to_type(term) when is_map(term), do: {:map, []}
+  def term_to_type(term) when is_tuple(term), do: :tuple
+
+  ## GUARDS
+
   def of_guard(guard, context) do
     expr_stack(guard, context, fn context ->
       # Flatten `foo and bar` to list of type constraints
@@ -194,6 +240,8 @@ defmodule Module.Types.Infer do
 
   defp type_guard(_other), do: {:error, :unknown_guard}
 
+  ## BINARY PATTERNS
+
   # binary-pattern :: specifier
   defp of_binary({:"::", _meta, [expr, specifiers]} = full_expr, context) do
     specifiers = List.flatten(collect_specifiers(specifiers))
diff --git a/lib/elixir/test/elixir/module/types/infer_test.exs b/lib/elixir/test/elixir/module/types/infer_test.exs
index 228ea667c..eb968b70b 100644
--- a/lib/elixir/test/elixir/module/types/infer_test.exs
+++ b/lib/elixir/test/elixir/module/types/infer_test.exs
@@ -85,10 +85,49 @@ defmodule Module.Types.InferTest do
       assert quoted_pattern(%{a: :b}) == {:ok, {:map, [{{:literal, :a}, {:literal, :b}}]}}
       assert quoted_pattern(%{123 => a}) == {:ok, {:map, [{:integer, {:var, 0}}]}}
 
+      # TODO
+      # assert quoted_pattern(%{123 => :foo, 456 => :bar}) ==
+      #          {:ok, {:map, [{:integer, {:union, [{:literal, :foo}, {:literal, :bar}]}}]}}
+
       assert {:error, {{:unable_unify, {:literal, :foo}, :integer, _, _}, _}} =
                quoted_pattern(%{a: a = 123, b: a = :foo})
     end
 
+    test "struct" do
+      defmodule :"Elixir.Module.Types.InferTest.Struct" do
+        defstruct foo: :atom, bar: 123, baz: %{}
+      end
+
+      assert quoted_pattern(%:"Elixir.Module.Types.InferTest.Struct"{}) ==
+               {:ok,
+                {:map,
+                 [
+                   {{:literal, :bar}, :integer},
+                   {{:literal, :baz}, {:map, []}},
+                   {{:literal, :foo}, {:literal, :atom}}
+                 ]}}
+
+      # TODO
+      # assert quoted_pattern(%:"Elixir.Module.Types.InferTest.Struct"{foo: 123, bar: :atom}) ==
+      #          {:ok,
+      #           {:map,
+      #            [
+      #              {{:literal, :bar}, {:literal, :atom}},
+      #              {{:literal, :baz}, {:map, []}},
+      #              {{:literal, :foo}, :integer}
+      #            ]}}
+    end
+
+    test "struct var" do
+      assert quoted_pattern(%var{}) == {:ok, {:map, [{{:literal, :__struct__}, :atom}]}}
+
+      assert quoted_pattern(%var{foo: 123}) ==
+               {:ok, {:map, [{{:literal, :__struct__}, :atom}, {{:literal, :foo}, :integer}]}}
+
+      assert quoted_pattern(%var{foo: var}) ==
+               {:ok, {:map, [{{:literal, :__struct__}, :atom}, {{:literal, :foo}, :atom}]}}
+    end
+
     test "binary" do
       assert quoted_pattern(<<"foo"::binary>>) == {:ok, :binary}
       assert quoted_pattern(<<123::integer>>) == {:ok, :binary}
@@ -366,4 +405,12 @@ defmodule Module.Types.InferTest do
 
     # assert to_union([{:tuple, [:boolean]}, {:tuple, [:atom]}], new_context()) == {:tuple, [:atom]}
   end
+
+  test "term_to_type/1" do
+    assert term_to_type(true) == {:literal, true}
+    assert term_to_type(:foo) == {:literal, :foo}
+    assert term_to_type(%{}) == {:map, []}
+    assert term_to_type({}) == :tuple
+    assert term_to_type(make_ref()) == :reference
+  end
 end
diff --git a/lib/elixir/test/elixir/module/types_test.exs b/lib/elixir/test/elixir/module/types_test.exs
index 1120a8e23..ec120a787 100644
--- a/lib/elixir/test/elixir/module/types_test.exs
+++ b/lib/elixir/test/elixir/module/types_test.exs
@@ -117,6 +117,14 @@ defmodule Module.TypesTest do
       assert {:error, {{:unable_unify, {:literal, true}, {:literal, false}, _, _}, _}} =
                quoted_clause([%{true: false} = foo, %{true: true} = foo])
     end
+
+    test "struct var guard" do
+      assert quoted_clause([%var{}], [:erlang.is_atom(var)]) ==
+               {:ok, [{:map, [{{:literal, :__struct__}, :atom}]}]}
+
+      assert {:error, {{:unable_unify, :integer, :atom, _, _}, _}} =
+               quoted_clause([%var{}], [:erlang.is_integer(var)])
+    end
   end
 
   test "format_type/1" do