Allow captures to be reconstructed on type system pretty printing

Author: josevalim

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

@@ -141,7 +141,8 @@ defmodule Module.Types.Helpers do
     end)
   end
 
-  defp hint, do: :elixir_errors.prefix(:hint)
+  @doc "The hint prefix"
+  def hint, do: :elixir_errors.prefix(:hint)
 
   @doc """
   Collect traces from variables in expression.
@@ -340,6 +341,13 @@ defmodule Module.Types.Helpers do
       {{:., _, [mod, fun]}, meta, args} ->
         erl_to_ex(mod, fun, args, meta)
 
+      {:fn, meta, [{:->, _, [_args, return]}]} = expr ->
+        if meta[:capture] do
+          {:&, meta, [return]}
+        else
+          expr
+        end
+
       {:&, amp_meta, [{:/, slash_meta, [{{:., dot_meta, [mod, fun]}, call_meta, []}, arity]}]} ->
         {mod, fun} =
           case :elixir_rewrite.erl_to_ex(mod, fun, arity) do
@@ -385,6 +393,13 @@ defmodule Module.Types.Helpers do
           case
         end
 
+      {var, meta, context} = expr when is_atom(var) and is_atom(context) ->
+        if is_integer(meta[:capture]) do
+          {:&, meta, [meta[:capture]]}
+        else
+          expr
+        end
+
       other ->
         other
     end)

lib/elixir/src/elixir_fn.erl

@@ -131,7 +131,7 @@ capture_expr(Meta, Expr, S, E, Escaped, ArgsType) ->
       {expand, Fn, S, E};
     {EExpr, EDict} ->
       EVars = validate(Meta, EDict, 1, E),
-      Fn = {fn, Meta, [{'->', Meta, [EVars, EExpr]}]},
+      Fn = {fn, [{capture, true} | Meta], [{'->', Meta, [EVars, EExpr]}]},
       {expand, Fn, S, E}
   end.
 
@@ -154,7 +154,7 @@ escape({'&', Meta, [Pos]}, E, Dict) when is_integer(Pos), Pos > 0 ->
       {Var, Dict};
     error ->
       Next = elixir_module:next_counter(?key(E, module)),
-      Var = {capture, [{counter, Next} | Meta], nil},
+      Var = {capture, [{counter, Next}, {capture, Pos} | Meta], nil},
       {Var, orddict:store(Pos, Var, Dict)}
   end;
 escape({'&', Meta, [Pos]}, E, _Dict) when is_integer(Pos) ->

lib/elixir/test/elixir/kernel/expansion_test.exs

@@ -1221,8 +1221,14 @@ defmodule Kernel.ExpansionTest do
 
     test "keeps position meta on & variables" do
       assert expand(Code.string_to_quoted!("& &1")) |> clean_meta([:counter]) ==
-               {:fn, [{:line, 1}],
-                [{:->, [{:line, 1}], [[{:capture, [line: 1], nil}], {:capture, [line: 1], nil}]}]}
+               {:fn, [capture: true, line: 1],
+                [
+                  {:->, [line: 1],
+                   [
+                     [{:capture, [capture: 1, line: 1], nil}],
+                     {:capture, [capture: 1, line: 1], nil}
+                   ]}
+                ]}
     end
 
     test "removes no_parens when expanding 0-arity capture to fn" do

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

@@ -256,6 +256,18 @@ defmodule Module.Types.ExprTest do
                  (dynamic(map()) -> :map)
              """
     end
+
+    test "capture printing" do
+      assert typeerror!(123 = &{:ok, &1}) == """
+             the following pattern will never match:
+
+                 123 = &{:ok, &1}
+
+             because the right-hand side has type:
+
+                 (dynamic() -> dynamic({:ok, term()}))
+             """
+    end
   end
 
   describe "remotes" do