From c1f3856306c158252c39b290bf833ec414a37b2c Mon Sep 17 00:00:00 2001
From: Zach Allaun <zach.allaun@gmail.com>
Date: Fri, 26 Jul 2024 11:04:52 -0400
Subject: [PATCH] Infer function arg names from named types in spec

---
 .../completion/translations/macro.ex          | 32 +++++++++++++------
 .../completion/translations/macro_test.exs    | 19 +++++++++++
 2 files changed, 42 insertions(+), 9 deletions(-)

diff --git a/apps/server/lib/lexical/server/code_intelligence/completion/translations/macro.ex b/apps/server/lib/lexical/server/code_intelligence/completion/translations/macro.ex
index 79a252c33..bc24ce976 100644
--- a/apps/server/lib/lexical/server/code_intelligence/completion/translations/macro.ex
+++ b/apps/server/lib/lexical/server/code_intelligence/completion/translations/macro.ex
@@ -573,12 +573,19 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.Macro do
 
     snippet =
       with {:ok, %Position{} = position} <- Env.prev_significant_position(env),
-           {:ok, name, arity} <- extract_spec_name_arity(env, position) do
+           {:ok, name, args} <- extract_spec_name_and_args(env, position) do
         args_snippet =
-          if arity == 0 do
-            ""
-          else
-            "(" <> Enum.map_join(1..arity, ", ", &"${#{&1}:arg_#{&1}}") <> ")"
+          case suggest_arg_names(args) do
+            [] ->
+              ""
+
+            names ->
+              placeholders =
+                names
+                |> Enum.with_index(1)
+                |> Enum.map_join(", ", fn {name, i} -> "${#{i}:#{name}}" end)
+
+              "(" <> placeholders <> ")"
           end
 
         """
@@ -605,21 +612,28 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.Macro do
     |> builder.set_sort_scope(SortScope.global())
   end
 
-  defp extract_spec_name_arity(%Env{} = env, %Position{} = position) do
+  defp extract_spec_name_and_args(%Env{} = env, %Position{} = position) do
     with {:ok, [maybe_spec | _]} <- Ast.path_at(env.analysis, position),
          {:@, _, [{:spec, _, [typespec]}]} <- maybe_spec,
          {:"::", _, [{name, _, args}, _return]} <- typespec do
       if is_list(args) do
-        {:ok, name, length(args)}
+        {:ok, name, args}
       else
-        {:ok, name, 0}
+        {:ok, name, []}
       end
     else
       _ -> :error
     end
   end
 
-  def suggest_module_name(%Document{} = document) do
+  defp suggest_arg_names(args) do
+    Enum.with_index(args, fn
+      {:"::", _, [{name, _, nil}, _]}, _i when is_atom(name) -> name
+      _, i -> "arg_#{i + 1}"
+    end)
+  end
+
+  defp suggest_module_name(%Document{} = document) do
     result =
       document.path
       |> Path.split()
diff --git a/apps/server/test/lexical/server/code_intelligence/completion/translations/macro_test.exs b/apps/server/test/lexical/server/code_intelligence/completion/translations/macro_test.exs
index ad40cdc68..e95de440e 100644
--- a/apps/server/test/lexical/server/code_intelligence/completion/translations/macro_test.exs
+++ b/apps/server/test/lexical/server/code_intelligence/completion/translations/macro_test.exs
@@ -65,6 +65,25 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.MacroTest do
       ]
     end
 
+    test "def preceeded by a @spec with named args", %{project: project} do
+      source = ~q[
+        @spec my_function(x :: term(), y :: term(), term()) :: term()
+        def|
+      ]
+
+      assert {:ok, completion} =
+               project
+               |> complete(source)
+               |> fetch_completion("def ")
+
+      assert apply_completion(completion) == ~q[
+        @spec my_function(x :: term(), y :: term(), term()) :: term()
+        def my_function(${1:x}, ${2:y}, ${3:arg_3}) do
+          $0
+        end
+      ]
+    end
+
     test "def preceeded by a @spec without args", %{project: project} do
       source = ~q[
         @spec my_function :: term()