Support local variable on highlight

lib/rucoa/handlers/text_document_document_highlight_handler.rb

@@ -116,6 +116,8 @@ def call
               IfMapper.call(@node)
             when Nodes::IvarNode, Nodes::IvasgnNode
               InstanceVariableMapper.call(@node)
+            when Nodes::ArgNode, Nodes::LvarNode, Nodes::LvasgnNode
+              LocalVariableMapper.call(@node)
             when Nodes::SendNode
               SendMapper.call(@node)
             when Nodes::UntilNode, Nodes::WhileNode
@@ -356,6 +358,101 @@ def nodes
           end
         end
 
+        class LocalVariableMapper < Base
+          # @return [Array]
+          def call
+            return [] unless nodes.any?
+
+            nodes.map do |node|
+              case node
+              when Nodes::LvarNode
+                Highlights::ReadHighlight
+              when Nodes::ArgNode, Nodes::LvasgnNode
+                Highlights::WriteHighlight
+              end.new(parser_range: node.location.name)
+            end
+          end
+
+          private
+
+          # @return [Rucoa::Nodes::ArgNode, Rucoa::Nodes::LvasgnNode, nil]
+          def assignment_node
+            @assignment_node ||=
+              case @node
+              when Nodes::ArgNode, Nodes::LvasgnNode
+                @node
+              when Nodes::LvarNode
+                scope_boundary_node_types = ::Set[
+                  :class,
+                  :def,
+                  :defs,
+                  :module,
+                ]
+                (
+                  [@node] + @node.ancestors.take_while do |node|
+                    !scope_boundary_node_types.include?(node.type)
+                  end
+                ).find do |node|
+                  case node
+                  when Nodes::ArgNode, Nodes::LvasgnNode
+                    node.name == @node.name
+                  when Nodes::BlockNode, Nodes::DefNode
+                    target = node.arguments.find do |argument|
+                      argument.name == @node.name
+                    end
+                    break target if target
+                  else
+                    target = node.previous_siblings.find do |sibling|
+                      case sibling
+                      when Nodes::LvasgnNode
+                        sibling.name == @node.name
+                      end
+                    end
+                    break target if target
+                  end
+                end || @node.each_ancestor(:def, :defs).first&.then do |node|
+                  break node.arguments.find do |argument|
+                    argument.name == @node.name
+                  end
+                end
+              end
+          end
+
+          # @return [Enumerable<Rucoa::Nodes::ArgNode, Rucoa::Nodes::LvarNode, Rucoa::Nodes::LvasgnNode>]
+          def nodes
+            [
+              assignment_node,
+              *reference_nodes
+            ].compact
+          end
+
+          # @todo Support shadowing.
+          # @return [Enumerable<Rucoa::Nodes::LvarNode>]
+          def reference_nodes
+            return [] unless assignment_node
+
+            case assignment_node
+            when Nodes::ArgNode
+              assignment_node.each_ancestor(:block, :def, :defs).first.body_children
+            when Nodes::LvasgnNode
+              [
+                assignment_node,
+                *assignment_node.next_siblings
+              ]
+            end.flat_map do |node|
+              [
+                node,
+                *node.descendants
+              ]
+            end.select do |node|
+              case node
+              when Nodes::LvarNode
+                node.name == @node.name
+              end
+            end
+          end
+        end
+
         class ModuleMapper < Base
           # @return [Array]
           def call

lib/rucoa/nodes.rb

@@ -2,6 +2,8 @@
 
 module Rucoa
   module Nodes
+    autoload :ArgNode, 'rucoa/nodes/arg_node'
+    autoload :ArgsNode, 'rucoa/nodes/args_node'
     autoload :Base, 'rucoa/nodes/base'
     autoload :BeginNode, 'rucoa/nodes/begin_node'
     autoload :BlockNode, 'rucoa/nodes/block_node'
@@ -21,6 +23,7 @@ module Nodes
     autoload :IvarNode, 'rucoa/nodes/ivar_node'
     autoload :IvasgnNode, 'rucoa/nodes/ivasgn_node'
     autoload :LvarNode, 'rucoa/nodes/lvar_node'
+    autoload :LvasgnNode, 'rucoa/nodes/lvasgn_node'
     autoload :ModuleNode, 'rucoa/nodes/module_node'
     autoload :ResbodyNode, 'rucoa/nodes/resbody_node'
     autoload :RescueNode, 'rucoa/nodes/rescue_node'

lib/rucoa/nodes/arg_node.rb

@@ -0,0 +1,26 @@
+# frozen_string_literal: true
+
+module Rucoa
+  module Nodes
+    class ArgNode < Base
+      # @return [String]
+      # @example returns variable name
+      #   node = Rucoa::Source.new(
+      #     content: <<~RUBY,
+      #       def foo(bar)
+      #       end
+      #     RUBY
+      #     uri: 'file:///path/to/example.rb'
+      #   ).node_at(
+      #     Rucoa::Position.new(
+      #       column: 8,
+      #       line: 1
+      #     )
+      #   )
+      #   expect(node.name).to eq('bar')
+      def name
+        children[0].to_s
+      end
+    end
+  end
+end

lib/rucoa/nodes/args_node.rb

@@ -0,0 +1,14 @@
+# frozen_string_literal: true
+
+module Rucoa
+  module Nodes
+    class ArgsNode < Base
+      include ::Enumerable
+
+      # @note For `Enumerable`.
+      def each(&block)
+        children.each(&block)
+      end
+    end
+  end
+end

lib/rucoa/nodes/base.rb

@@ -126,6 +126,32 @@ def namespace
         module_nesting.first || 'Object'
       end
 
+      # @return [Array<Rucoa::Nodes::Base>]
+      # @example returns next siblings
+      #   node = Rucoa::Source.new(
+      #     content: <<~RUBY,
+      #       def foo
+      #         a
+      #         b
+      #         c
+      #         d
+      #         e
+      #       end
+      #     RUBY
+      #     uri: 'file:///path/to/example.rb'
+      #   ).node_at(
+      #     Rucoa::Position.new(
+      #       column: 2,
+      #       line: 4
+      #     )
+      #   )
+      #   expect(node.next_siblings.map(&:name)).to eq(%w[d e])
+      def next_siblings
+        return [] unless parent
+
+        parent.children[(sibling_index + 1)..]
+      end
+
       # @return [Rucoa::Nodes::Base, nil]
       def parent
         @mutable_attributes[:parent]
@@ -136,6 +162,32 @@ def parent=(node)
         @mutable_attributes[:parent] = node
       end
 
+      # @return [Array<Rucoa::Nodes::Base>]
+      # @example returns previous siblings
+      #   node = Rucoa::Source.new(
+      #     content: <<~RUBY,
+      #       def foo
+      #         a
+      #         b
+      #         c
+      #         d
+      #         e
+      #       end
+      #     RUBY
+      #     uri: 'file:///path/to/example.rb'
+      #   ).node_at(
+      #     Rucoa::Position.new(
+      #       column: 2,
+      #       line: 4
+      #     )
+      #   )
+      #   expect(node.previous_siblings.map(&:name)).to eq(%w[a b])
+      def previous_siblings
+        return [] unless parent
+
+        parent.children[0...sibling_index]
+      end
+
       # @note Override.
       #   Some nodes change their type depending on the context.
       #   For example, `const` node can be `casgn` node.
@@ -170,6 +222,13 @@ def visit_descendants(
 
       private
 
+      # @return [Integer, nil]
+      def sibling_index
+        parent&.children&.index do |child|
+          child.eql?(self)
+        end
+      end
+
       # Visits all ancestors.
       # @param types [Array<Symbol>]
       # @return [void]

lib/rucoa/nodes/block_node.rb

@@ -6,6 +6,20 @@ class BlockNode < Base
       include NodeConcerns::Body
       include NodeConcerns::Rescue
 
+      # @return [Array<Rucoa::Nodes::ArgNode>]
+      # @example returns arguments
+      #   node = Rucoa::Source.new(
+      #     content: <<~RUBY,
+      #       foo do |bar, baz|
+      #       end
+      #     RUBY
+      #     uri: 'file:///path/to/example.rb'
+      #   ).root_node
+      #   expect(node.arguments.map(&:name)).to eq(%w[bar baz])
+      def arguments
+        children[-2]
+      end
+
       # @return [Rucoa::Nodes::SendNode]
       # @example returns send node
       #   node = Rucoa::Source.new(

lib/rucoa/nodes/def_node.rb

@@ -6,6 +6,20 @@ class DefNode < Base
       include NodeConcerns::Body
       include NodeConcerns::Rescue
 
+      # @return [Array<Rucoa::Nodes::ArgNode>]
+      # @example returns arguments
+      #   node = Rucoa::Source.new(
+      #     content: <<~RUBY,
+      #       def foo(bar, baz)
+      #       end
+      #     RUBY
+      #     uri: 'file:///path/to/example.rb'
+      #   ).root_node
+      #   expect(node.arguments.map(&:name)).to eq(%w[bar baz])
+      def arguments
+        children[-2]
+      end
+
       # @return [String]
       def method_marker
         if singleton?

lib/rucoa/nodes/lvasgn_node.rb

@@ -0,0 +1,12 @@
+# frozen_string_literal: true
+
+module Rucoa
+  module Nodes
+    class LvasgnNode < Base
+      # @return [String]
+      def name
+        children[0].to_s
+      end
+    end
+  end
+end

lib/rucoa/parser_builder.rb

@@ -5,6 +5,8 @@
 module Rucoa
   class ParserBuilder < ::Parser::Builders::Default
     NODE_CLASS_BY_TYPE = {
+      arg: Nodes::ArgNode,
+      args: Nodes::ArgsNode,
       begin: Nodes::BeginNode,
       block: Nodes::BlockNode,
       case: Nodes::CaseNode,
@@ -26,6 +28,7 @@ class ParserBuilder < ::Parser::Builders::Default
       ivasgn: Nodes::IvasgnNode,
       kwbegin: Nodes::BeginNode,
       lvar: Nodes::LvarNode,
+      lvasgn: Nodes::LvasgnNode,
       module: Nodes::ModuleNode,
       resbody: Nodes::ResbodyNode,
       rescue: Nodes::RescueNode,