Completion: Use existing specs to infer function signatures and vice-versa

apps/common/lib/lexical/ast/env.ex

@@ -193,4 +193,46 @@ defmodule Lexical.Ast.Env do
   def empty?(string) when is_binary(string) do
     String.trim(string) == ""
   end
+
+  @doc """
+  Returns the position of the next non-whitespace token on a line after `env.position`.
+  """
+  @spec next_significant_position(t) :: {:ok, Position.t()} | :error
+  def next_significant_position(%__MODULE__{} = env) do
+    find_significant_position(env.document, env.position.line + 1, 1)
+  end
+
+  @doc """
+  Returns the position of the next non-whitespace token on a line before `env.position`.
+  """
+  @spec prev_significant_position(t) :: {:ok, Position.t()} | :error
+  def prev_significant_position(%__MODULE__{} = env) do
+    find_significant_position(env.document, env.position.line - 1, -1)
+  end
+
+  defp find_significant_position(%Document{} = document, line, inc_by) do
+    case Document.fetch_text_at(document, line) do
+      {:ok, text} ->
+        case fetch_leading_whitespace_count(text) do
+          {:ok, count} ->
+            {:ok, Position.new(document, line, count + 1)}
+
+          :error ->
+            find_significant_position(document, line + inc_by, inc_by)
+        end
+
+      :error ->
+        :error
+    end
+  end
+
+  defp fetch_leading_whitespace_count(string, count \\ 0)
+
+  defp fetch_leading_whitespace_count(<<" ", rest::binary>>, count) do
+    fetch_leading_whitespace_count(rest, count + 1)
+  end
+
+  defp fetch_leading_whitespace_count(<<>>, _count), do: :error
+  defp fetch_leading_whitespace_count(<<"\n" <> _::binary>>, _count), do: :error
+  defp fetch_leading_whitespace_count(<<_non_whitespace::binary>>, count), do: {:ok, count}
 end

apps/server/lib/lexical/server/code_intelligence/completion/translations/macro.ex

@@ -1,6 +1,8 @@
 defmodule Lexical.Server.CodeIntelligence.Completion.Translations.Macro do
+  alias Lexical.Ast
   alias Lexical.Ast.Env
   alias Lexical.Document
+  alias Lexical.Document.Position
   alias Lexical.RemoteControl.Completion.Candidate
   alias Lexical.Server.CodeIntelligence.Completion.SortScope
   alias Lexical.Server.CodeIntelligence.Completion.Translatable
@@ -23,41 +25,11 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.Macro do
   end
 
   def translate(%Candidate.Macro{name: "def", arity: 2} = macro, builder, env) do
-    label = "#{macro.name} (define a function)"
-
-    snippet = """
-    def ${1:name}($2) do
-      $0
-    end
-    """
-
-    env
-    |> builder.snippet(snippet,
-      detail: macro.spec,
-      kind: :class,
-      label: label,
-      filter_text: macro.name
-    )
-    |> builder.set_sort_scope(SortScope.global())
+    function_snippet("def", "define a function", macro, builder, env)
   end
 
   def translate(%Candidate.Macro{name: "defp", arity: 2} = macro, builder, env) do
-    label = "#{macro.name} (define a private function)"
-
-    snippet = """
-    defp ${1:name}($2) do
-      $0
-    end
-    """
-
-    env
-    |> builder.snippet(snippet,
-      detail: macro.spec,
-      kind: :class,
-      label: label,
-      filter_text: macro.name
-    )
-    |> builder.set_sort_scope(SortScope.global())
+    function_snippet("defp", "define a private function", macro, builder, env)
   end
 
   def translate(%Candidate.Macro{name: "defmodule"} = macro, builder, env) do
@@ -596,7 +568,72 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.Macro do
     :skip
   end
 
-  def suggest_module_name(%Document{} = document) do
+  defp function_snippet(kind, label, %Candidate.Macro{} = macro, builder, env) do
+    label = "#{macro.name} (#{label})"
+
+    snippet =
+      with {:ok, %Position{} = position} <- Env.prev_significant_position(env),
+           {:ok, name, args} <- extract_spec_name_and_args(env, position) do
+        args_snippet =
+          case suggest_arg_names(args) do
+            [] ->
+              ""
+
+            names ->
+              placeholders =
+                names
+                |> Enum.with_index(1)
+                |> Enum.map_join(", ", fn {name, i} -> "${#{i}:#{name}}" end)
+
+              "(" <> placeholders <> ")"
+          end
+
+        """
+        #{kind} #{name}#{args_snippet} do
+          $0
+        end
+        """
+      else
+        _ ->
+          """
+          #{kind} ${1:name}($2) do
+            $0
+          end
+          """
+      end
+
+    env
+    |> builder.snippet(snippet,
+      detail: macro.spec,
+      kind: :class,
+      label: label,
+      filter_text: macro.name
+    )
+    |> builder.set_sort_scope(SortScope.global())
+  end
+
+  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, args}
+      else
+        {:ok, name, []}
+      end
+    else
+      _ -> :error
+    end
+  end
+
+  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()

apps/server/lib/lexical/server/code_intelligence/completion/translations/module_attribute.ex

@@ -1,6 +1,9 @@
 defmodule Lexical.Server.CodeIntelligence.Completion.Translations.ModuleAttribute do
+  alias Lexical.Ast
   alias Lexical.Ast.Env
+  alias Lexical.Document.Position
   alias Lexical.RemoteControl.Completion.Candidate
+  alias Lexical.Server.CodeIntelligence.Completion.SortScope
   alias Lexical.Server.CodeIntelligence.Completion.Translatable
   alias Lexical.Server.CodeIntelligence.Completion.Translations
 
@@ -70,6 +73,19 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.ModuleAttribut
     end
   end
 
+  def translate(%Candidate.ModuleAttribute{name: "@spec"}, builder, env) do
+    case fetch_range(env) do
+      {:ok, range} ->
+        [
+          maybe_specialized_spec_snippet(builder, env, range),
+          basic_spec_snippet(builder, env, range)
+        ]
+
+      :error ->
+        :skip
+    end
+  end
+
   def translate(%Candidate.ModuleAttribute{} = attribute, builder, env) do
     case fetch_range(env) do
       {:ok, range} ->
@@ -106,4 +122,57 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.ModuleAttribut
         {:cont, acc}
     end)
   end
+
+  defp maybe_specialized_spec_snippet(builder, %Env{} = env, range) do
+    with {:ok, %Position{} = position} <- Env.next_significant_position(env),
+         {:ok, [{maybe_def, _, [call, _]} | _]} when maybe_def in [:def, :defp] <-
+           Ast.path_at(env.analysis, position),
+         {function_name, _, args} <- call do
+      specialized_spec_snippet(builder, env, range, function_name, args)
+    else
+      _ -> nil
+    end
+  end
+
+  defp specialized_spec_snippet(builder, env, range, function_name, args) do
+    name = to_string(function_name)
+
+    args_snippet =
+      case args do
+        nil ->
+          ""
+
+        list ->
+          Enum.map_join(1..length(list), ", ", &"${#{&1}:term()}")
+      end
+
+    snippet = ~s"""
+    @spec #{name}(#{args_snippet}) :: ${0:term()}
+    """
+
+    env
+    |> builder.text_edit_snippet(snippet, range,
+      detail: "Typespec",
+      kind: :property,
+      label: "@spec #{name}"
+    )
+    |> builder.set_sort_scope(SortScope.global(false, 0))
+  end
+
+  defp basic_spec_snippet(builder, env, range) do
+    snippet = ~S"""
+    @spec ${1:function}(${2:term()}) :: ${3:term()}
+    def ${1:function}(${4:args}) do
+      $0
+    end
+    """
+
+    env
+    |> builder.text_edit_snippet(snippet, range,
+      detail: "Typespec",
+      kind: :property,
+      label: "@spec"
+    )
+    |> builder.set_sort_scope(SortScope.global(false, 1))
+  end
 end

apps/server/test/lexical/server/code_intelligence/completion/translations/macro_test.exs

@@ -46,6 +46,82 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.MacroTest do
       assert apply_completion(completion) == "def ${1:name}($2) do\n  $0\nend"
     end
 
+    test "def preceeded by a @spec with args", %{project: project} do
+      source = ~q[
+        @spec my_function(term(), term()) :: term()
+        def|
+      ]
+
+      assert {:ok, completion} =
+               project
+               |> complete(source)
+               |> fetch_completion("def ")
+
+      assert apply_completion(completion) == ~q[
+        @spec my_function(term(), term()) :: term()
+        def my_function(${1:arg_1}, ${2:arg_2}) do
+          $0
+        end
+      ]
+    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()
+        def|
+      ]
+
+      assert {:ok, completion} =
+               project
+               |> complete(source)
+               |> fetch_completion("def ")
+
+      assert apply_completion(completion) == ~q[
+        @spec my_function :: term()
+        def my_function do
+          $0
+        end
+      ]
+    end
+
+    test "defp preceeded by a @spec with args", %{project: project} do
+      source = ~q[
+        @spec my_function(term(), term()) :: term()
+        def|
+      ]
+
+      assert {:ok, completion} =
+               project
+               |> complete(source)
+               |> fetch_completion("defp ")
+
+      assert apply_completion(completion) == ~q[
+        @spec my_function(term(), term()) :: term()
+        defp my_function(${1:arg_1}, ${2:arg_2}) do
+          $0
+        end
+      ]
+    end
+
     test "defp only has a single completion", %{project: project} do
       assert {:ok, completion} =
                project