From a015889b399bf215763e0306eae48c90fe2a4d89 Mon Sep 17 00:00:00 2001
From: Steve Cohen <scohen@scohen.org>
Date: Tue, 12 Mar 2024 11:18:24 -0700
Subject: [PATCH 1/4] Find references for variables

This PR implements find references for variables and fixes a couple
bugs and missed cases in the variable indexer.

There's a notable edge case present in the implementation. A rebound
variable that's assigned to a block will effectively shadow the
variable of that name inside the block. I believe this is a limitation
of the information we get back from the indexer, so another approach
might be necessary in the future if this becomes an issue. An example
follows:

```elixir
name = "Stinky"
name = if some_condition() do
  "Smelly"
else
  name
end
```

Searching for references for the `name` variable on the first line
will return no results, but searching for references to the rebind on
the second line will return the reference in the `else` clause in the if.
---
 .../code_intelligence/entity.ex               |   5 +-
 .../code_intelligence/references.ex           |  40 +-
 .../code_intelligence/variable.ex             | 258 +++++++++
 .../search/indexer/extractors/variable.ex     |  55 +-
 .../search/indexer/source/reducer.ex          |   1 +
 .../code_intelligence/references_test.exs     |  56 +-
 .../code_intelligence/variable_test.exs       | 489 ++++++++++++++++++
 .../indexer/extractors/variable_test.exs      | 150 +++++-
 .../support/lexical/test/extractor_case.ex    |  12 +-
 .../provider/handlers/find_references.ex      |  10 +-
 .../handlers/find_references_test.exs         |   4 +-
 .../lib/lexical/document/range.ex             |   3 +
 12 files changed, 1023 insertions(+), 60 deletions(-)
 create mode 100644 apps/remote_control/lib/lexical/remote_control/code_intelligence/variable.ex
 create mode 100644 apps/remote_control/test/lexical/remote_control/code_intelligence/variable_test.exs

diff --git a/apps/remote_control/lib/lexical/remote_control/code_intelligence/entity.ex b/apps/remote_control/lib/lexical/remote_control/code_intelligence/entity.ex
index 20af38811..92ea0c76e 100644
--- a/apps/remote_control/lib/lexical/remote_control/code_intelligence/entity.ex
+++ b/apps/remote_control/lib/lexical/remote_control/code_intelligence/entity.ex
@@ -16,6 +16,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.Entity do
           | {:call, module(), fun_name :: atom(), arity :: non_neg_integer()}
           | {:type, module(), type_name :: atom(), arity :: non_neg_integer()}
           | {:module_attribute, container_module :: module(), attribut_name :: atom()}
+          | {:variable, variable_name :: atom()}
 
   defguardp is_call(form) when Sourceror.Identifier.is_call(form) and elem(form, 0) != :.
 
@@ -84,11 +85,11 @@ defmodule Lexical.RemoteControl.CodeIntelligence.Entity do
 
         case fetch_module_for_function(analysis, position, maybe_fun, arity) do
           {:ok, module} -> {:ok, {:call, module, maybe_fun, arity}, node_range}
-          _ -> {:error, :not_found}
+          _ -> {:ok, {:variable, List.to_atom(chars)}, node_range}
         end
 
       _ ->
-        {:error, :not_found}
+        {:ok, {:variable, List.to_atom(chars)}, node_range}
     end
   end
 
diff --git a/apps/remote_control/lib/lexical/remote_control/code_intelligence/references.ex b/apps/remote_control/lib/lexical/remote_control/code_intelligence/references.ex
index a0bc3667b..b32a949f2 100644
--- a/apps/remote_control/lib/lexical/remote_control/code_intelligence/references.ex
+++ b/apps/remote_control/lib/lexical/remote_control/code_intelligence/references.ex
@@ -5,6 +5,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.References do
   alias Lexical.Document.Position
   alias Lexical.RemoteControl.Analyzer
   alias Lexical.RemoteControl.CodeIntelligence.Entity
+  alias Lexical.RemoteControl.CodeIntelligence.Variable
   alias Lexical.RemoteControl.Search.Indexer.Entry
   alias Lexical.RemoteControl.Search.Store
   alias Lexical.RemoteControl.Search.Subject
@@ -13,48 +14,59 @@ defmodule Lexical.RemoteControl.CodeIntelligence.References do
 
   def references(%Analysis{} = analysis, %Position{} = position, include_definitions?) do
     with {:ok, resolved, _range} <- Entity.resolve(analysis, position) do
-      resolved = maybe_rewrite_resolution(resolved, analysis, position)
-
-      references =
-        resolved
-        |> maybe_rewrite_resolution(analysis, position)
-        |> find_references(include_definitions?)
-
-      {:ok, references}
+      resolved
+      |> maybe_rewrite_resolution(analysis, position)
+      |> find_references(analysis, position, include_definitions?)
     end
   end
 
-  defp find_references({:module, module}, include_definitions?) do
+  defp find_references({:module, module}, _analysis, _position, include_definitions?) do
     subject = Subject.module(module)
     subtype = subtype(include_definitions?)
 
     query(subject, type: :module, subtype: subtype)
   end
 
-  defp find_references({:struct, struct_module}, include_definitions?) do
+  defp find_references({:struct, struct_module}, _analysis, _position, include_definitions?) do
     subject = Subject.module(struct_module)
     subtype = subtype(include_definitions?)
 
     query(subject, type: :struct, subtype: subtype)
   end
 
-  defp find_references({:call, module, function_name, arity}, include_definitions?) do
+  defp find_references(
+         {:call, module, function_name, arity},
+         _analysis,
+         _position,
+         include_definitions?
+       ) do
     subject = Subject.mfa(module, function_name, arity)
     subtype = subtype(include_definitions?)
 
     query(subject, type: :function, subtype: subtype)
   end
 
-  defp find_references({:module_attribute, module, attribute_name}, include_definitions?) do
+  defp find_references(
+         {:module_attribute, module, attribute_name},
+         _analysis,
+         _position,
+         include_definitions?
+       ) do
     subject = Subject.module_attribute(module, attribute_name)
     subtype = subtype(include_definitions?)
 
     query(subject, type: :module_attribute, subtype: subtype)
   end
 
-  defp find_references(resolved, _include_definitions?) do
+  defp find_references({:variable, var_name}, analysis, position, include_definitions?) do
+    analysis
+    |> Variable.references(position, var_name, include_definitions?)
+    |> Enum.map(&to_location/1)
+  end
+
+  defp find_references(resolved, _, _, _include_definitions?) do
     Logger.info("Not attempting to find references for unhandled type: #{inspect(resolved)}")
-    []
+    :error
   end
 
   def maybe_rewrite_resolution({:call, Kernel, :defstruct, 1}, analysis, position) do
diff --git a/apps/remote_control/lib/lexical/remote_control/code_intelligence/variable.ex b/apps/remote_control/lib/lexical/remote_control/code_intelligence/variable.ex
new file mode 100644
index 000000000..2a17e6851
--- /dev/null
+++ b/apps/remote_control/lib/lexical/remote_control/code_intelligence/variable.ex
@@ -0,0 +1,258 @@
+defmodule Lexical.RemoteControl.CodeIntelligence.Variable do
+  alias Lexical.Ast.Analysis
+  alias Lexical.Document.Position
+  alias Lexical.Document.Range
+  alias Lexical.RemoteControl.Search.Indexer
+  alias Lexical.RemoteControl.Search.Indexer.Entry
+
+  require Logger
+
+  @extractors [Indexer.Extractors.Variable]
+
+  @spec definition(Analysis.t(), Position.t(), atom()) :: {:ok, Entry.t()} | :error
+  def definition(%Analysis{} = analysis, %Position{} = position, variable_name) do
+    with {:ok, block_structure, entries} <- index_variables(analysis),
+         {:ok, %Entry{} = definition_entry} <-
+           do_find_definition(variable_name, block_structure, entries, position) do
+      {:ok, definition_entry}
+    else
+      _ ->
+        :error
+    end
+  end
+
+  @spec references(Analysis.t(), Position.t(), charlist(), boolean()) :: [Range.t()]
+  def references(
+        %Analysis{} = analysis,
+        %Position{} = position,
+        variable_name,
+        include_definitions? \\ false
+      ) do
+    with {:ok, block_structure, entries} <- index_variables(analysis),
+         {:ok, %Entry{} = definition_entry} <-
+           do_find_definition(variable_name, block_structure, entries, position) do
+      references = search_for_references(entries, definition_entry, block_structure)
+
+      entries =
+        if include_definitions? do
+          [definition_entry | references]
+        else
+          references
+        end
+
+      Enum.sort_by(entries, fn %Entry{} = entry ->
+        {entry.range.start.line, entry.range.start.character}
+      end)
+    else
+      _ ->
+        []
+    end
+  end
+
+  defp index_variables(%Analysis{} = analysis) do
+    with {:ok, entries} <- Indexer.Quoted.index(analysis, @extractors),
+         {[block_structure], entries} <- Enum.split_with(entries, &(&1.type == :metadata)) do
+      {:ok, block_structure.subject, entries}
+    end
+  end
+
+  defp do_find_definition(variable_name, block_structure, entries, position) do
+    with {:ok, entry} <- fetch_entry(entries, variable_name, position) do
+      search_for_definition(entries, entry, block_structure)
+    end
+  end
+
+  defp fetch_entry(entries, variable_name, position) do
+    entries
+    |> Enum.find(fn %Entry{} = entry ->
+      entry.subject == variable_name and entry.type == :variable and
+        Range.contains?(entry.range, position)
+    end)
+    |> case do
+      %Entry{} = entry ->
+        {:ok, entry}
+
+      _ ->
+        :error
+    end
+  end
+
+  defp search_for_references(entries, %Entry{} = definition_entry, block_structure) do
+    block_id_to_children = block_id_to_children(block_structure)
+
+    definition_children = Map.get(block_id_to_children, definition_entry.block_id, [])
+
+    definition_start = definition_entry.range.start
+
+    # The algorithm here is to first clean up the entries so they either are definitions or references to a
+    # variable with the given name. We sort them by their occurrence in the file, working backwards on a line, so
+    # definitions earlier in the line shadow definitions later in the line.
+    # Then we start at the definition entry, and then for each entry after that,
+    # if it's a definition, we mark the state as being shadowed, but reset the state if the block
+    # id isn't in the children of the current block id. If we're not in a child of the current block
+    # id, then we're no longer shadowed
+    #
+    # Note, this algorithm doesn't work when we have a block definition whose result rebinds a variable.
+    # For example:
+    # entries = [4, 5, 6]
+    # entries =
+    #  if something() do
+    #    [1 | entries]
+    #  else
+    #    entries
+    # end
+    # Searching for the references to the initial variable won't find anything inside the block, but
+    # searching for the rebound variable will.
+
+    {entries, _, _} =
+      entries
+      |> Enum.filter(fn %Entry{} = entry ->
+        entry_start = entry.range.start
+
+        after_definition? =
+          if entry_start.line == definition_start.line do
+            entry_start.character > definition_entry.range.end.character
+          else
+            entry_start.line > definition_start.line
+          end
+
+        variable_type? = entry.type == :variable
+        correct_subject? = entry.subject == definition_entry.subject
+        child_of_definition_block? = entry.block_id in definition_children
+
+        variable_type? and correct_subject? and child_of_definition_block? and after_definition?
+      end)
+      |> Enum.sort_by(fn %Entry{} = entry ->
+        start = entry.range.start
+        {start.line, -start.character, entry.block_id}
+      end)
+      |> Enum.reduce({[], false, definition_entry.block_id}, fn
+        %Entry{subtype: :definition} = entry, {entries, _, _} ->
+          # we have a definition that's shadowing our definition entry
+          {entries, true, entry.block_id}
+
+        %Entry{subtype: :reference} = entry, {entries, true, current_block_id} ->
+          shadowed? = entry.block_id in Map.get(block_id_to_children, current_block_id, [])
+
+          entries =
+            if shadowed? do
+              entries
+            else
+              [entry | entries]
+            end
+
+          {entries, shadowed?, entry.block_id}
+
+        %Entry{} = entry, {entries, false, _} ->
+          # we're a reference and we're not being shadowed; collect it and move on.
+          {[entry | entries], false, entry.block_id}
+      end)
+
+    entries
+  end
+
+  defp search_for_definition(entries, %Entry{} = entry, block_structure) do
+    block_id_to_parents = collect_parents(block_structure)
+    block_path = Map.get(block_id_to_parents, entry.block_id)
+    entries_by_block_id = entries_by_block_id(entries)
+
+    Enum.reduce_while([entry.block_id | block_path], :error, fn block_id, _ ->
+      block_entries =
+        entries_by_block_id
+        |> Map.get(block_id, [])
+        |> then(fn entries ->
+          # In the current block, reject all entries that come after the entry whose definition
+          # we're searching for. This prevents us from finding definitions who are shadowing
+          # our entry. For example, the definition on the left of the equals in: `param = param + 1`.
+
+          if block_id == entry.block_id do
+            Enum.drop_while(entries, &(&1.id != entry.id))
+          else
+            entries
+          end
+        end)
+
+      case Enum.find(block_entries, &definition_of?(entry, &1)) do
+        %Entry{} = definition ->
+          {:halt, {:ok, definition}}
+
+        nil ->
+          {:cont, :error}
+      end
+    end)
+  end
+
+  defp definition_of?(%Entry{} = needle, %Entry{} = compare) do
+    compare.type == :variable and compare.subtype == :definition and
+      compare.subject == needle.subject
+  end
+
+  defp entries_by_block_id(entries) do
+    entries
+    |> Enum.reduce(%{}, fn %Entry{} = entry, acc ->
+      Map.update(acc, entry.block_id, [entry], &[entry | &1])
+    end)
+    |> Map.new(fn {block_id, entries} ->
+      entries =
+        Enum.sort_by(
+          entries,
+          fn %Entry{} = entry ->
+            {entry.range.start.line, -entry.range.start.character}
+          end,
+          :desc
+        )
+
+      {block_id, entries}
+    end)
+  end
+
+  def block_id_to_parents(hierarchy) do
+    hierarchy
+    |> flatten_hierarchy()
+    |> Enum.reduce(%{}, fn {parent_id, child_id}, acc ->
+      old_parents = [parent_id | Map.get(acc, parent_id, [])]
+      Map.update(acc, child_id, old_parents, &Enum.concat(&1, old_parents))
+    end)
+    |> Map.put(:root, [])
+  end
+
+  def block_id_to_children(hierarchy) do
+    # Note: Parent ids are included in their children list in order to simplify
+    # checks for "is this id in one of its children"
+
+    hierarchy
+    |> flatten_hierarchy()
+    |> Enum.reverse()
+    |> Enum.reduce(%{root: [:root]}, fn {parent_id, child_id}, current_mapping ->
+      current_children = [child_id | Map.get(current_mapping, child_id, [parent_id])]
+
+      current_mapping
+      |> Map.put_new(child_id, [child_id])
+      |> Map.update(parent_id, current_children, &Enum.concat(&1, current_children))
+    end)
+  end
+
+  def flatten_hierarchy(hierarchy) do
+    Enum.flat_map(hierarchy, fn
+      {k, v} when is_map(v) and map_size(v) > 0 ->
+        v
+        |> Map.keys()
+        |> Enum.map(&{k, &1})
+        |> Enum.concat(flatten_hierarchy(v))
+
+      _ ->
+        []
+    end)
+  end
+
+  defp collect_parents(block_structure) do
+    do_collect_parents(block_structure, %{}, [])
+  end
+
+  defp do_collect_parents(hierarchy, parent_map, path) do
+    Enum.reduce(hierarchy, parent_map, fn {block_id, children}, acc ->
+      parent_map = Map.put(acc, block_id, path)
+      do_collect_parents(children, parent_map, [block_id | path])
+    end)
+  end
+end
diff --git a/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex b/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
index 16bdf4ccc..17bb0f626 100644
--- a/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
+++ b/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
@@ -19,11 +19,30 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
     {:ok, entries, List.wrap(body)}
   end
 
+  # with match operator. with {:ok, var} <- something()
+  def extract({:<-, _, [left, right]}, %Reducer{} = reducer) do
+    entries = extract_definitions(left, reducer)
+
+    {:ok, entries, right}
+  end
+
   # Match operator left = right
   def extract({:=, _, [left, right]}, %Reducer{} = reducer) do
     definitions = extract_definitions(left, reducer)
-    references = extract_references(right, reducer)
-    {:ok, definitions ++ references, nil}
+
+    {:ok, definitions, [right]}
+  end
+
+  # String interpolations "#{foo}"
+  def extract(
+        {:<<>>, _, [{:"::", _, [{{:., _, [Kernel, :to_string]}, _, body}, {:binary, _, _}]}]},
+        %Reducer{}
+      ) do
+    {:ok, [], body}
+  end
+
+  def extract({:binary, _, _}, %Reducer{}) do
+    :ignored
   end
 
   # Generic variable reference
@@ -37,6 +56,7 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
   # Pin operator ^pinned_variable
   def extract({:^, _, [reference]}, %Reducer{} = reducer) do
     reference = extract_reference(reference, reducer, get_current_app(reducer))
+
     {:ok, reference, nil}
   end
 
@@ -72,6 +92,16 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
     {reference, nil}
   end
 
+  # unquote(expression)
+  defp extract_definition({:unquote, _, [expr]}, %Reducer{} = reducer, current_app) do
+    reference = extract_reference(expr, reducer, current_app)
+    {reference, nil}
+  end
+
+  defp extract_definition({:binary, _metadata, nil}, _reducer, _current_app) do
+    nil
+  end
+
   defp extract_definition({var_name, _metadata, nil} = ast, reducer, current_app) do
     if used_variable?(var_name) do
       document = reducer.analysis.document
@@ -90,23 +120,6 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
 
   defp extract_definition(_, _reducer, _current_app), do: nil
 
-  defp extract_references(ast, reducer) do
-    current_app = get_current_app(reducer)
-
-    {_ast, entries} =
-      Macro.prewalk(ast, [], fn ast, acc ->
-        case extract_reference(ast, reducer, current_app) do
-          %Entry{} = entry ->
-            {ast, [entry | acc]}
-
-          _ ->
-            {ast, acc}
-        end
-      end)
-
-    Enum.reverse(entries)
-  end
-
   defp extract_reference({var_name, _metadata, nil} = ast, reducer, current_app) do
     if used_variable?(var_name) do
       document = reducer.analysis.document
@@ -123,7 +136,9 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
     end
   end
 
-  defp extract_reference(_, _, _), do: nil
+  defp extract_reference(_, _, _) do
+    nil
+  end
 
   # extracts definitions like e in SomeException ->
   defp extract_in_definitions(ast, %Reducer{} = reducer) do
diff --git a/apps/remote_control/lib/lexical/remote_control/search/indexer/source/reducer.ex b/apps/remote_control/lib/lexical/remote_control/search/indexer/source/reducer.ex
index b4fb8b191..eeb9ac90b 100644
--- a/apps/remote_control/lib/lexical/remote_control/search/indexer/source/reducer.ex
+++ b/apps/remote_control/lib/lexical/remote_control/search/indexer/source/reducer.ex
@@ -87,6 +87,7 @@ defmodule Lexical.RemoteControl.Search.Indexer.Source.Reducer do
       {:expression, position} ->
         reducer
         |> update_position(position)
+        |> maybe_pop_block()
         |> apply_extractors(element)
     end
   end
diff --git a/apps/remote_control/test/lexical/remote_control/code_intelligence/references_test.exs b/apps/remote_control/test/lexical/remote_control/code_intelligence/references_test.exs
index e185f3e29..487e296d5 100644
--- a/apps/remote_control/test/lexical/remote_control/code_intelligence/references_test.exs
+++ b/apps/remote_control/test/lexical/remote_control/code_intelligence/references_test.exs
@@ -61,7 +61,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
         end
       ]
 
-      assert {:ok, [%Location{} = location]} = references(project, "ReferencedModule|", code)
+      assert [%Location{} = location] = references(project, "ReferencedModule|", code)
       assert decorate(code, location.range) =~ ~s[alias «ReferencedModule»]
     end
 
@@ -72,7 +72,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
         end
       ]
 
-      assert {:ok, [%Location{} = location]} = references(project, "ReferencedModule|", code)
+      assert [%Location{} = location] = references(project, "ReferencedModule|", code)
       assert decorate(code, location.range) =~ ~s[@attr «ReferencedModule»]
     end
 
@@ -81,7 +81,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
         some_module = ReferencedModule
       ]
 
-      assert {:ok, [%Location{} = location]} = references(project, "ReferencedModule|", code)
+      assert [%Location{} = location] = references(project, "ReferencedModule|", code)
       assert decorate(code, location.range) =~ ~s[some_module = «ReferencedModule»]
     end
 
@@ -91,7 +91,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
         end
       ]
 
-      assert {:ok, [%Location{} = location]} = references(project, "ReferencedModule|", code)
+      assert [%Location{} = location] = references(project, "ReferencedModule|", code)
       assert decorate(code, location.range) =~ ~s[def some_fn(«ReferencedModule») do]
     end
 
@@ -100,7 +100,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
         %ReferencedModule{} = something_else
       ]
 
-      assert {:ok, [%Location{} = location]} = references(project, "ReferencedModule|", code)
+      assert [%Location{} = location] = references(project, "ReferencedModule|", code)
       assert decorate(code, location.range) =~ ~s[%«ReferencedModule»{} = something_else]
     end
 
@@ -114,7 +114,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
         end
       ]
 
-      assert {:ok, [location_1, location_2]} = references(project, "DefinedModule|", code, true)
+      assert [location_1, location_2] = references(project, "DefinedModule|", code, true)
       assert decorate(code, location_1.range) =~ ~s[defmodule «DefinedModule» do]
       assert decorate(code, location_2.range) =~ ~s[@attr «DefinedModule»]
     end
@@ -128,7 +128,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
       end
       )
 
-      assert {:ok, [location]} = references(project, "%Struct|{}", code, true)
+      assert [location] = references(project, "%Struct|{}", code, true)
       assert decorate(code, location.range) =~ "«defstruct [:field]»"
     end
 
@@ -145,7 +145,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
         defstruc|t [:name, :value]
       end
       )
-      assert {:ok, [location]} = references(project, selector, code)
+      assert [location] = references(project, selector, code)
       assert decorate(code, location.range) =~ "def something(«%Struct{}») do"
     end
 
@@ -156,7 +156,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
       end
       )
 
-      assert {:ok, []} = references(project, "%Struct|{}", code)
+      assert [] = references(project, "%Struct|{}", code)
     end
   end
 
@@ -177,7 +177,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
 
       ]
 
-      assert {:ok, [reference]} = references(project, query, code)
+      assert [reference] = references(project, query, code)
       assert decorate(code, reference.range) =~ "  def fun(«@attr»), do: true"
     end
 
@@ -197,12 +197,46 @@ defmodule Lexical.RemoteControl.CodeIntelligence.ReferencesTest do
 
       ]
 
-      assert {:ok, [definition, reference]} = references(project, query, code, true)
+      assert [definition, reference] = references(project, query, code, true)
       assert decorate(code, definition.range) =~ "«@attr 3»"
       assert decorate(code, reference.range) =~ "  def fun(«@attr»), do: true"
     end
   end
 
+  describe "variable references" do
+    test "are found in a function body", %{project: project} do
+      query = ~S[
+        def my_fun do
+          first| = 4
+          y = first * 2
+          z = y * 3 + first
+        end
+      ]
+
+      {_, code} = pop_cursor(query)
+
+      assert [ref_1, ref_2] = references(project, query, code)
+
+      assert decorate(code, ref_1.range) =~ "  y = «first» * 2"
+      assert decorate(code, ref_2.range) =~ "  z = y * 3 + «first»"
+    end
+
+    test "can include definitions", %{project: project} do
+      query = ~S[
+        def my_fun do
+          first = 4
+          y = first| * 2
+          z = y * 3 + first
+        end
+      ]
+
+      {_, code} = pop_cursor(query)
+
+      assert [definition, _ref_1, _ref_2] = references(project, query, code, true)
+      assert decorate(code, definition.range) =~ "  «first» = 4"
+    end
+  end
+
   defp references(project, referenced, code, include_definitions? \\ false) do
     with {position, referenced} <- pop_cursor(referenced, as: :document),
          {:ok, document} <- project_module(project, code),
diff --git a/apps/remote_control/test/lexical/remote_control/code_intelligence/variable_test.exs b/apps/remote_control/test/lexical/remote_control/code_intelligence/variable_test.exs
new file mode 100644
index 000000000..7a97d871e
--- /dev/null
+++ b/apps/remote_control/test/lexical/remote_control/code_intelligence/variable_test.exs
@@ -0,0 +1,489 @@
+defmodule Lexical.RemoteControl.CodeIntelligence.VariableTest do
+  alias Lexical.Ast
+  alias Lexical.RemoteControl.CodeIntelligence.Variable
+
+  use ExUnit.Case
+
+  import Lexical.Test.CodeSigil
+  import Lexical.Test.CursorSupport
+  import Lexical.Test.RangeSupport
+
+  def find_definition(code) do
+    {position, document} = pop_cursor(code, as: :document)
+    analysis = Ast.analyze(document)
+    {:ok, {:local_or_var, var_name}} = Ast.cursor_context(analysis, position)
+
+    case Variable.definition(analysis, position, List.to_atom(var_name)) do
+      {:ok, entry} -> {:ok, entry.range, document}
+      error -> error
+    end
+  end
+
+  def find_references(code, include_definition? \\ false) do
+    {position, document} = pop_cursor(code, as: :document)
+    analysis = Ast.analyze(document)
+    {:ok, {:local_or_var, var_name}} = Ast.cursor_context(analysis, position)
+
+    ranges =
+      analysis
+      |> Variable.references(position, List.to_atom(var_name), include_definition?)
+      |> Enum.map(& &1.range)
+
+    {:ok, ranges, document}
+  end
+
+  describe "definitions in a single scope" do
+    test "are returned if it is selected" do
+      {:ok, range, doc} =
+        ~q[
+          def foo(param|) do
+            param
+          end
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "def foo(«param») do"
+    end
+
+    test "are found in a parameter" do
+      {:ok, range, doc} =
+        ~q[
+          def foo(param) do
+            param|
+          end
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "def foo(«param») do"
+    end
+
+    test "are found in a parameter list" do
+      {:ok, range, doc} =
+        ~q[
+          def foo(other_param, param) do
+            param|
+          end
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "def foo(other_param, «param») do"
+    end
+
+    test "are found when shadowed" do
+      {:ok, range, doc} =
+        ~q[
+          def foo(param) do
+            param = param + 1
+            param|
+          end
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "«param» = param + 1"
+    end
+
+    test "are found when shadowing a parameter" do
+      {:ok, range, doc} =
+        ~q[
+          def foo(param) do
+            param = param| + 1
+          end
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "def foo(«param») do"
+    end
+
+    test "when there are multiple definitions on one line" do
+      {:ok, range, doc} =
+        ~q[
+            param = 3
+            foo = param = param + 1
+            param|
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "= «param» = param + 1"
+    end
+
+    test "when the definition is in a map key" do
+      {:ok, range, doc} =
+        ~q[
+          %{key: value} = map
+          value|
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "%{key: «value»} = map"
+    end
+  end
+
+  describe "definitions across scopes" do
+    test "works in an if in a function" do
+      {:ok, range, doc} =
+        ~q[
+          def my_fun do
+            foo = 3
+            if something do
+              foo|
+            end
+          end
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "«foo» = 3"
+    end
+
+    test "works for variables defined in a module" do
+      {:ok, range, doc} =
+        ~q[
+          defmodule Parent do
+            x = 3
+            def fun do
+              unquote(x|)
+            end
+          end
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "«x» = 3"
+    end
+
+    test "works for variables defined outside module" do
+      {:ok, range, doc} =
+        ~q[
+          x = 3
+          defmodule Parent do
+            def fun do
+              unquote(x|)
+            end
+          end
+        ]
+        |> find_definition()
+
+      assert decorate(doc, range) =~ "«x» = 3"
+    end
+  end
+
+  describe "references" do
+    test "in a function parameter" do
+      {:ok, [range], doc} =
+        ~q[
+          def something(param|) do
+            param
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, range) =~ "«param»"
+    end
+
+    test "can include definitions" do
+      {:ok, [definition, reference], doc} =
+        ~q[
+          def something(param|) do
+            param
+          end
+        ]
+        |> find_references(true)
+
+      assert decorate(doc, definition) =~ "def something(«param») do"
+      assert decorate(doc, reference) =~ "  «param»"
+    end
+
+    test "can be found via a usage" do
+      {:ok, [first, second, third], doc} =
+        ~q[
+          def something(param) do
+           y = param + 3
+           z = param + 4
+           param| + y + z
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, first) =~ " y = «param» + 3"
+      assert decorate(doc, second) =~ " z = «param» + 4"
+      assert decorate(doc, third) =~ " «param» + y + z"
+    end
+
+    test "are found in a function body" do
+      {:ok, [first, second, third, fourth, fifth], doc} =
+        ~q[
+          def something(param|) do
+            x = param + param + 3
+            y = param + x
+            z = 10 + param
+            x + y + z + param
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, first) =~ "  x = «param» + param + 3"
+      assert decorate(doc, second) =~ "  x = param + «param» + 3"
+      assert decorate(doc, third) =~ "  y = «param» + x"
+      assert decorate(doc, fourth) =~ "  z = 10 + «param»"
+      assert decorate(doc, fifth) =~ " x + y + z + «param»"
+    end
+
+    test "are constrained to their definition function" do
+      {:ok, [range], doc} =
+        ~q[
+          def something(param|) do
+            param
+          end
+
+          def other_fn(param) do
+            param + 1
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, range) =~ "«param»"
+    end
+
+    test "are visible across blocks" do
+      {:ok, [first, second], doc} =
+        ~q[
+          def something(param|) do
+            if something() do
+              param + 1
+            else
+              param + 2
+            end
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, first) =~ "  «param» + 1"
+      assert decorate(doc, second) =~ "  «param» + 2"
+    end
+
+    test "dont leak out of blocks" do
+      {:ok, [range], doc} =
+        ~q[
+          def something(param) do
+
+            if something() do
+              param| = 3
+              param + 1
+            end
+            param + 1
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, range) =~ "«param»"
+    end
+
+    test "are found in the head of a case statement" do
+      {:ok, [range], doc} =
+        ~q[
+          def something(param|) do
+            case param do
+             _ -> :ok
+            end
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, range) =~ "  case «param» do"
+    end
+
+    test "are constrained to a single arm of a case statement" do
+      {:ok, [range], doc} =
+        ~q[
+          def something(param) do
+            case param do
+             param| when is_number(param) -> param + 1
+             param -> 0
+            end
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, range) =~ "  param when is_number(param) -> «param» + 1"
+    end
+
+    test "are found in a module body" do
+      {:ok, [range], doc} =
+        ~q[
+        defmodule Outer do
+          something| = 3
+          def foo(unquote(something)) do
+          end
+        end
+        ]
+        |> find_references()
+
+      assert decorate(doc, range) =~ "def foo(unquote(«something»)) do"
+    end
+
+    test "are found in anonymous function parameters" do
+      {:ok, [first, second], doc} =
+        ~q[
+        def outer do
+          fn param| ->
+            y = param + 1
+            x = param + 2
+            x + y
+          end
+        end
+        ]
+        |> find_references()
+
+      assert decorate(doc, first) =~ "y = «param» + 1"
+      assert decorate(doc, second) =~ "x = «param» + 2"
+    end
+
+    test "are found in a pin operator" do
+      {:ok, [ref], doc} =
+        ~q[
+        def outer(param|) do
+          fn ^param ->
+            nil
+          end
+        end
+        ]
+        |> find_references()
+
+      assert decorate(doc, ref) =~ "fn ^«param» ->"
+    end
+
+    test "are found inside of string interpolation" do
+      {:ok, [ref], doc} =
+        ~S[
+          name| = "Stinky"
+          "#{name} Stinkman"
+        ]
+        |> find_references()
+
+      assert decorate(doc, ref) =~ "\#{«name»} Stinkman"
+    end
+
+    # Note: This test needs to pass before we can implement renaming variables reliably
+    @tag :skip
+    test "works for variables defined outside of an if while being shadowed" do
+      {:ok, [first, second], doc} =
+        ~q{
+          entries| = [1, 2, 3]
+          entries =
+            if something() do
+              [4 | entries]
+            else
+              entries
+            end
+        }
+        |> find_references()
+
+      assert decorate(doc, first) =~ "[4 | «entries»]"
+      assert decorate(doc, second) =~ "«entries»"
+    end
+
+    test "finds variables defined in anonymous function arms" do
+      {:ok, [first, second], doc} =
+        ~q"
+          shadowed? = false
+          fn
+          {:foo, entries|} ->
+            if shadowed? do
+              [1, entries]
+            else
+              entries
+            end
+          {:bar, entries} ->
+            entries
+          end
+        "
+        |> find_references()
+
+      assert decorate(doc, first) =~ "[1, «entries»]"
+      assert decorate(doc, second) =~ "«entries»"
+    end
+  end
+
+  describe "reference shadowing" do
+    test "on a single line" do
+      {:ok, [], _doc} =
+        ~q[
+          def something(param) do
+            other = other = other| = param
+          end
+      ]
+        |> find_references()
+    end
+
+    test "in a function body" do
+      {:ok, [], _doc} =
+        ~q[
+          def something(param|) do
+           param = 3
+           param
+          end
+      ]
+        |> find_references()
+    end
+
+    test "in anonymous function arguments" do
+      {:ok, [], _doc} =
+        ~q[
+          def something(param|) do
+           fn param ->
+             param + 1
+           end
+           :ok
+          end
+        ]
+        |> find_references()
+    end
+
+    test "inside of a block" do
+      {:ok, [range], doc} =
+        ~q[
+          def something do
+           shadow| = 4
+           if true do
+             shadow = shadow + 1
+             shadow
+           end
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, range) == "   shadow = «shadow» + 1"
+    end
+
+    test "exiting a block" do
+      {:ok, [range], doc} =
+        ~q[
+          def something do
+           shadow| = 4
+           if true do
+             shadow = :ok
+             shadow
+           end
+           shadow + 1
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, range) == " «shadow» + 1"
+    end
+
+    test "exiting nested blocks" do
+      {:ok, [range], doc} =
+        ~q[
+          def something(param| = arg) do
+            case arg do
+              param when is_number(n) ->
+                param + 4
+            end
+            param + 5
+          end
+        ]
+        |> find_references()
+
+      assert decorate(doc, range) == "  «param» + 5"
+    end
+  end
+end
diff --git a/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs b/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
index 4d57179bb..12b46e406 100644
--- a/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
+++ b/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
@@ -36,9 +36,9 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
       test "in a plain parameter" do
         {:ok, [param], doc} =
           ~q[
-        #{unquote(def_type)} my_fun(var) do
-        end
-        ]
+          #{unquote(def_type)} my_fun(var) do
+          end
+          ]
           |> index_definitions()
 
         assert_definition(param, :var)
@@ -88,6 +88,20 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
         assert decorate(doc, var_1.range) =~ "#{unquote(def_type)} my_fun(%«my_module»{})"
       end
 
+      test "in a bitstrings" do
+        {:ok, [var], doc} =
+          ~q[
+            #{unquote(def_type)} my_fun(<<foo::binary-size(3)>>) do
+            end
+          ]
+          |> index_definitions()
+
+        assert_definition(var, :foo)
+
+        assert decorate(doc, var.range) =~
+                 "#{unquote(def_type)} my_fun(<<«foo»::binary-size(3)>>) do"
+      end
+
       test "in list elements" do
         {:ok, [var_1, var_2], doc} =
           ~q{
@@ -640,6 +654,71 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
       assert decorate(doc, access_ref.range) =~ "3 = foo[«bar»]"
     end
 
+    test "inside string interpolations" do
+      quoted =
+        quote file: "foo.ex", line: 1 do
+          foo = 3
+          "#{foo}"
+        end
+
+      assert {:ok, [ref], doc} = index_references(quoted)
+
+      assert_reference(ref, :foo)
+      assert decorate(doc, ref.range) =~ ~S["#{«foo»}"]
+    end
+
+    test "inside string interpolations that have a statement" do
+      quoted =
+        quote file: "foo.ex", line: 1 do
+          foo = 3
+          "#{foo + 3}"
+        end
+
+      assert {:ok, [ref], doc} = index_references(quoted)
+
+      assert_reference(ref, :foo)
+      assert decorate(doc, ref.range) =~ ~S["#{«foo» + 3}"]
+    end
+
+    test "inside string interpolations that have a literal prefix" do
+      quoted =
+        quote file: "foo.ex", line: 1 do
+          foo = 3
+          "prefix #{foo}"
+        end
+
+      assert {:ok, [ref], doc} = index_references(quoted)
+
+      assert_reference(ref, :foo)
+      assert decorate(doc, ref.range) =~ ~S["prefix #{«foo»}"]
+    end
+
+    test "inside string interpolations that have a literal suffix" do
+      quoted =
+        quote file: "foo.ex", line: 1 do
+          foo = 3
+          "#{foo} suffix"
+        end
+
+      assert {:ok, [ref], doc} = index_references(quoted)
+
+      assert_reference(ref, :foo)
+      assert decorate(doc, ref.range) =~ ~S["#{«foo»} suffix"]
+    end
+
+    test "inside string interpolations that have a literal prefix and suffix" do
+      quoted =
+        quote file: "foo.ex", line: 1 do
+          foo = 3
+          "prefix #{foo} suffix"
+        end
+
+      assert {:ok, [ref], doc} = index_references(quoted)
+
+      assert_reference(ref, :foo)
+      assert decorate(doc, ref.range) =~ ~S["prefix #{«foo»} suffix"]
+    end
+
     test "when inside a rescue block in a try" do
       {:ok, [ref], doc} =
         ~q[
@@ -719,6 +798,18 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
       assert decorate(doc, ref.range) =~ "  «var»"
     end
 
+    test "when unquote is used in a function definition" do
+      {:ok, [ref], doc} =
+        ~q[
+        def my_fun(unquote(other_var)) do
+        end
+        ]
+        |> index_references()
+
+      assert_reference(ref, :other_var)
+      assert decorate(doc, ref.range) =~ "def my_fun(unquote(«other_var»)) do"
+    end
+
     test "unless it begins with underscore" do
       assert {:ok, [], _} = index_references("_")
       assert {:ok, [], _} = index_references("_unused")
@@ -728,6 +819,26 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
   end
 
   describe "variable and references are extracted" do
+    test "in a multiple match" do
+      {:ok, [foo_def, param_def, bar_def, other_ref], doc} =
+        ~q[
+          foo = param = bar = other
+        ]
+        |> index()
+
+      assert_definition(foo_def, :foo)
+      assert decorate(doc, foo_def.range) =~ "«foo» = param = bar = other"
+
+      assert_definition(param_def, :param)
+      assert decorate(doc, param_def.range) =~ "foo = «param» = bar = other"
+
+      assert_definition(bar_def, :bar)
+      assert decorate(doc, bar_def.range) =~ "foo = param = «bar» = other"
+
+      assert_reference(other_ref, :other)
+      assert decorate(doc, other_ref.range) =~ "foo = param = bar = «other»"
+    end
+
     test "in an anoymous function" do
       {:ok, [pin_param, var_param, first_def, pin_pin, var_ref, first_ref], doc} =
         ~q{
@@ -756,5 +867,38 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
       assert_reference(first_ref, :first)
       assert decorate(doc, first_ref.range) =~ "  «first»"
     end
+
+    test "in the match arms of a with" do
+      {:ok, [var_def, var_2_def, var_ref], doc} =
+        ~q[
+          with {:ok, var} <- something(),
+               {:ok, var_2} <- something_else(var) do
+            :bad
+          end
+        ]
+        |> index()
+
+      assert_definition(var_def, :var)
+      assert decorate(doc, var_def.range) =~ "{:ok, «var»} <- something(),"
+
+      assert_definition(var_2_def, :var_2)
+      assert decorate(doc, var_2_def.range) =~ "     {:ok, «var_2»} <- something_else(var) do"
+
+      assert_reference(var_ref, :var)
+      assert decorate(doc, var_ref.range) =~ "     {:ok, var_2} <- something_else(«var») do"
+    end
+
+    test "in the body of a with" do
+      {:ok, [_var_def, var_ref], doc} =
+        ~q[
+          with {:ok, var} <- something() do
+            var + 1
+          end
+        ]
+        |> index()
+
+      assert_reference(var_ref, :var)
+      assert decorate(doc, var_ref.range) =~ "  «var» + 1"
+    end
   end
 end
diff --git a/apps/remote_control/test/support/lexical/test/extractor_case.ex b/apps/remote_control/test/support/lexical/test/extractor_case.ex
index 188207118..865a6a69c 100644
--- a/apps/remote_control/test/support/lexical/test/extractor_case.ex
+++ b/apps/remote_control/test/support/lexical/test/extractor_case.ex
@@ -1,6 +1,7 @@
 defmodule Lexical.Test.ExtractorCase do
   alias Lexical.Document
   alias Lexical.RemoteControl.Search.Indexer
+
   use ExUnit.CaseTemplate
   import Lexical.Test.CodeSigil
 
@@ -12,11 +13,13 @@ defmodule Lexical.Test.ExtractorCase do
     end
   end
 
-  def do_index(source, filter, extractors \\ nil) do
+  def do_index(source, filter, extractors \\ nil)
+
+  def do_index(source, filter, extractors) when is_binary(source) do
     path = "/foo/bar/baz.ex"
     doc = Document.new("file:///#{path}", source, 1)
 
-    case Indexer.Source.index("/foo/bar/baz.ex", source, extractors) do
+    case Indexer.Source.index(path, source, extractors) do
       {:ok, indexed_items} ->
         indexed_items = Enum.filter(indexed_items, filter)
         {:ok, indexed_items, doc}
@@ -26,6 +29,11 @@ defmodule Lexical.Test.ExtractorCase do
     end
   end
 
+  def do_index(quoted_source, filter, extractors) do
+    source_string = Macro.to_string(quoted_source)
+    do_index(source_string, filter, extractors)
+  end
+
   def in_a_module(code, module_name \\ "Parent") do
     ~q[
     defmodule #{module_name} do
diff --git a/apps/server/lib/lexical/server/provider/handlers/find_references.ex b/apps/server/lib/lexical/server/provider/handlers/find_references.ex
index c7827fce5..54f4ced36 100644
--- a/apps/server/lib/lexical/server/provider/handlers/find_references.ex
+++ b/apps/server/lib/lexical/server/provider/handlers/find_references.ex
@@ -12,12 +12,10 @@ defmodule Lexical.Server.Provider.Handlers.FindReferences do
     include_declaration? = !!request.context.include_declaration
 
     locations =
-      with {:ok, _document, %Ast.Analysis{} = analysis} <-
-             Document.Store.fetch(request.document.uri, :analysis),
-           {:ok, locations} <-
-             Api.references(env.project, analysis, request.position, include_declaration?) do
-        locations
-      else
+      case Document.Store.fetch(request.document.uri, :analysis) do
+        {:ok, _document, %Ast.Analysis{} = analysis} ->
+          Api.references(env.project, analysis, request.position, include_declaration?)
+
         _ ->
           nil
       end
diff --git a/apps/server/test/lexical/server/provider/handlers/find_references_test.exs b/apps/server/test/lexical/server/provider/handlers/find_references_test.exs
index 1e64def6f..2e7a5f869 100644
--- a/apps/server/test/lexical/server/provider/handlers/find_references_test.exs
+++ b/apps/server/test/lexical/server/provider/handlers/find_references_test.exs
@@ -62,7 +62,7 @@ defmodule Lexical.Server.Provider.Handlers.FindReferencesTest do
           )
         ]
 
-        {:ok, locations}
+        locations
       end)
 
       {:ok, request} = build_request(uri, 5, 6)
@@ -73,7 +73,7 @@ defmodule Lexical.Server.Provider.Handlers.FindReferencesTest do
     end
 
     test "returns nothing if the entity can't resolve it", %{project: project, uri: uri} do
-      patch(RemoteControl.Api, :references, {:error, :not_found})
+      patch(RemoteControl.Api, :references, nil)
 
       {:ok, request} = build_request(uri, 1, 5)
 
diff --git a/projects/lexical_shared/lib/lexical/document/range.ex b/projects/lexical_shared/lib/lexical/document/range.ex
index 7661ae019..17d65813c 100644
--- a/projects/lexical_shared/lib/lexical/document/range.ex
+++ b/projects/lexical_shared/lib/lexical/document/range.ex
@@ -41,6 +41,9 @@ defmodule Lexical.Document.Range do
     %__MODULE__{start: start_pos, end: end_pos} = range
 
     cond do
+      position.line == start_pos.line and position.line == end_pos.line ->
+        position.character >= start_pos.character and position.character <= end_pos.character
+
       position.line == start_pos.line ->
         position.character >= start_pos.character
 

From 643d6608f5a86328aa58b55063b35a82ab216324 Mon Sep 17 00:00:00 2001
From: Steve Cohen <scohen@scohen.org>
Date: Mon, 18 Mar 2024 13:36:50 -0700
Subject: [PATCH 2/4] handled variable definitions for test cases

---
 .../search/indexer/extractors/variable.ex        | 14 ++++++++++++++
 .../search/indexer/extractors/variable_test.exs  | 16 ++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex b/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
index 17bb0f626..42fddc882 100644
--- a/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
+++ b/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
@@ -41,6 +41,20 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
     {:ok, [], body}
   end
 
+  # Test declarations
+  def extract(
+        {:test, _metadata,
+         [
+           {:__block__, [delimiter: "\"", line: 3, column: 8], ["my test"]},
+           args,
+           body
+         ]},
+        %Reducer{} = reducer
+      ) do
+    entries = extract_definitions(args, reducer)
+    {:ok, entries, body}
+  end
+
   def extract({:binary, _, _}, %Reducer{}) do
     :ignored
   end
diff --git a/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs b/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
index 12b46e406..efd2721eb 100644
--- a/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
+++ b/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
@@ -558,6 +558,22 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
       assert decorate(doc, tuple_second.range) =~
                "%struct_module{key: [list_elem, {tuple_first, «tuple_second»}]} = whatever"
     end
+
+    test "from test arguments" do
+      {:ok, [test_def], doc} =
+        ~q[
+        defmodule TestCase do
+          use ExUnit.Case
+          test "my test", %{var: var} do
+            var
+          end
+        end
+        ]
+        |> index_definitions()
+
+      assert_definition(test_def, :var)
+      assert decorate(doc, test_def.range) =~ "%{var: «var»} do"
+    end
   end
 
   describe "variable references are extracted" do

From 0a8b1d2ecf2dfddb5f0392351d68941b3daf9771 Mon Sep 17 00:00:00 2001
From: Steve Cohen <scohen@scohen.org>
Date: Tue, 19 Mar 2024 07:56:45 -0700
Subject: [PATCH 3/4] Added tests for comprehensions and fixed a bug with the
 left stab operator

---
 .../search/indexer/extractors/variable.ex     |  4 +-
 .../indexer/extractors/variable_test.exs      | 58 +++++++++++++++++++
 2 files changed, 60 insertions(+), 2 deletions(-)

diff --git a/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex b/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
index 42fddc882..f019742d8 100644
--- a/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
+++ b/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
@@ -23,14 +23,14 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
   def extract({:<-, _, [left, right]}, %Reducer{} = reducer) do
     entries = extract_definitions(left, reducer)
 
-    {:ok, entries, right}
+    {:ok, entries, List.wrap(right)}
   end
 
   # Match operator left = right
   def extract({:=, _, [left, right]}, %Reducer{} = reducer) do
     definitions = extract_definitions(left, reducer)
 
-    {:ok, definitions, [right]}
+    {:ok, definitions, List.wrap(right)}
   end
 
   # String interpolations "#{foo}"
diff --git a/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs b/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
index efd2721eb..1950d257c 100644
--- a/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
+++ b/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
@@ -472,6 +472,29 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
       assert decorate(doc, value.range) =~ "else «var» ->"
     end
 
+    test "from comprehensions" do
+      {:ok, [var, thing, field_1, field_2], doc} =
+        ~q[
+          for var <- things,
+              {:ok, thing} = var,
+              {:record, field_1, field_2} <- thing do
+          end
+        ]
+        |> index_definitions()
+
+      assert_definition(var, :var)
+      assert decorate(doc, var.range) =~ "for «var» <- things,"
+
+      assert_definition(thing, :thing)
+      assert decorate(doc, thing.range) =~ "{:ok, «thing»} = var,"
+
+      assert_definition(field_1, :field_1)
+      assert decorate(doc, field_1.range) =~ "{:record, «field_1», field_2} <- thing do"
+
+      assert_definition(field_2, :field_2)
+      assert decorate(doc, field_2.range) =~ "{:record, field_1, «field_2»} <- thing do"
+    end
+
     test "in an else block in a try" do
       {:ok, [value], doc} =
         ~q[
@@ -916,5 +939,40 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
       assert_reference(var_ref, :var)
       assert decorate(doc, var_ref.range) =~ "  «var» + 1"
     end
+
+    test "in a comprehension" do
+      {:ok, extracted, doc} =
+        ~q[
+          for {:ok, var} <- list,
+              {:record, elem_1, elem_2} <- var do
+            {:ok, elem_1 + elem_2}
+          end
+        ]
+        |> index()
+
+      assert [var_def, list_ref, elem_1_def, elem_2_def, var_ref, elem_1_ref, elem_2_ref] =
+               extracted
+
+      assert_definition(var_def, :var)
+      assert decorate(doc, var_def.range) =~ "for {:ok, «var»} <- list,"
+
+      assert_reference(list_ref, :list)
+      assert decorate(doc, list_ref.range) =~ "for {:ok, var} <- «list»,"
+
+      assert_definition(elem_1_def, :elem_1)
+      assert decorate(doc, elem_1_def.range) =~ "{:record, «elem_1», elem_2} <- var do"
+
+      assert_definition(elem_2_def, :elem_2)
+      assert decorate(doc, elem_2_def.range) =~ "{:record, elem_1, «elem_2»} <- var do"
+
+      assert_reference(var_ref, :var)
+      assert decorate(doc, var_ref.range) =~ "{:record, elem_1, elem_2} <- «var» do"
+
+      assert_reference(elem_1_ref, :elem_1)
+      assert decorate(doc, elem_1_ref.range) =~ "  {:ok, «elem_1» + elem_2}"
+
+      assert_reference(elem_2_ref, :elem_2)
+      assert decorate(doc, elem_2_ref.range) =~ "  {:ok, elem_1 + «elem_2»}"
+    end
   end
 end

From 14310cce9f33f4d5d08ae7fbfd98855a0c5ea26a Mon Sep 17 00:00:00 2001
From: Steve Cohen <scohen@scohen.org>
Date: Tue, 19 Mar 2024 10:12:17 -0700
Subject: [PATCH 4/4] Handled guard clauses

---
 .../search/indexer/extractors/variable.ex     | 42 +++++++++++++++++++
 .../code_intelligence/variable_test.exs       |  5 ++-
 .../indexer/extractors/variable_test.exs      | 37 ++++++++++++++++
 3 files changed, 82 insertions(+), 2 deletions(-)

diff --git a/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex b/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
index f019742d8..010a16375 100644
--- a/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
+++ b/apps/remote_control/lib/lexical/remote_control/search/indexer/extractors/variable.ex
@@ -5,6 +5,16 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
   alias Lexical.RemoteControl.Search.Indexer.Source.Reducer
 
   @defs [:def, :defmacro, :defp, :defmacrop]
+
+  def extract(
+        {def, _, [{:when, _, [{_fn_name, _, params} | when_args]}, body]},
+        %Reducer{} = reducer
+      )
+      when def in @defs do
+    entries = extract_definitions(params, reducer) ++ extract_references(when_args, reducer)
+    {:ok, entries, body}
+  end
+
   def extract({def, _, [{_fn_name, _, params}, body]}, %Reducer{} = reducer)
       when def in @defs do
     entries = extract_definitions(params, reducer)
@@ -90,6 +100,9 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
           {%Entry{} = entry, ast} ->
             {ast, [entry | acc]}
 
+          {entries, ast} when is_list(entries) ->
+            {ast, entries ++ acc}
+
           _ ->
             {ast, acc}
         end
@@ -112,6 +125,18 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
     {reference, nil}
   end
 
+  # when clauses actually contain parameters and references
+  defp extract_definition({:when, _, when_args}, %Reducer{} = reducer, _current_app) do
+    {definitions, references} =
+      Enum.split_with(when_args, fn {_, _, context} -> is_atom(context) end)
+
+    definitions = extract_definitions(definitions, reducer)
+    references = extract_references(references, reducer)
+
+    {Enum.reverse(definitions ++ references), nil}
+  end
+
+  # This is an effect of string interpolation
   defp extract_definition({:binary, _metadata, nil}, _reducer, _current_app) do
     nil
   end
@@ -134,6 +159,23 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.Variable do
 
   defp extract_definition(_, _reducer, _current_app), do: nil
 
+  defp extract_references(ast, reducer) do
+    current_app = get_current_app(reducer)
+
+    {_ast, entries} =
+      Macro.prewalk(ast, [], fn ast, acc ->
+        case extract_reference(ast, reducer, current_app) do
+          %Entry{} = entry ->
+            {ast, [entry | acc]}
+
+          _ ->
+            {ast, acc}
+        end
+      end)
+
+    Enum.reverse(entries)
+  end
+
   defp extract_reference({var_name, _metadata, nil} = ast, reducer, current_app) do
     if used_variable?(var_name) do
       document = reducer.analysis.document
diff --git a/apps/remote_control/test/lexical/remote_control/code_intelligence/variable_test.exs b/apps/remote_control/test/lexical/remote_control/code_intelligence/variable_test.exs
index 7a97d871e..4ede9e690 100644
--- a/apps/remote_control/test/lexical/remote_control/code_intelligence/variable_test.exs
+++ b/apps/remote_control/test/lexical/remote_control/code_intelligence/variable_test.exs
@@ -291,7 +291,7 @@ defmodule Lexical.RemoteControl.CodeIntelligence.VariableTest do
     end
 
     test "are constrained to a single arm of a case statement" do
-      {:ok, [range], doc} =
+      {:ok, [guard_range, usage_range], doc} =
         ~q[
           def something(param) do
             case param do
@@ -302,7 +302,8 @@ defmodule Lexical.RemoteControl.CodeIntelligence.VariableTest do
         ]
         |> find_references()
 
-      assert decorate(doc, range) =~ "  param when is_number(param) -> «param» + 1"
+      assert decorate(doc, guard_range) =~ "  param when is_number(«param») -> param + 1"
+      assert decorate(doc, usage_range) =~ "  param when is_number(param) -> «param» + 1"
     end
 
     test "are found in a module body" do
diff --git a/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs b/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
index 1950d257c..c5053d1ff 100644
--- a/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
+++ b/apps/remote_control/test/lexical/remote_control/search/indexer/extractors/variable_test.exs
@@ -974,5 +974,42 @@ defmodule Lexical.RemoteControl.Search.Indexer.Extractors.VariableTest do
       assert_reference(elem_2_ref, :elem_2)
       assert decorate(doc, elem_2_ref.range) =~ "  {:ok, elem_1 + «elem_2»}"
     end
+
+    test "in guards in def functions" do
+      {:ok, [param_def, param_2_def, param_ref], doc} =
+        ~q[
+          def something(param, param_2) when param > 1 do
+          end
+        ]
+        |> index()
+
+      assert_definition(param_def, :param)
+      assert decorate(doc, param_def.range) =~ "def something(«param», param_2) when param > 1 do"
+
+      assert_definition(param_2_def, :param_2)
+
+      assert decorate(doc, param_2_def.range) =~
+               "def something(param, «param_2») when param > 1 do"
+
+      assert_reference(param_ref, :param)
+      assert decorate(doc, param_ref.range) =~ "def something(param, param_2) when «param» > 1 do"
+    end
+
+    test "in guards in anonymous functions" do
+      {:ok, [param_def, param_2_def, param_ref], doc} =
+        ~q[
+         fn param, param_2 when param > 1 -> :ok end
+        ]
+        |> index()
+
+      assert_definition(param_def, :param)
+      assert decorate(doc, param_def.range) =~ "fn «param», param_2 when param > 1 -> :ok end"
+
+      assert_definition(param_2_def, :param_2)
+      assert decorate(doc, param_2_def.range) =~ "fn param, «param_2» when param > 1 -> :ok end"
+
+      assert_reference(param_ref, :param)
+      assert decorate(doc, param_ref.range) =~ "fn param, param_2 when «param» > 1 -> :ok end"
+    end
   end
 end