fix(engine): unify import lookup to correctly resolve local calls (#547)

When working on macro expansion indexing, I figured out that
`Entity.resolve` is ignoring imports, leading to wrong mfa. This causes
index search to find nothing, in which case we fallback to elixirsense
(this may be the reason some cases seem to work). At least function
capture case did not work.

Author: Draggu

apps/engine/lib/engine/analyzer.ex

@@ -13,6 +13,7 @@ defmodule Engine.Analyzer do
 
   defdelegate aliases_at(analysis, position), to: Aliases, as: :at
   defdelegate imports_at(analysis, position), to: Imports, as: :at
+  defdelegate import_module_for(analysis, position, fun, arity), to: Imports, as: :module_for
 
   @spec requires_at(Analysis.t(), Position.t()) :: [module()]
   def requires_at(%Analysis{} = analysis, %Position{} = position) do
@@ -39,20 +40,14 @@ defmodule Engine.Analyzer do
   end
 
   def resolve_local_call(%Analysis{} = analysis, %Position{} = position, function_name, arity) do
-    maybe_imported_mfa =
-      analysis
-      |> imports_at(position)
-      |> Enum.find(fn
-        {_, ^function_name, ^arity} -> true
-        _ -> false
-      end)
-
-    if is_nil(maybe_imported_mfa) do
-      aliases = aliases_at(analysis, position)
-      current_module = aliases[[:__MODULE__]]
-      {current_module, function_name, arity}
-    else
-      maybe_imported_mfa
+    case import_module_for(analysis, position, function_name, arity) do
+      {:ok, module} ->
+        {module, function_name, arity}
+
+      :error ->
+        aliases = aliases_at(analysis, position)
+        current_module = aliases[[:__MODULE__]]
+        {current_module, function_name, arity}
     end
   end
 

apps/engine/lib/engine/analyzer/imports.ex

@@ -19,6 +19,68 @@ defmodule Engine.Analyzer.Imports do
     end
   end
 
+  @doc """
+  Returns the source module for an imported function by walking import declarations in scope.
+  """
+  @spec module_for(Analysis.t(), Position.t(), atom(), non_neg_integer()) ::
+          {:ok, module()} | :error
+  def module_for(%Analysis{} = analysis, %Position{} = position, function_name, arity) do
+    case Analysis.scopes_at(analysis, position) do
+      [%Scope{} = scope | _] ->
+        end_line = Scope.end_line(scope, position)
+
+        scope.imports
+        |> Enum.sort_by(& &1.range.start.line)
+        |> Enum.take_while(&(&1.range.start.line <= end_line))
+        |> Enum.reverse(kernel_imports(scope))
+        |> Enum.find_value(:error, fn %Import{} = import ->
+          module = Aliases.resolve_at(scope, import.module, import.range.start.line)
+
+          if import_allows?(module, import.selector, function_name, arity) do
+            {:ok, module}
+          else
+            false
+          end
+        end)
+
+      _ ->
+        :error
+    end
+  end
+
+  defp import_allows?(module, :all, fun, arity) do
+    not Loader.ensure_loaded?(module) or
+      {fun, arity} in function_and_arities_for_module(module, :functions) or
+      {fun, arity} in function_and_arities_for_module(module, :macros)
+  end
+
+  defp import_allows?(module, [only: :functions], fun, arity) do
+    not Loader.ensure_loaded?(module) or
+      {fun, arity} in function_and_arities_for_module(module, :functions)
+  end
+
+  defp import_allows?(module, [only: :macros], fun, arity) do
+    not Loader.ensure_loaded?(module) or
+      {fun, arity} in function_and_arities_for_module(module, :macros)
+  end
+
+  defp import_allows?(module, [only: :sigils], fun, arity) do
+    not Loader.ensure_loaded?(module) or
+      {fun, arity} in function_and_arities_for_module(module, :sigils)
+  end
+
+  defp import_allows?(_module, [only: fns], fun, arity) when is_list(fns),
+    do: {fun, arity} in fns
+
+  defp import_allows?(module, [except: fns], fun, arity) when is_list(fns) do
+    {fun, arity} not in fns and
+      (not Loader.ensure_loaded?(module) or
+         {fun, arity} in function_and_arities_for_module(module, :functions) or
+         {fun, arity} in function_and_arities_for_module(module, :macros))
+  end
+
+  defp import_allows?(_module, _selector, _fun, _arity), do: false
+
   @spec imports(Scope.t(), Scope.scope_position()) :: [Scope.import_mfa()]
   def imports(%Scope{} = scope, position \\ :end) do
     scope

apps/engine/lib/engine/code_intelligence/entity.ex

@@ -117,12 +117,18 @@ defmodule Engine.CodeIntelligence.Entity do
   end
 
   defp resolve({:local_arity, chars}, node_range, analysis, position) do
-    current_module = current_module(analysis, position)
+    fun = List.to_atom(chars)
 
     with {:ok, %Zipper{node: {:/, _, [_, {:__block__, _, [arity]}]}} = zipper} <-
            Ast.zipper_at(analysis.document, position),
          true <- inside_capture?(zipper) do
-      {:ok, {:call, current_module, List.to_atom(chars), arity}, node_range}
+      module =
+        case fetch_module_for_function(analysis, position, fun, arity) do
+          {:ok, mod} -> mod
+          _ -> current_module(analysis, position)
+        end
+
+      {:ok, {:call, module, fun, arity}, node_range}
     else
       _ ->
         {:error, :not_found}
@@ -491,22 +497,10 @@ defmodule Engine.CodeIntelligence.Entity do
 
   defp fetch_module_for_function(analysis, position, function_name, arity) do
     with :error <- fetch_module_for_local_function(analysis, position, function_name, arity) do
-      fetch_module_for_imported_function(analysis, position, function_name, arity)
+      Engine.Analyzer.import_module_for(analysis, position, function_name, arity)
     end
   end
 
-  defp fetch_module_for_imported_function(analysis, position, function_name, arity) do
-    analysis
-    |> Engine.Analyzer.imports_at(position)
-    |> Enum.find_value({:error, :not_found}, fn
-      {imported_module, ^function_name, ^arity} ->
-        {:ok, imported_module}
-
-      _ ->
-        false
-    end)
-  end
-
   defp fetch_module_for_local_function(analysis, position, function_name, arity) do
     with {:ok, current_module} <- Engine.Analyzer.current_module(analysis, position),
          true <- function_exported?(current_module, function_name, arity) do
@@ -537,8 +531,13 @@ defmodule Engine.CodeIntelligence.Entity do
       range = Ast.Range.fetch!(zipper.node, analysis.document)
       range = put_in(range.end.character, range.start.character + function_name_length)
 
-      current_module = current_module(analysis, position)
-      {:ok, {:call, current_module, local_func_name, arity}, range}
+      module =
+        case fetch_module_for_function(analysis, position, local_func_name, arity) do
+          {:ok, mod} -> mod
+          _ -> current_module(analysis, position)
+        end
+
+      {:ok, {:call, module, local_func_name, arity}, range}
     else
       _ ->
         {:error, :not_found}

apps/engine/test/engine/analyzer/imports_test.exs

@@ -37,6 +37,11 @@ defmodule WithSigils do
   end
 end
 
+defmodule ShadowsKernel do
+  # to_string/1 exists in Kernel — used to test that explicit imports shadow Kernel
+  def to_string(x), do: x
+end
+
 defmodule Engine.Ast.Analysis.ImportsTest do
   use ExUnit.Case
 
@@ -55,6 +60,14 @@ defmodule Engine.Ast.Analysis.ImportsTest do
     |> Analyzer.imports_at(position)
   end
 
+  def module_for_cursor(text, function, arity) do
+    {position, document} = pop_cursor(text, as: :document)
+
+    document
+    |> Ast.analyze()
+    |> Analyzer.import_module_for(position, function, arity)
+  end
+
   def assert_imported(imports, module) do
     functions = module.__info__(:functions)
     macros = module.__info__(:macros)
@@ -373,6 +386,31 @@ defmodule Engine.Ast.Analysis.ImportsTest do
     end
   end
 
+  describe "module_for/4" do
+    test "finds kernel functions without any explicit import" do
+      assert {:ok, Kernel} = module_for_cursor(~q[
+        defmodule Foo do
+          |
+        end
+      ], :to_string, 1)
+    end
+
+    test "explicit import shadows kernel for the same function" do
+      assert {:ok, ShadowsKernel} = module_for_cursor(~q[
+        import ShadowsKernel
+        |
+      ], :to_string, 1)
+    end
+
+    test "returns error when function does not exist in any import" do
+      assert :error = module_for_cursor(~q[
+        defmodule Foo do
+          |
+        end
+      ], :does_not_exist_function, 0)
+    end
+  end
+
   describe "import scopes" do
     test "an import defined in a named function" do
       imports =

apps/engine/test/engine/code_intelligence/entity_test.exs

@@ -899,6 +899,106 @@ defmodule Engine.CodeIntelligence.EntityTest do
       assert {:ok, {:call, Parent, :from, 1}, resolved_range} = resolve(code)
       assert resolved_range =~ ~S[«from»(doc)]
     end
+
+    test "imported from unloaded module (local_call)" do
+      code = ~q[
+        defmodule MyProject.B do
+          def foo(), do: 1
+        end
+
+        defmodule MyProject do
+          import MyProject.B
+
+          def test() do
+            |foo()
+          end
+        end
+      ]
+
+      assert {:ok, {:call, MyProject.B, :foo, 0}, _} = resolve(code)
+    end
+
+    test "imported from unloaded module (local_or_var)" do
+      code = ~q[
+        defmodule MyProject.B do
+          def foo(), do: 1
+        end
+
+        defmodule MyProject do
+          import MyProject.B
+
+          def test() do
+            |foo
+          end
+        end
+      ]
+
+      assert {:ok, {:call, MyProject.B, :foo, 0}, _} = resolve(code)
+    end
+
+    test "imported from unloaded module (module-level call)" do
+      code = ~q[
+        defmodule MyProject.B do
+          def foo(), do: 1
+        end
+
+        defmodule MyProject do
+          import MyProject.B
+
+          |foo()
+        end
+      ]
+
+      assert {:ok, {:call, MyProject.B, :foo, 0}, _} = resolve(code)
+    end
+
+    test "imported from unloaded module (capture)" do
+      code = ~q[
+        defmodule MyProject.B do
+          def foo(), do: 1
+        end
+
+        defmodule MyProject do
+          import MyProject.B
+
+          def test() do
+            &foo|/0
+          end
+        end
+      ]
+
+      assert {:ok, {:call, MyProject.B, :foo, 0}, _} = resolve(code)
+    end
+
+    test "imported function in a capture" do
+      code = ~q[
+        defmodule Parent do
+          import Forge.Ast
+
+          def function do
+            &from|/1
+          end
+        end
+      ]
+
+      assert {:ok, {:call, Forge.Ast, :from, 1}, resolved_range} = resolve(code)
+      assert resolved_range =~ ~S[&«from»/1]
+    end
+
+    test "imported function in a capture (cursor before name)" do
+      code = ~q[
+        defmodule Parent do
+          import Forge.Ast
+
+          def function do
+            &|from/1
+          end
+        end
+      ]
+
+      assert {:ok, {:call, Forge.Ast, :from, 1}, resolved_range} = resolve(code)
+      assert resolved_range =~ ~S[&«from»/1]
+    end
   end
 
   describe "type resolve/2" do