Add support for local variable references

lib/ruby_indexer/lib/ruby_indexer/reference_finder.rb

@@ -43,6 +43,17 @@ def initialize(name, owner_ancestors)
       end
     end
 
+    class LocalVariableTarget < Target
+      #: String
+      attr_reader :name
+
+      #: (String name) -> void
+      def initialize(name)
+        super()
+        @name = name
+      end
+    end
+
     class Reference
       #: String
       attr_reader :name
@@ -98,6 +109,13 @@ def initialize(target, index, dispatcher, uri, include_declarations: true)
         :on_instance_variable_operator_write_node_enter,
         :on_instance_variable_or_write_node_enter,
         :on_instance_variable_target_node_enter,
+        :on_local_variable_and_write_node_enter,
+        :on_local_variable_operator_write_node_enter,
+        :on_local_variable_or_write_node_enter,
+        :on_local_variable_read_node_enter,
+        :on_local_variable_target_node_enter,
+        :on_local_variable_write_node_enter,
+        :on_parameters_node_enter,
         :on_call_node_enter,
       )
     end
@@ -280,6 +298,45 @@ def on_instance_variable_target_node_enter(node)
       collect_instance_variable_references(node.name.to_s, node.location, true)
     end
 
+    #: (Prism::LocalVariableAndWriteNode node) -> void
+    def on_local_variable_and_write_node_enter(node)
+      collect_local_variable_references(node.name.to_s, node.name_loc, true)
+    end
+
+    #: (Prism::LocalVariableOperatorWriteNode node) -> void
+    def on_local_variable_operator_write_node_enter(node)
+      collect_local_variable_references(node.name.to_s, node.name_loc, true)
+    end
+
+    #: (Prism::LocalVariableOrWriteNode node) -> void
+    def on_local_variable_or_write_node_enter(node)
+      collect_local_variable_references(node.name.to_s, node.name_loc, true)
+    end
+
+    #: (Prism::LocalVariableReadNode node) -> void
+    def on_local_variable_read_node_enter(node)
+      collect_local_variable_references(node.name.to_s, node.location, false)
+    end
+
+    #: (Prism::LocalVariableTargetNode node) -> void
+    def on_local_variable_target_node_enter(node)
+      collect_local_variable_references(node.name.to_s, node.location, true)
+    end
+
+    #: (Prism::LocalVariableWriteNode node) -> void
+    def on_local_variable_write_node_enter(node)
+      collect_local_variable_references(node.name.to_s, node.name_loc, true)
+    end
+
+    #: (Prism::ParametersNode node) -> void
+    def on_parameters_node_enter(node)
+      node.requireds.each do |parameter|
+        next unless parameter.is_a?(Prism::RequiredParameterNode)
+
+        collect_local_variable_references(parameter.name.to_s, parameter.location, true)
+      end
+    end
+
     #: (Prism::CallNode node) -> void
     def on_call_node_enter(node)
       if @target.is_a?(MethodTarget) && (name = node.name.to_s) == @target.method_name
@@ -331,5 +388,12 @@ def collect_instance_variable_references(name, location, declaration)
         @references << Reference.new(name, location, declaration: declaration)
       end
     end
+
+    #: (String name, Prism::Location location, bool declaration) -> void
+    def collect_local_variable_references(name, location, declaration)
+      return unless @target.is_a?(LocalVariableTarget) && name == @target.name
+
+      @references << Reference.new(name, location, declaration: declaration)
+    end
   end
 end

lib/ruby_indexer/test/reference_finder_test.rb

@@ -363,6 +363,69 @@ class Bar
       assert_equal(8, refs[2].location.start_line)
     end
 
+    def test_finds_local_variable_references
+      refs = find_local_variable_references("foo", <<~RUBY)
+        class Foo
+          def method
+            foo = 1
+            foo &&= 2
+            foo ||= 3
+            foo += 4
+            foo, bar = []
+            foo
+          end
+        end
+      RUBY
+      assert_equal(6, refs.size)
+
+      assert_equal(["foo"], refs.map(&:name).uniq)
+      assert_equal(3, refs[0].location.start_line)
+      assert_equal(4, refs[1].location.start_line)
+      assert_equal(5, refs[2].location.start_line)
+      assert_equal(6, refs[3].location.start_line)
+      assert_equal(7, refs[4].location.start_line)
+      assert_equal(8, refs[5].location.start_line)
+    end
+
+    def test_finds_local_variable_references_in_nested_scopes
+      refs = find_local_variable_references("foo", <<~RUBY)
+        class Foo
+          def method_with_block
+            foo = 1
+
+            block do
+              foo = 2
+            end
+          end
+        end
+      RUBY
+      assert_equal(2, refs.size)
+
+      assert_equal(["foo"], refs.map(&:name).uniq)
+      assert_equal(3, refs[0].location.start_line)
+      assert_equal(6, refs[1].location.start_line)
+    end
+
+    def test_finds_local_variable_references_in_nested_scopes_with_shadowing
+      refs = find_local_variable_references("foo", <<~RUBY)
+        class Foo
+          def method_with_block_and_shadowing
+            foo = 1
+
+            block do |foo, bar|
+              foo = 2
+            end
+          end
+        end
+      RUBY
+      assert_equal(3, refs.size)
+
+      assert_equal(["foo"], refs.map(&:name).uniq)
+      assert_equal(3, refs[0].location.start_line)
+      assert_equal(5, refs[1].location.start_line)
+      assert_equal(6, refs[2].location.start_line)
+    end
+
     private
 
     def find_const_references(const_name, source)
@@ -380,6 +443,11 @@ def find_instance_variable_references(instance_variable_name, owner_ancestors, s
       find_references(target, source)
     end
 
+    def find_local_variable_references(local_variable_name, source)
+      target = ReferenceFinder::LocalVariableTarget.new(local_variable_name)
+      find_references(target, source)
+    end
+
     def find_references(target, source)
       file_path = "/fake.rb"
       uri = URI::Generic.from_path(path: file_path)

lib/ruby_lsp/requests/prepare_rename.rb

@@ -24,7 +24,18 @@ def perform
         node_context = RubyDocument.locate(
           @document.parse_result.value,
           char_position,
-          node_types: [Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode],
+          node_types: [
+            Prism::ConstantReadNode,
+            Prism::ConstantPathNode,
+            Prism::ConstantPathTargetNode,
+            Prism::LocalVariableAndWriteNode,
+            Prism::LocalVariableOperatorWriteNode,
+            Prism::LocalVariableOrWriteNode,
+            Prism::LocalVariableReadNode,
+            Prism::LocalVariableTargetNode,
+            Prism::LocalVariableWriteNode,
+            Prism::RequiredParameterNode,
+          ],
           code_units_cache: @document.code_units_cache,
         )
         target = node_context.node

lib/ruby_lsp/requests/references.rb

@@ -38,6 +38,13 @@ def perform
             Prism::InstanceVariableReadNode,
             Prism::InstanceVariableTargetNode,
             Prism::InstanceVariableWriteNode,
+            Prism::LocalVariableAndWriteNode,
+            Prism::LocalVariableOperatorWriteNode,
+            Prism::LocalVariableOrWriteNode,
+            Prism::LocalVariableReadNode,
+            Prism::LocalVariableTargetNode,
+            Prism::LocalVariableWriteNode,
+            Prism::RequiredParameterNode,
             Prism::CallNode,
             Prism::DefNode,
           ],
@@ -60,28 +67,34 @@ def perform
         reference_target = create_reference_target(target, node_context)
         return @locations unless reference_target
 
+        # for LocalVariableTarget, we only need to collect references from the current document
+        if reference_target.is_a?(RubyIndexer::ReferenceFinder::LocalVariableTarget)
+          collect_references(reference_target, parent, @document.uri) if parent
+          return @locations
+        end
+
         Dir.glob(File.join(@global_state.workspace_path, "**/*.rb")).each do |path|
           uri = URI::Generic.from_path(path: path)
           # If the document is being managed by the client, then we should use whatever is present in the store instead
           # of reading from disk
           next if @store.key?(uri)
 
           parse_result = Prism.parse_file(path)
-          collect_references(reference_target, parse_result, uri)
+          collect_references(reference_target, parse_result.value, uri)
         rescue Errno::EISDIR, Errno::ENOENT
           # If `path` is a directory, just ignore it and continue. If the file doesn't exist, then we also ignore it.
         end
 
         @store.each do |_uri, document|
-          collect_references(reference_target, document.parse_result, document.uri)
+          collect_references(reference_target, document.parse_result.value, document.uri)
         end
 
         @locations
       end
 
       private
 
-      #: ((Prism::ConstantReadNode | Prism::ConstantPathNode | Prism::ConstantPathTargetNode | Prism::InstanceVariableAndWriteNode | Prism::InstanceVariableOperatorWriteNode | Prism::InstanceVariableOrWriteNode | Prism::InstanceVariableReadNode | Prism::InstanceVariableTargetNode | Prism::InstanceVariableWriteNode | Prism::CallNode | Prism::DefNode) target_node, NodeContext node_context) -> RubyIndexer::ReferenceFinder::Target?
+      #: ((Prism::ConstantReadNode | Prism::ConstantPathNode | Prism::ConstantPathTargetNode | Prism::InstanceVariableAndWriteNode | Prism::InstanceVariableOperatorWriteNode | Prism::InstanceVariableOrWriteNode | Prism::InstanceVariableReadNode | Prism::InstanceVariableTargetNode | Prism::InstanceVariableWriteNode | Prism::LocalVariableAndWriteNode | Prism::LocalVariableOperatorWriteNode | Prism::LocalVariableOrWriteNode | Prism::LocalVariableReadNode | Prism::LocalVariableTargetNode | Prism::LocalVariableWriteNode | Prism::RequiredParameterNode | Prism::CallNode | Prism::DefNode) target_node, NodeContext node_context) -> RubyIndexer::ReferenceFinder::Target?
       def create_reference_target(target_node, node_context)
         case target_node
         when Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode
@@ -106,13 +119,23 @@ def create_reference_target(target_node, node_context)
 
           ancestors = @global_state.index.linearized_ancestors_of(receiver_type.name)
           RubyIndexer::ReferenceFinder::InstanceVariableTarget.new(target_node.name.to_s, ancestors)
+        when
+          Prism::LocalVariableAndWriteNode,
+          Prism::LocalVariableOperatorWriteNode,
+          Prism::LocalVariableOrWriteNode,
+          Prism::LocalVariableReadNode,
+          Prism::LocalVariableTargetNode,
+          Prism::LocalVariableWriteNode
+          RubyIndexer::ReferenceFinder::LocalVariableTarget.new(target_node.name.to_s)
+        when Prism::RequiredParameterNode
+          RubyIndexer::ReferenceFinder::LocalVariableTarget.new(target_node.name.to_s)
         when Prism::CallNode, Prism::DefNode
           RubyIndexer::ReferenceFinder::MethodTarget.new(target_node.name.to_s)
         end
       end
 
-      #: (RubyIndexer::ReferenceFinder::Target target, Prism::ParseResult parse_result, URI::Generic uri) -> void
-      def collect_references(target, parse_result, uri)
+      #: (RubyIndexer::ReferenceFinder::Target target, Prism::Node node, URI::Generic uri) -> void
+      def collect_references(target, node, uri)
         dispatcher = Prism::Dispatcher.new
         finder = RubyIndexer::ReferenceFinder.new(
           target,
@@ -121,7 +144,7 @@ def collect_references(target, parse_result, uri)
           uri,
           include_declarations: @params.dig(:context, :includeDeclaration) || true,
         )
-        dispatcher.visit(parse_result.value)
+        dispatcher.visit(node)
 
         finder.references.each do |reference|
           @locations << Interface::Location.new(

lib/ruby_lsp/requests/rename.rb

@@ -36,7 +36,18 @@ def perform
         node_context = RubyDocument.locate(
           @document.parse_result.value,
           char_position,
-          node_types: [Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode],
+          node_types: [
+            Prism::ConstantReadNode,
+            Prism::ConstantPathNode,
+            Prism::ConstantPathTargetNode,
+            Prism::LocalVariableAndWriteNode,
+            Prism::LocalVariableOperatorWriteNode,
+            Prism::LocalVariableOrWriteNode,
+            Prism::LocalVariableReadNode,
+            Prism::LocalVariableTargetNode,
+            Prism::LocalVariableWriteNode,
+            Prism::RequiredParameterNode,
+          ],
           code_units_cache: @document.code_units_cache,
         )
         target = node_context.node
@@ -51,15 +62,27 @@ def perform
           )
         end
 
-        target = target #: as Prism::ConstantReadNode | Prism::ConstantPathNode | Prism::ConstantPathTargetNode
+        case target
+        when Prism::ConstantReadNode, Prism::ConstantPathNode, Prism::ConstantPathTargetNode
+          perform_constant_rename(target, node_context.nesting)
+        when Prism::LocalVariableAndWriteNode, Prism::LocalVariableOperatorWriteNode, Prism::LocalVariableOrWriteNode,
+             Prism::LocalVariableReadNode, Prism::LocalVariableTargetNode, Prism::LocalVariableWriteNode,
+             Prism::RequiredParameterNode
+          perform_local_variable_rename(target, parent, node_context.nesting) if parent
+        end
+      end
+
+      private
 
+      #: (Prism::ConstantReadNode | Prism::ConstantPathNode | Prism::ConstantPathTargetNode, Array[String]) -> Interface::WorkspaceEdit?
+      def perform_constant_rename(target, nesting)
         name = RubyIndexer::Index.constant_name(target)
         return unless name
 
-        entries = @global_state.index.resolve(name, node_context.nesting)
+        entries = @global_state.index.resolve(name, nesting)
         return unless entries
 
-        if (conflict_entries = @global_state.index.resolve(@new_name, node_context.nesting))
+        if (conflict_entries = @global_state.index.resolve(@new_name, nesting))
           raise InvalidNameError, "The new name is already in use by #{conflict_entries.first&.name}"
         end
 
@@ -87,7 +110,26 @@ def perform
         Interface::WorkspaceEdit.new(document_changes: document_changes)
       end
 
-      private
+      #: (Prism::LocalVariableAndWriteNode | Prism::LocalVariableOperatorWriteNode | Prism::LocalVariableOrWriteNode | Prism::LocalVariableReadNode | Prism::LocalVariableTargetNode | Prism::LocalVariableWriteNode | Prism::RequiredParameterNode, Prism::Node, Array[String]) -> Interface::WorkspaceEdit?
+      def perform_local_variable_rename(target, parent, nesting)
+        name = target.name.to_s
+
+        reference_target = RubyIndexer::ReferenceFinder::LocalVariableTarget.new(name)
+        dispatcher = Prism::Dispatcher.new
+        finder = RubyIndexer::ReferenceFinder.new(
+          reference_target,
+          @global_state.index,
+          dispatcher,
+          @document.uri,
+        )
+        dispatcher.visit(parent)
+
+        edits = finder.references.map do |reference|
+          adjust_reference_for_edit(name, reference)
+        end
+
+        Interface::WorkspaceEdit.new(changes: edits)
+      end
 
       #: (String fully_qualified_name, Array[(Interface::RenameFile | Interface::TextDocumentEdit)] document_changes) -> void
       def collect_file_renames(fully_qualified_name, document_changes)

test/fixtures/local_var_examples.rb

@@ -0,0 +1,26 @@
+def method_with_local_var
+  # on_local_variable_write_node_enter
+  a_local_var = 1
+  # on_local_variable_and_write_node_enter
+  a_local_var &&= 1
+  # on_local_variable_operator_write_node_enter
+  a_local_var += 1
+  # on_local_variable_or_write_node_enter
+  a_local_var ||= 1
+  # on_local_variable_target_node_enter
+  a_local_var, _ = [1, 2]
+  # on_local_variable_read_node_enter
+  a_local_var
+end
+
+def another_method
+  a_local_var = 2
+end
+
+def method_with_block_and_shadowing
+  a_local_var = 1
+
+  block do |a_local_var, bar|
+    a_local_var = 2
+  end
+end