From 9e4f69bb55d3c69a1e61d071fb8647c0a7789160 Mon Sep 17 00:00:00 2001
From: Mitchell Hanberg <mitch@mitchellhanberg.com>
Date: Thu, 16 May 2024 09:30:00 -0400
Subject: [PATCH] fix(completions): more accurate inside with/for

---
 flake.nix                                |  2 +-
 priv/monkey/_next_ls_private_compiler.ex | 15 +++--
 test/next_ls/completions_test.exs        | 73 ++++++++++++++++++++++++
 3 files changed, 85 insertions(+), 5 deletions(-)

diff --git a/flake.nix b/flake.nix
index a5ae57ab..4fee3fde 100644
--- a/flake.nix
+++ b/flake.nix
@@ -115,7 +115,7 @@
             src = self.outPath;
             inherit version elixir;
             pname = "next-ls-deps";
-            hash = "sha256-sIV8/KS5hcIiqk5sSwcIEnPFZNvefzvNyt6af829ri8=";
+            hash = "sha256-qdJf3A+k28J2EDtaC8pLE7HgqzKuCjCWmfezx62wyUs=";
             mixEnv = "prod";
           };
 
diff --git a/priv/monkey/_next_ls_private_compiler.ex b/priv/monkey/_next_ls_private_compiler.ex
index d673b3ea..031dbf0a 100644
--- a/priv/monkey/_next_ls_private_compiler.ex
+++ b/priv/monkey/_next_ls_private_compiler.ex
@@ -1538,8 +1538,8 @@ if Version.match?(System.version(), ">= 1.17.0-dev") do
     defp expand_local(_meta, fun, args, state, env) when fun in [:for, :with] do
       {params, blocks} =
         Enum.split_while(args, fn
-          {:<-, _, _} -> true
-          _ -> false
+          [{:do, _} | _] -> false
+          _ -> true
         end)
 
       {_, state, penv} =
@@ -1549,9 +1549,16 @@ if Version.match?(System.version(), ">= 1.17.0-dev") do
         end
 
       {blocks, state} =
-        for {type, block} <- blocks, reduce: {[], state} do
+        for {type, block} <- List.first(blocks, []), reduce: {[], state} do
           {acc, state} ->
-            {res, state, _env} = expand(block, state, penv)
+            env =
+              if type == :do do
+                penv
+              else
+                env
+              end
+
+            {res, state, env} = expand(block, state, env)
             {[{type, res} | acc], state}
         end
 
diff --git a/test/next_ls/completions_test.exs b/test/next_ls/completions_test.exs
index 6c24ce65..bcfce92e 100644
--- a/test/next_ls/completions_test.exs
+++ b/test/next_ls/completions_test.exs
@@ -700,6 +700,79 @@ defmodule NextLS.CompletionsTest do
     assert_match %{"kind" => 6, "label" => "ast1"} not in results
   end
 
+  test "with else", %{client: client, foo: foo} do
+    uri = uri(foo)
+
+    did_open(client, foo, """
+    defmodule Foo do
+      def run(doug) do
+        completion_item =
+          with {:ok, darrel} <- completion_item.data do
+            darrel
+          else
+            %{"uri" => uri, "data" => data} ->
+              d
+
+      end
+    end
+    """)
+
+    request client, %{
+      method: "textDocument/completion",
+      id: 2,
+      jsonrpc: "2.0",
+      params: %{
+        textDocument: %{uri: uri},
+        position: %{
+          line: 7,
+          character: 11
+        }
+      }
+    }
+
+    assert_result 2, results
+
+    assert_match %{"kind" => 6, "label" => "doug"} in results
+    assert_match %{"kind" => 6, "label" => "data"} in results
+
+    assert_match %{"kind" => 6, "label" => "darrel"} not in results
+  end
+
+  test "for comprehension", %{client: client, foo: foo} do
+    uri = uri(foo)
+
+    did_open(client, foo, """
+    defmodule Foo do
+      def run(items) do
+        for item <- items,
+            iname = item.name,
+            String.starts_with?(name, "Mitch") do
+          i
+      end
+    end
+    """)
+
+    request client, %{
+      method: "textDocument/completion",
+      id: 2,
+      jsonrpc: "2.0",
+      params: %{
+        textDocument: %{uri: uri},
+        position: %{
+          line: 5,
+          character: 7
+        }
+      }
+    }
+
+    assert_result 2, results
+
+    assert_match %{"kind" => 6, "label" => "item"} in results
+    assert_match %{"kind" => 6, "label" => "iname"} in results
+
+    assert_match %{"kind" => 6, "label" => "items"} in results
+  end
+
   test "variables show up in test blocks", %{client: client, foo: foo} do
     uri = uri(foo)