Support local variables for completion and definition requests

lib/core_ext/prism.rb

@@ -0,0 +1,22 @@
+# typed: strict
+# frozen_string_literal: true
+
+module Prism
+  LocalVariableNode = T.type_alias do
+    T.any(
+      Prism::LocalVariableWriteNode,
+      Prism::LocalVariableTargetNode,
+      Prism::LocalVariableAndWriteNode,
+      Prism::LocalVariableOrWriteNode,
+      Prism::LocalVariableReadNode,
+      Prism::LocalVariableOperatorWriteNode,
+      Prism::RequiredParameterNode,
+      Prism::OptionalParameterNode,
+      Prism::RequiredKeywordParameterNode,
+      Prism::OptionalKeywordParameterNode,
+      Prism::BlockLocalVariableNode,
+      Prism::RestParameterNode,
+      Prism::KeywordRestParameterNode,
+    )
+  end
+end

lib/ruby_lsp/document.rb

@@ -173,6 +173,61 @@ def locate(node, char_position, node_types: [])
       [closest, parent, nesting.map { |n| n.constant_path.location.slice }]
     end
 
+    sig do
+      params(position: T::Hash[Symbol, T.untyped]).returns(T::Array[Prism::LocalVariableNode])
+    end
+    def locate_local_variable_nodes(position)
+      locate_local_variables(@parse_result.value, create_scanner.find_char_position(position))
+    end
+
+    sig do
+      params(
+        node: Prism::Node,
+        char_position: Integer,
+      ).returns(T::Array[Prism::LocalVariableNode])
+    end
+    def locate_local_variables(node, char_position)
+      queue = T.let(node.compact_child_nodes, T::Array[Prism::Node])
+      variable_nodes = T.let([], T::Array[Prism::LocalVariableNode])
+
+      until queue.empty?
+        current = T.must(queue.shift)
+        loc = current.location
+
+        # Don't take into account variables after given position
+        break if char_position < loc.start_offset
+
+        # Local variables defined in def, begin, classes, modules and block nodes are visible only in their own scope,
+        # so we can skip it if our position is outside.
+        # Other nodes (like if ... end) don't create its own scope.
+        next if [Prism::DefNode, Prism::BeginNode, Prism::BlockNode, Prism::ClassNode, Prism::ModuleNode]
+          .any? { |t| current.is_a?(t) } && char_position > loc.end_offset
+
+        T.unsafe(queue).unshift(*current.compact_child_nodes)
+
+        case current
+        when Prism::DefNode, Prism::ClassNode # in this nodes we don't have access to variables outside
+          variable_nodes.clear
+        when Prism::LocalVariableWriteNode,
+          Prism::LocalVariableTargetNode,
+          Prism::LocalVariableAndWriteNode,
+          Prism::LocalVariableOrWriteNode,
+          Prism::RequiredParameterNode,
+          Prism::OptionalParameterNode,
+          Prism::RequiredKeywordParameterNode,
+          Prism::OptionalKeywordParameterNode,
+          Prism::BlockLocalVariableNode,
+          Prism::RestParameterNode,
+          Prism::KeywordRestParameterNode
+
+          variable_nodes << current
+        end
+
+      end
+
+      variable_nodes
+    end
+
     sig { returns(T::Boolean) }
     def sorbet_sigil_is_true_or_higher
       parse_result.magic_comments.any? do |comment|

lib/ruby_lsp/internal.rb

@@ -23,6 +23,7 @@
 require "ruby_lsp/base_server"
 require "ruby_indexer/ruby_indexer"
 require "core_ext/uri"
+require "core_ext/prism"
 require "ruby_lsp/utils"
 require "ruby_lsp/parameter_scope"
 require "ruby_lsp/global_state"

lib/ruby_lsp/listeners/completion.rb

@@ -12,16 +12,18 @@ class Completion
           response_builder: ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem],
           global_state: GlobalState,
           nesting: T::Array[String],
+          local_variables: T::Array[Prism::LocalVariableNode],
           typechecker_enabled: T::Boolean,
           dispatcher: Prism::Dispatcher,
           uri: URI::Generic,
         ).void
       end
-      def initialize(response_builder, global_state, nesting, typechecker_enabled, dispatcher, uri) # rubocop:disable Metrics/ParameterLists
+      def initialize(response_builder, global_state, nesting, local_variables, typechecker_enabled, dispatcher, uri) # rubocop:disable Metrics/ParameterLists
         @response_builder = response_builder
         @global_state = global_state
         @index = T.let(global_state.index, RubyIndexer::Index)
         @nesting = nesting
+        @local_variables = local_variables
         @typechecker_enabled = typechecker_enabled
         @uri = uri
 
@@ -166,6 +168,12 @@ def complete_require_relative(node)
 
       sig { params(node: Prism::CallNode, name: String).void }
       def complete_self_receiver_method(node, name)
+        @local_variables.uniq(&:name).each do |variable_node|
+          next unless variable_node.name.to_s.start_with?(node.name.to_s)
+
+          @response_builder << build_local_variable_completion(variable_node, node)
+        end
+
         receiver_entries = @index[@nesting.join("::")]
         return unless receiver_entries
 
@@ -204,6 +212,26 @@ def build_method_completion(entry, node)
         )
       end
 
+      sig do
+        params(
+          variable_node: Prism::LocalVariableNode,
+          node: Prism::CallNode,
+        ).returns(Interface::CompletionItem)
+      end
+      def build_local_variable_completion(variable_node, node)
+        name = variable_node.name.to_s
+
+        Interface::CompletionItem.new(
+          label: name,
+          filter_text: name,
+          text_edit: Interface::TextEdit.new(
+            range: range_from_location(T.must(node.message_loc)),
+            new_text: name,
+          ),
+          kind: Constant::CompletionItemKind::VARIABLE,
+        )
+      end
+
       sig { params(label: String, node: Prism::StringNode).returns(Interface::CompletionItem) }
       def build_completion(label, node)
         # We should use the content location as we only replace the content and not the delimiters of the string

lib/ruby_lsp/listeners/definition.rb

@@ -13,23 +13,30 @@ class Definition
           global_state: GlobalState,
           uri: URI::Generic,
           nesting: T::Array[String],
+          local_variables: T::Array[Prism::LocalVariableNode],
           dispatcher: Prism::Dispatcher,
           typechecker_enabled: T::Boolean,
         ).void
       end
-      def initialize(response_builder, global_state, uri, nesting, dispatcher, typechecker_enabled) # rubocop:disable Metrics/ParameterLists
+      def initialize(response_builder, global_state, uri, nesting, local_variables, dispatcher, typechecker_enabled) # rubocop:disable Metrics/ParameterLists
         @response_builder = response_builder
         @global_state = global_state
         @index = T.let(global_state.index, RubyIndexer::Index)
         @uri = uri
         @nesting = nesting
         @typechecker_enabled = typechecker_enabled
+        @local_variables = local_variables
 
         dispatcher.register(
           self,
           :on_call_node_enter,
           :on_constant_read_node_enter,
           :on_constant_path_node_enter,
+          :on_local_variable_read_node_enter,
+          :on_local_variable_and_write_node_enter,
+          :on_local_variable_or_write_node_enter,
+          :on_local_variable_operator_write_node_enter,
+          :on_local_variable_write_node_enter,
         )
       end
 
@@ -60,6 +67,31 @@ def on_constant_read_node_enter(node)
         find_in_index(name)
       end
 
+      sig { params(node: Prism::LocalVariableReadNode).void }
+      def on_local_variable_read_node_enter(node)
+        find_local_variables(node)
+      end
+
+      sig { params(node: Prism::LocalVariableOrWriteNode).void }
+      def on_local_variable_or_write_node_enter(node)
+        find_local_variables(node)
+      end
+
+      sig { params(node: Prism::LocalVariableAndWriteNode).void }
+      def on_local_variable_and_write_node_enter(node)
+        find_local_variables(node)
+      end
+
+      sig { params(node: Prism::LocalVariableOperatorWriteNode).void }
+      def on_local_variable_operator_write_node_enter(node)
+        find_local_variables(node)
+      end
+
+      sig { params(node: Prism::LocalVariableWriteNode).void }
+      def on_local_variable_write_node_enter(node)
+        find_local_variables(node)
+      end
+
       private
 
       sig { params(node: Prism::CallNode).void }
@@ -156,6 +188,24 @@ def find_in_index(value)
           )
         end
       end
+
+      sig { params(node: Prism::LocalVariableNode).void }
+      def find_local_variables(node)
+        @local_variables.reverse.find do |variable_node|
+          next if variable_node == node
+          next unless variable_node.name == node.name
+
+          location = variable_node.location
+
+          @response_builder << Interface::Location.new(
+            uri: @uri.to_s,
+            range: Interface::Range.new(
+              start: Interface::Position.new(line: location.start_line - 1, character: location.start_column),
+              end: Interface::Position.new(line: location.end_line - 1, character: location.end_column),
+            ),
+          )
+        end
+      end
     end
   end
 end

lib/ruby_lsp/requests/completion.rb

@@ -16,6 +16,7 @@ module Requests
     # - Constants
     # - Require paths
     # - Methods invoked on self only
+    # - Local variables
     #
     # # Example
     #
@@ -63,6 +64,8 @@ def initialize(document, global_state, position, typechecker_enabled, dispatcher
           char_position,
           node_types: [Prism::CallNode, Prism::ConstantReadNode, Prism::ConstantPathNode],
         )
+        local_variables = document.locate_local_variables(document.tree, char_position)
+
         @response_builder = T.let(
           ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem].new,
           ResponseBuilders::CollectionResponseBuilder[Interface::CompletionItem],
@@ -72,6 +75,7 @@ def initialize(document, global_state, position, typechecker_enabled, dispatcher
           @response_builder,
           global_state,
           nesting,
+          local_variables,
           typechecker_enabled,
           dispatcher,
           document.uri,

lib/ruby_lsp/requests/definition.rb

@@ -17,6 +17,7 @@ module Requests
     # - Constants
     # - Require paths
     # - Methods invoked on self only
+    # - Local variables
     #
     # # Example
     #
@@ -46,7 +47,16 @@ def initialize(document, global_state, position, dispatcher, typechecker_enabled
 
         target, parent, nesting = document.locate_node(
           position,
-          node_types: [Prism::CallNode, Prism::ConstantReadNode, Prism::ConstantPathNode],
+          node_types: [
+            Prism::CallNode,
+            Prism::ConstantReadNode,
+            Prism::ConstantPathNode,
+            Prism::LocalVariableReadNode,
+            Prism::LocalVariableOrWriteNode,
+            Prism::LocalVariableAndWriteNode,
+            Prism::LocalVariableOperatorWriteNode,
+            Prism::LocalVariableWriteNode,
+          ],
         )
 
         if target.is_a?(Prism::ConstantReadNode) && parent.is_a?(Prism::ConstantPathNode)
@@ -57,11 +67,14 @@ def initialize(document, global_state, position, dispatcher, typechecker_enabled
           )
         end
 
+        local_variables = document.locate_local_variable_nodes(position)
+
         Listeners::Definition.new(
           @response_builder,
           global_state,
           document.uri,
           nesting,
+          local_variables,
           dispatcher,
           typechecker_enabled,
         )

test/requests/completion_test.rb

@@ -510,6 +510,60 @@ def process
     end
   end
 
+  def test_completion_for_local_variables
+    source = +<<~RUBY
+      class Foo
+        class_var = 1
+        def foo(a_required_param, an_optional_param = 1, *a_rest, a_required_named_param:, an_optional_named_param:, **an_options)
+          a_method_var = 1
+          a_method_var = 1
+          a_if_var = 1 if false
+          proc { a_proc_var_outside = 1 }
+
+          an_and_var &&= 1
+          an_or_var ||= 1
+          an_array_var1, an_array_var2 = [1, 2]
+
+          proc do |a_proc_arg|
+            a_proc_var = 1
+            a
+          end
+        end
+      end
+    RUBY
+
+    expected_variables = [
+      "a_required_param",
+      "an_optional_param",
+      "a_rest",
+      "a_required_named_param",
+      "an_optional_named_param",
+      "an_options",
+      "a_method_var",
+      "a_if_var",
+      "an_and_var",
+      "an_or_var",
+      "an_array_var1",
+      "an_array_var2",
+      "a_proc_arg",
+      "a_proc_var",
+    ]
+
+    with_server(source) do |server, uri|
+      with_file_structure(server) do
+        server.process_message(id: 1, method: "textDocument/completion", params: {
+          textDocument: { uri: uri },
+          position: { line: 14, character: 7 },
+        })
+
+        result = server.pop_response.response
+        assert_equal(expected_variables, result.map(&:label))
+        assert_equal(expected_variables, result.map(&:filter_text))
+        assert_equal(expected_variables, result.map { |completion| completion.text_edit.new_text })
+      end
+    end
+  end
+
   def test_completion_for_attributes
     source = +<<~RUBY
       class Foo