From 9ff5b907aa3aff70a215195f30bf49801def1385 Mon Sep 17 00:00:00 2001
From: Mitchell Hanberg <mitch@mitchellhanberg.com>
Date: Wed, 17 Apr 2024 18:19:57 -0400
Subject: [PATCH] fix(completions): correctly accumulate variables in `<-`
 expressions

---
 lib/next_ls/helpers/ast_helpers.ex            |  2 +-
 lib/next_ls/helpers/ast_helpers/env.ex        | 20 ++++--
 test/next_ls/completions_test.exs             | 69 +++++++++++++++++++
 test/next_ls/helpers/ast_helpers/env_test.exs | 44 ++++++++++++
 4 files changed, 128 insertions(+), 7 deletions(-)

diff --git a/lib/next_ls/helpers/ast_helpers.ex b/lib/next_ls/helpers/ast_helpers.ex
index b0364c6a..d3ad83e4 100644
--- a/lib/next_ls/helpers/ast_helpers.ex
+++ b/lib/next_ls/helpers/ast_helpers.ex
@@ -160,7 +160,7 @@ defmodule NextLS.ASTHelpers do
            |> Zipper.zip()
            |> Zipper.find(fn
              {:@, _, [{:__cursor__, _, []}]} -> true
-             {:__cursor__, _, []} -> true
+             {:__cursor__, _, _} -> true
              {{:., _, [_, :__cursor__]}, _, _} -> true
              _ -> false
            end) do
diff --git a/lib/next_ls/helpers/ast_helpers/env.ex b/lib/next_ls/helpers/ast_helpers/env.ex
index f200d2be..d1cc376b 100644
--- a/lib/next_ls/helpers/ast_helpers/env.ex
+++ b/lib/next_ls/helpers/ast_helpers/env.ex
@@ -3,7 +3,8 @@ defmodule NextLS.ASTHelpers.Env do
   alias Sourceror.Zipper
 
   defp inside?(range, position) do
-    Sourceror.compare_positions(range.start, position) == :lt && Sourceror.compare_positions(range.end, position) == :gt
+    Sourceror.compare_positions(range.start, position) in [:lt, :eq] &&
+      Sourceror.compare_positions(range.end, position) in [:gt, :eq]
   end
 
   def build(nil) do
@@ -40,21 +41,28 @@ defmodule NextLS.ASTHelpers.Env do
 
             Map.update!(acc, :variables, &(vars ++ &1))
 
-          {match_op, _, [pm | _]} when match_op in [:<-] ->
+          {match_op, _, [pm | rhs]} when match_op in [:<-] ->
             up_node = zipper |> Zipper.up() |> Zipper.node()
 
             # in_match operator comes with for and with normally, so we need to 
             # check if we are inside the parent node, which is the for/with
-            is_inside =
-              with {_, _, _} <- up_node,
-                   range when not is_nil(range) <- Sourceror.get_range(up_node) do
+            is_inside_p =
+              with {_, _, _} <- up_node, range when not is_nil(range) <- Sourceror.get_range(up_node) do
                 inside?(range, position)
               else
                 _ ->
                   false
               end
 
-            if is_inside do
+            is_inside_rhs =
+              with range when not is_nil(range) <- Sourceror.get_range(rhs) do
+                inside?(range, position)
+              else
+                _ ->
+                  false
+              end
+
+            if is_inside_p and not is_inside_rhs do
               {_, vars} =
                 Macro.prewalk(pm, [], fn node, acc ->
                   case node do
diff --git a/test/next_ls/completions_test.exs b/test/next_ls/completions_test.exs
index 88481f17..f4a03385 100644
--- a/test/next_ls/completions_test.exs
+++ b/test/next_ls/completions_test.exs
@@ -19,6 +19,20 @@ defmodule NextLS.CompletionsTest do
     end
   end
 
+  defmacrop assert_match({:not, _, [{:in, _, [left, right]}]}) do
+    quote do
+      refute Enum.any?(unquote(right), fn x ->
+               match?(unquote(left), x)
+             end),
+             """
+             found a match inside of list, expected none
+
+             left: #{unquote(Macro.to_string(left))}
+             right: #{inspect(unquote(right), pretty: true)}
+             """
+    end
+  end
+
   @moduletag tmp_dir: true, root_paths: ["my_proj"]
 
   setup %{tmp_dir: tmp_dir} do
@@ -482,4 +496,59 @@ defmodule NextLS.CompletionsTest do
       %{"data" => _, "documentation" => _, "insertText" => "capture_log", "kind" => 3, "label" => "capture_log/2"} in results
     )
   end
+
+  test "completions inside generator rhs", %{client: client, foo: foo} do
+    uri = uri(foo)
+
+    did_open(client, foo, """
+    defmodule Foo do
+      def run() do
+        var = "hi"
+
+        for thing <- v do
+        end
+
+      end
+    end
+    """)
+
+    request client, %{
+      method: "textDocument/completion",
+      id: 2,
+      jsonrpc: "2.0",
+      params: %{
+        textDocument: %{
+          uri: uri
+        },
+        position: %{
+          line: 4,
+          character: 18
+        }
+      }
+    }
+
+    assert_result 2, [
+      %{
+        "data" => nil,
+        "documentation" => "",
+        "insertText" => "var",
+        "kind" => 6,
+        "label" => "var"
+      },
+      %{
+        "data" => nil,
+        "documentation" => _,
+        "insertText" => "var!",
+        "kind" => 3,
+        "label" => "var!/1"
+      },
+      %{
+        "data" => nil,
+        "documentation" => _,
+        "insertText" => "var!",
+        "kind" => 3,
+        "label" => "var!/2"
+      }
+    ]
+  end
 end
diff --git a/test/next_ls/helpers/ast_helpers/env_test.exs b/test/next_ls/helpers/ast_helpers/env_test.exs
index bf68ee20..1be03d13 100644
--- a/test/next_ls/helpers/ast_helpers/env_test.exs
+++ b/test/next_ls/helpers/ast_helpers/env_test.exs
@@ -156,6 +156,50 @@ defmodule NextLS.ASTHelpers.EnvTest do
 
       assert actual.variables == ["entries"]
     end
+
+    test "comprehension lhs of generator do not leak into rhs " do
+      code = """
+      defmodule Foo do
+        def one(entries) do
+          for entry <- entries,
+              not_me <- __cursor__() do
+            :ok
+          end
+        end
+
+        def two do
+          baz = :bar
+        end
+      end
+      """
+
+      actual = run(code)
+
+      assert actual.variables == ["entries", "entry"]
+    end
+
+    test "multiple generators and filters in comprehension" do
+      code = """
+      defmodule Foo do
+        def one(entries) do
+          for entry <- entries,
+              foo = do_something(),
+              bar <- foo do
+            __cursor__()
+            :ok
+          end
+        end
+
+        def two do
+          baz = :bar
+        end
+      end
+      """
+
+      actual = run(code)
+
+      assert actual.variables == ["entries", "entry", "foo", "bar"]
+    end
   end
 
   defp run(code) do