From ed875265f54994911774ac05c3a1b9adb65b80e7 Mon Sep 17 00:00:00 2001
From: Tim Gent <tim.gent@gmail.com>
Date: Sun, 8 Jan 2023 08:35:01 +0000
Subject: [PATCH] Find definition to respect function arity (#173)

* Go to definition respects arity of function

* Refactor code for finding function definition

* Remove unrequired TODO

* Remove unrequired column arg to Definition.find

* Simplify test case for find definition with arity

Co-authored-by: Tim Gent <tim.gent@multiverse.io>
---
 lib/elixir_sense.ex                      |  1 +
 lib/elixir_sense/providers/definition.ex | 41 +++++++++++++++++++++++-
 test/elixir_sense/definition_test.exs    | 26 +++++++++++++++
 3 files changed, 67 insertions(+), 1 deletion(-)

diff --git a/lib/elixir_sense.ex b/lib/elixir_sense.ex
index 7b13ce05..e224ace2 100644
--- a/lib/elixir_sense.ex
+++ b/lib/elixir_sense.ex
@@ -111,6 +111,7 @@ defmodule ElixirSense do
 
         Definition.find(
           subject,
+          line,
           env,
           buffer_file_metadata.mods_funs_to_positions,
           calls,
diff --git a/lib/elixir_sense/providers/definition.ex b/lib/elixir_sense/providers/definition.ex
index 7a435eb3..3c171536 100644
--- a/lib/elixir_sense/providers/definition.ex
+++ b/lib/elixir_sense/providers/definition.ex
@@ -20,6 +20,7 @@ defmodule ElixirSense.Providers.Definition do
   """
   @spec find(
           String.t(),
+          pos_integer,
           State.Env.t(),
           State.mods_funs_to_positions_t(),
           list(State.CallInfo.t()),
@@ -27,6 +28,7 @@ defmodule ElixirSense.Providers.Definition do
         ) :: %Location{} | nil
   def find(
         subject,
+        line,
         %State.Env{
           aliases: aliases,
           module: module,
@@ -63,6 +65,8 @@ defmodule ElixirSense.Providers.Definition do
         subject
         |> Source.split_module_and_func(module, aliases)
         |> find_function_or_module(
+          line,
+          calls,
           mods_funs_to_positions,
           env,
           metadata_types,
@@ -95,6 +99,8 @@ defmodule ElixirSense.Providers.Definition do
 
   defp find_function_or_module(
          {module, function},
+         line,
+         calls,
          mods_funs_to_positions,
          env,
          metadata_types,
@@ -104,6 +110,8 @@ defmodule ElixirSense.Providers.Definition do
     unless {module, function} in visited do
       do_find_function_or_module(
         {module, function},
+        line,
+        calls,
         mods_funs_to_positions,
         env,
         metadata_types,
@@ -115,6 +123,8 @@ defmodule ElixirSense.Providers.Definition do
 
   defp do_find_function_or_module(
          {{:attribute, _attr} = type, function},
+         line,
+         calls,
          mods_funs_to_positions,
          env,
          metadata_types,
@@ -125,6 +135,8 @@ defmodule ElixirSense.Providers.Definition do
       {:atom, module} ->
         do_find_function_or_module(
           {Introspection.expand_alias(module, env.aliases), function},
+          line,
+          calls,
           mods_funs_to_positions,
           env,
           metadata_types,
@@ -139,6 +151,8 @@ defmodule ElixirSense.Providers.Definition do
 
   defp do_find_function_or_module(
          {nil, :super},
+         line,
+         calls,
          mods_funs_to_positions,
          %State.Env{scope: {function, arity}, module: module} = env,
          metadata_types,
@@ -150,6 +164,8 @@ defmodule ElixirSense.Providers.Definition do
         # overridable function is most likely defined by __using__ macro
         do_find_function_or_module(
           {origin, :__using__},
+          line,
+          calls,
           mods_funs_to_positions,
           env,
           metadata_types,
@@ -164,6 +180,8 @@ defmodule ElixirSense.Providers.Definition do
 
   defp do_find_function_or_module(
          {module, function},
+         line,
+         calls,
          mods_funs_to_positions,
          env,
          metadata_types,
@@ -188,7 +206,10 @@ defmodule ElixirSense.Providers.Definition do
         nil
 
       {mod, fun, true} ->
-        case mods_funs_to_positions[{mod, fun, nil}] || metadata_types[{mod, fun, nil}] do
+        fn_definition = find_fn_definition(mod, fun, line, mods_funs_to_positions, calls)
+
+        case fn_definition || mods_funs_to_positions[{mod, fun, nil}] ||
+               metadata_types[{mod, fun, nil}] do
           nil ->
             Location.find_source({mod, fun}, current_module)
 
@@ -216,4 +237,22 @@ defmodule ElixirSense.Providers.Definition do
         end
     end
   end
+
+  defp find_fn_definition(mod, fun, line, mods_funs_to_positions, calls) do
+    mods_funs_to_positions
+    |> Enum.find(fn
+      {{^mod, ^fun, fn_arity}, %{positions: fn_positions}} when not is_nil(fn_arity) ->
+        case calls do
+          [] -> Enum.any?(fn_positions, fn {fn_line, _fn_col} -> fn_line == line end)
+          [%{arity: call_arity} | _] -> fn_arity == call_arity
+        end
+
+      _ ->
+        false
+    end)
+    |> case do
+      {_, mod_fun_info} -> mod_fun_info
+      nil -> nil
+    end
+  end
 end
diff --git a/test/elixir_sense/definition_test.exs b/test/elixir_sense/definition_test.exs
index b9f68c52..99a5b1e7 100644
--- a/test/elixir_sense/definition_test.exs
+++ b/test/elixir_sense/definition_test.exs
@@ -157,6 +157,32 @@ defmodule ElixirSense.Providers.DefinitionTest do
     assert read_line(file, {line, column}) =~ "ElixirSenseExample.Macros.go"
   end
 
+  test "find definition for the correct arity of function - on fn call" do
+    buffer = """
+    defmodule MyModule do
+      def main, do: my_func("a", "b")
+      #               ^ 
+      def my_func, do: "not this one"
+      def my_func(a, b), do: a <> b
+    end
+    """
+
+    assert %Location{type: :function, file: nil, line: 5, column: 7} =
+             ElixirSense.definition(buffer, 2, 18)
+  end
+
+  test "find definition for the correct arity of function - on fn definition" do
+    buffer = """
+    defmodule MyModule do
+      def my_func, do: "not this one"
+      def my_func(a, b), do: a <> b
+    end
+    """
+
+    assert %Location{type: :function, file: nil, line: 3, column: 7} =
+             ElixirSense.definition(buffer, 3, 9)
+  end
+
   test "find definition of delegated functions" do
     buffer = """
     defmodule MyModule do