Perform type inference using reverse arrows on all non-branching constructs (#14145)

Author: josevalim

lib/elixir/lib/module/parallel_checker.ex

@@ -208,7 +208,7 @@ defmodule Module.ParallelChecker do
   or if the function does not exist return `{:error, :function}`.
   """
   @spec fetch_export(cache(), module(), atom(), arity()) ::
-          {:ok, mode(), binary() | nil, {:infer, [term()]} | :none}
+          {:ok, mode(), binary() | nil, {:infer, [term()] | nil, [term()]} | :none}
           | :badmodule
           | {:badfunction, mode()}
   def fetch_export({checker, table}, module, fun, arity) do
@@ -451,10 +451,15 @@ defmodule Module.ParallelChecker do
 
   defp cache_chunk(table, module, exports) do
     Enum.each(exports, fn {{fun, arity}, info} ->
-      # TODO: Match on signature directly in Elixir v1.22+
+      sig =
+        case info do
+          %{sig: {key, _, _} = sig} when key in [:infer, :strong] -> sig
+          _ -> :none
+        end
+
       :ets.insert(
         table,
-        {{module, {fun, arity}}, Map.get(info, :deprecated), Map.get(info, :sig, :none)}
+        {{module, {fun, arity}}, Map.get(info, :deprecated), sig}
       )
     end)
 

lib/elixir/lib/module/types.ex

@@ -281,11 +281,18 @@ defmodule Module.Types do
           context = fresh_context(context)
 
           try do
-            {args_types, context} =
+            {trees, context} =
               Pattern.of_head(args, guards, expected, {:infer, expected}, meta, stack, context)
 
             {return_type, context} =
-              Expr.of_expr(body, stack, context)
+              Expr.of_expr(body, Descr.term(), body, stack, context)
+
+            args_types =
+              if stack.mode == :traversal do
+                expected
+              else
+                Pattern.of_domain(trees, expected, context)
+              end
 
             {type_index, inferred} =
               add_inferred(inferred, args_types, return_type, total - 1, [])
@@ -301,8 +308,19 @@ defmodule Module.Types do
           end
       end)
 
-    inferred = {:infer, Enum.reverse(clauses_types)}
-    {inferred, mapping, restore_context(context, clauses_context)}
+    domain =
+      case clauses_types do
+        [_] ->
+          nil
+
+        _ ->
+          clauses_types
+          |> Enum.map(fn {args, _} -> args end)
+          |> Enum.zip_with(fn types -> Enum.reduce(types, &Descr.union/2) end)
+      end
+
+    inferred = {:infer, domain, Enum.reverse(clauses_types)}
+    {inferred, mapping, restore_context(clauses_context, context)}
   end
 
   # We check for term equality of types as an optimization
@@ -399,7 +417,7 @@ defmodule Module.Types do
     %{context | vars: %{}, failed: false}
   end
 
-  defp restore_context(%{vars: vars, failed: failed}, later_context) do
+  defp restore_context(later_context, %{vars: vars, failed: failed}) do
     %{later_context | vars: vars, failed: failed}
   end