Find references for variables

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
 

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

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