From 9c2ff68a7a0ead32bb1c356742b992903b41c440 Mon Sep 17 00:00:00 2001
From: Mitchell Hanberg <mitch@mitchellhanberg.com>
Date: Wed, 24 Apr 2024 21:22:36 -0400
Subject: [PATCH] fix: accuracy of get_surrounding_module (#440)

---
 lib/next_ls/helpers/ast_helpers.ex        | 71 ++++++++++++++++++-----
 test/next_ls/helpers/ast_helpers_test.exs | 40 ++++---------
 2 files changed, 66 insertions(+), 45 deletions(-)

diff --git a/lib/next_ls/helpers/ast_helpers.ex b/lib/next_ls/helpers/ast_helpers.ex
index 55fca7f6..5f78c464 100644
--- a/lib/next_ls/helpers/ast_helpers.ex
+++ b/lib/next_ls/helpers/ast_helpers.ex
@@ -155,26 +155,52 @@ defmodule NextLS.ASTHelpers do
     end)
   end
 
+  defp sourceror_inside?(range, position) do
+    Sourceror.compare_positions(range.start, position) in [:lt, :eq] &&
+      Sourceror.compare_positions(range.end, position) in [:gt, :eq]
+  end
+
   @spec get_surrounding_module(ast :: Macro.t(), position :: Position.t()) :: {:ok, Macro.t()} | {:error, String.t()}
   def get_surrounding_module(ast, position) do
-    defm =
+    # TODO: this should take elixir positions and not LSP positions
+    position = [line: position.line + 1, column: position.character + 1]
+
+    {_zipper, acc} =
       ast
-      |> Macro.prewalker()
-      |> Enum.filter(fn node -> match?({:defmodule, _, _}, node) end)
-      |> Enum.filter(fn {_, ctx, _} ->
-        position.line + 1 - ctx[:line] >= 0
+      |> Zipper.zip()
+      |> Zipper.traverse_while(nil, fn tree, acc ->
+        node = Zipper.node(tree)
+        node_range = Sourceror.get_range(node)
+
+        is_inside =
+          with nil <- node_range do
+            false
+          else
+            _ -> sourceror_inside?(node_range, position)
+          end
+
+        acc =
+          with true <- is_inside,
+               {:defmodule, _, _} <- node do
+            node
+          else
+            _ -> acc
+          end
+
+        cond do
+          is_inside and match?({_, _, [_ | _]}, node) ->
+            {:cont, tree, acc}
+
+          is_inside and match?({_, _, []}, node) ->
+            {:halt, tree, acc}
+
+          true ->
+            {:cont, tree, acc}
+        end
       end)
-      |> Enum.min_by(
-        fn {_, ctx, _} ->
-          abs(ctx[:line] - 1 - position.line)
-        end,
-        fn -> nil end
-      )
-
-    if defm do
-      {:ok, defm}
-    else
-      {:error, "no defmodule definition"}
+
+    with {:ok, nil} <- {:ok, acc} do
+      {:error, :not_found}
     end
   end
 
@@ -193,4 +219,17 @@ defmodule NextLS.ASTHelpers do
       zipper -> {:ok, zipper}
     end
   end
+
+  def top(nil, acc, _callback), do: acc
+
+  def top(%Zipper{path: nil} = zipper, acc, callback), do: callback.(Zipper.node(zipper), zipper, acc)
+
+  def top(zipper, acc, callback) do
+    node = Zipper.node(zipper)
+    acc = callback.(node, zipper, acc)
+
+    zipper = Zipper.up(zipper)
+
+    top(zipper, acc, callback)
+  end
 end
diff --git a/test/next_ls/helpers/ast_helpers_test.exs b/test/next_ls/helpers/ast_helpers_test.exs
index c378da58..343739f1 100644
--- a/test/next_ls/helpers/ast_helpers_test.exs
+++ b/test/next_ls/helpers/ast_helpers_test.exs
@@ -91,26 +91,26 @@ defmodule NextLS.ASTHelpersTest do
         end
         """)
 
-      lines = 1..3
+      for {line, character} <- [{0, 2}, {1, 1}, {4, 0}, {5, 1}, {8, 2}] do
+        position = %Position{line: line, character: character}
 
-      for line <- lines do
-        position = %Position{line: line, character: 0}
+        assert {:ok, {:defmodule, _, [{:__aliases__, _, [:Test]} | _]}} =
+                 ASTHelpers.get_surrounding_module(ast, position)
+      end
+
+      for {line, character} <- [{1, 2}, {1, 6}, {2, 5}, {3, 3}] do
+        position = %Position{line: line, character: character}
 
         assert {:ok, {:defmodule, _, [{:__aliases__, _, [:Foo]} | _]}} =
                  ASTHelpers.get_surrounding_module(ast, position)
       end
 
-      lines = 5..7
-
-      for line <- lines do
-        position = %Position{line: line, character: 0}
+      for {line, character} <- [{5, 4}, {6, 1}, {7, 0}, {7, 3}] do
+        position = %Position{line: line, character: character}
 
         assert {:ok, {:defmodule, _, [{:__aliases__, _, [:Bar]} | _]}} =
                  ASTHelpers.get_surrounding_module(ast, position)
       end
-
-      position = %Position{line: 0, character: 0}
-      assert {:ok, {:defmodule, _, [{:__aliases__, _, [:Test]} | _]}} = ASTHelpers.get_surrounding_module(ast, position)
     end
 
     test "errors out when it can't find a module" do
@@ -120,25 +120,7 @@ defmodule NextLS.ASTHelpersTest do
         """)
 
       position = %Position{line: 0, character: 0}
-      assert {:error, "no defmodule definition"} = ASTHelpers.get_surrounding_module(ast, position)
-    end
-
-    test "it finds the nearest surrounding module" do
-      {:ok, ast} =
-        Spitfire.parse("""
-        defmodule Test do
-          alias Foo
-          alias Bar
-          alias Baz
-
-          defmodule Quix do
-            defstruct [:key]
-          end
-        end
-        """)
-
-      position = %Position{line: 4, character: 0}
-      assert {:ok, {:defmodule, _, [{:__aliases__, _, [:Test]} | _]}} = ASTHelpers.get_surrounding_module(ast, position)
+      assert {:error, :not_found} = ASTHelpers.get_surrounding_module(ast, position)
     end
   end
 end