Completion: infer spec signature from existing function

Closes #797

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

@@ -1,6 +1,10 @@
 defmodule Lexical.Server.CodeIntelligence.Completion.Translations.ModuleAttribute 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
   alias Lexical.Server.CodeIntelligence.Completion.Translations
 
@@ -70,6 +74,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 +123,76 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.ModuleAttribut
         {:cont, acc}
     end)
   end
+
+  defp maybe_specialized_spec_snippet(builder, %Env{} = env, range) do
+    next_line = env.position.line + 1
+
+    with %Position{} = position <- next_non_whitespace_position(env.document, next_line),
+         {: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 next_non_whitespace_position(%Document{} = document, line) do
+    case Document.fetch_text_at(document, line) do
+      {:ok, text} ->
+        case String.split(text, ~r/\w+/, parts: 2) do
+          [whitespace, _] ->
+            character = String.length(whitespace) + 1
+            Position.new(document, line, character)
+
+          _ ->
+            next_non_whitespace_position(document, line + 1)
+        end
+
+      :error ->
+        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.local(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.local(false, 1))
+  end
 end

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

@@ -147,4 +147,137 @@ defmodule Lexical.Server.CodeIntelligence.Completion.Translations.ModuleAttribut
       ]
     end
   end
+
+  describe "@spec completion" do
+    test "with no function following", %{project: project} do
+      source = ~q[
+        defmodule MyModule do
+          @spe|
+        end
+      ]
+
+      assert {:ok, completion} =
+               project
+               |> complete(source)
+               |> fetch_completion("@spec")
+
+      assert apply_completion(completion) == ~q[
+        defmodule MyModule do
+          @spec ${1:function}(${2:term()}) :: ${3:term()}
+        def ${1:function}(${4:args}) do
+          $0
+        end
+      end
+      ]
+    end
+
+    test "with a function with args after it", %{project: project} do
+      source = ~q[
+        defmodule MyModule do
+          @spe|
+          def my_function(arg1, arg2, arg3) do
+            :ok
+          end
+        end
+      ]
+
+      assert {:ok, [spec_my_function, spec]} =
+               project
+               |> complete(source)
+               |> fetch_completion(kind: :property)
+
+      assert spec_my_function.label == "@spec my_function"
+
+      assert apply_completion(spec_my_function) == ~q[
+        defmodule MyModule do
+          @spec my_function(${1:term()}, ${2:term()}, ${3:term()}) :: ${0:term()}
+          def my_function(arg1, arg2, arg3) do
+            :ok
+          end
+        end
+      ]
+
+      assert spec.label == "@spec"
+
+      assert apply_completion(spec) == ~q[
+        defmodule MyModule do
+          @spec ${1:function}(${2:term()}) :: ${3:term()}
+        def ${1:function}(${4:args}) do
+          $0
+        end
+          def my_function(arg1, arg2, arg3) do
+            :ok
+          end
+        end
+      ]
+    end
+
+    test "with a function without args after it", %{project: project} do
+      source = ~q[
+        defmodule MyModule do
+          @spe|
+          def my_function do
+            :ok
+          end
+        end
+      ]
+
+      assert {:ok, [spec_my_function, spec]} =
+               project
+               |> complete(source)
+               |> fetch_completion(kind: :property)
+
+      assert spec_my_function.label == "@spec my_function"
+
+      assert apply_completion(spec_my_function) == ~q[
+        defmodule MyModule do
+          @spec my_function() :: ${0:term()}
+          def my_function do
+            :ok
+          end
+        end
+      ]
+
+      assert spec.label == "@spec"
+
+      assert apply_completion(spec) == ~q[
+        defmodule MyModule do
+          @spec ${1:function}(${2:term()}) :: ${3:term()}
+        def ${1:function}(${4:args}) do
+          $0
+        end
+          def my_function do
+            :ok
+          end
+        end
+      ]
+    end
+
+    test "with a private function after it", %{project: project} do
+      source = ~q[
+        defmodule MyModule do
+          @spe|
+          defp my_function(arg1, arg2, arg3) do
+            :ok
+          end
+        end
+      ]
+
+      assert {:ok, [spec_my_function, _spec]} =
+               project
+               |> complete(source)
+               |> fetch_completion(kind: :property)
+
+      assert spec_my_function.label == "@spec my_function"
+
+      assert apply_completion(spec_my_function) == ~q[
+        defmodule MyModule do
+          @spec my_function(${1:term()}, ${2:term()}, ${3:term()}) :: ${0:term()}
+          defp my_function(arg1, arg2, arg3) do
+            :ok
+          end
+        end
+      ]
+    end
+  end
 end