More typing improvements (#10333)

1. Move push_expr_stack to unify and keep helpers typing free

2. Add get_var! when accessing variables so it helps us find
   variables that have not been properly processed

3. Properly process all variables found missing on step 2

Author: josevalim

lib/elixir/lib/module/types.ex

@@ -57,7 +57,7 @@ defmodule Module.Types do
   end
 
   defp warnings_from_clause(args, guards, body, def_expr, stack, context) do
-    head_stack = push_expr_stack(def_expr, stack)
+    head_stack = Unify.push_expr_stack(def_expr, stack)
 
     with {:ok, _types, context} <- Pattern.of_head(args, guards, head_stack, context),
          {:ok, _type, context} <- Expr.of_expr(body, stack, context) do

lib/elixir/lib/module/types/expr.ex

@@ -109,8 +109,7 @@ defmodule Module.Types.Expr do
 
   # var
   def of_expr(var, _stack, context) when is_var(var) do
-    {type, context} = new_var(var, context)
-    {:ok, type, context}
+    {:ok, get_var!(var, context), context}
   end
 
   # {left, right}
@@ -379,6 +378,13 @@ defmodule Module.Types.Expr do
          do: {:ok, context}
   end
 
+  defp for_clause({:<<>>, _, [{:<-, _, [pattern, expr]}]}, stack, context) do
+    # TODO: the compiler guarantees pattern is a binary but we need to check expr is a binary
+    with {:ok, _pattern_type, context} <- Pattern.of_pattern(pattern, stack, context),
+         {:ok, _expr_type, context} <- of_expr(expr, stack, context),
+         do: {:ok, context}
+  end
+
   defp for_clause(list, stack, context) when is_list(list) do
     reduce_ok(list, context, &for_option(&1, stack, &2))
   end

lib/elixir/lib/module/types/helpers.ex

@@ -25,16 +25,6 @@ defmodule Module.Types.Helpers do
   def get_meta({_, meta, _}), do: meta
   def get_meta(_other), do: []
 
-  @doc """
-  Push expression to stack.
-
-  The expression stack is used to give the context where a type variable
-  was refined when show a type conflict error.
-  """
-  def push_expr_stack(expr, stack) do
-    %{stack | last_expr: expr}
-  end
-
   @doc """
   Like `Enum.reduce/3` but only continues while `fun` returns `{:ok, acc}`
   and stops on `{:error, reason}`.

lib/elixir/lib/module/types/of.ex

@@ -6,10 +6,10 @@ defmodule Module.Types.Of do
   @prefix quote(do: ...)
   @suffix quote(do: ...)
 
-  alias Module.Types.Unify
   alias Module.ParallelChecker
 
   import Module.Types.Helpers
+  import Module.Types.Unify
 
   @doc """
   Handles open maps (with dynamic => dynamic).
@@ -40,7 +40,7 @@ defmodule Module.Types.Of do
   defp pairs_to_unions([{key, value}], _context), do: [{:required, key, value}]
 
   defp pairs_to_unions(pairs, context) do
-    case Enum.split_with(pairs, fn {key, _value} -> Unify.has_unbound_var?(key, context) end) do
+    case Enum.split_with(pairs, fn {key, _value} -> has_unbound_var?(key, context) end) do
       {[], pairs} -> pairs_to_unions(pairs, [], context)
       {[_ | _], pairs} -> pairs_to_unions([{:dynamic, :dynamic} | pairs], [], context)
     end
@@ -57,17 +57,17 @@ defmodule Module.Types.Of do
         find_subtype_values(ahead, key, context) ++
         find_subtype_values(behind, key, context)
 
-    pairs_to_unions(ahead, [{key, Unify.to_union(all_values, context)} | behind], context)
+    pairs_to_unions(ahead, [{key, to_union(all_values, context)} | behind], context)
   end
 
   defp pairs_to_unions([], acc, context) do
     acc
-    |> Enum.sort(&Unify.subtype?(elem(&1, 0), elem(&2, 0), context))
+    |> Enum.sort(&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, Unify.subtype?(pair_key, key, context), do: pair_value
+    for {pair_key, pair_value} <- pairs, subtype?(pair_key, key, context), do: pair_value
   end
 
   defp find_matching_values([{key, value} | ahead], key, acc, values) do
@@ -153,16 +153,11 @@ defmodule Module.Types.Of do
 
       true ->
         with {:ok, type, context} <- fun.(expr, stack, context),
-             {:ok, _type, context} <- Unify.unify(type, expected_type, stack, context),
+             {:ok, _type, context} <- unify(type, expected_type, stack, context),
              do: {:ok, context}
     end
   end
 
-  # TODO: Remove this clause once we properly handle comprehensions
-  defp binary_segment({:<-, _, _}, _stack, context, _fun) do
-    {:ok, context}
-  end
-
   # Collect binary type specifiers,
   # from `<<pattern::integer-size(10)>>` collect `integer`
   defp collect_binary_specifier({:-, _meta, [left, right]}, fun) do

lib/elixir/lib/module/types/pattern.ex

@@ -129,8 +129,8 @@ defmodule Module.Types.Pattern do
   end
 
   # ^var
-  def of_pattern({:^, _meta, [var]}, stack, context) do
-    of_pattern(var, stack, context)
+  def of_pattern({:^, _meta, [var]}, _stack, context) do
+    {:ok, get_var!(var, context), context}
   end
 
   # var
@@ -194,16 +194,11 @@ defmodule Module.Types.Pattern do
     end
   end
 
-  # %^var{...}
-  def of_pattern({:%, meta1, [{:^, _meta2, [var]}, args]}, stack, context) do
-    of_pattern({:%, meta1, [var, args]}, stack, context)
-  end
-
-  # %var{...}
+  # %var{...} and %^var{...}
   def of_pattern({:%, _meta1, [var, {:%{}, _meta2, args}]} = expr, stack, context) do
     stack = push_expr_stack(expr, stack)
 
-    with {var_type, context} = new_var(var, context),
+    with {:ok, var_type, context} = of_pattern(var, stack, context),
          {:ok, _, context} <- unify(var_type, :atom, stack, context),
          {:ok, {:map, pairs}, context} <- Of.open_map(args, stack, context, &of_pattern/3) do
       {:ok, {:map, [{:required, {:atom, :__struct__}, var_type} | pairs]}, context}

lib/elixir/lib/module/types/unify.ex

@@ -301,6 +301,23 @@ defmodule Module.Types.Unify do
 
   defp error(type, reason, context), do: {:error, {type, reason, context}}
 
+  @doc """
+  Push expression to stack.
+
+  The expression stack is used to give the context where a type variable
+  was refined when show a type conflict error.
+  """
+  def push_expr_stack(expr, stack) do
+    %{stack | last_expr: expr}
+  end
+
+  @doc """
+  Gets a variable.
+  """
+  def get_var!(var, context) do
+    Map.fetch!(context.vars, var_name(var))
+  end
+
   @doc """
   Adds a variable to the typing context and returns its type variable.
   If the variable has already been added, return the existing type variable.

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

@@ -297,7 +297,7 @@ defmodule Module.Types.MapTest do
     end
   end
 
-  test "with bound var keys" do
+  test "map creation with bound var keys" do
     assert quoted_expr(
              [atom, bool, true = var],
              [is_atom(atom) and is_boolean(bool)],
@@ -341,7 +341,7 @@ defmodule Module.Types.MapTest do
                ]}}
   end
 
-  test "with unbound var keys" do
+  test "map creation with unbound var keys" do
     assert quoted_expr(
              [var, struct],
              (