Generalize scope so that it can be used with any locals (#2603)

* Rename ParameterScope to Scope

* Rename << to add

* Generalize scope to be used with any locals

lib/ruby_lsp/internal.rb

@@ -26,7 +26,7 @@
 require "ruby_indexer/ruby_indexer"
 require "core_ext/uri"
 require "ruby_lsp/utils"
-require "ruby_lsp/parameter_scope"
+require "ruby_lsp/scope"
 require "ruby_lsp/global_state"
 require "ruby_lsp/server"
 require "ruby_lsp/type_inferrer"

lib/ruby_lsp/listeners/semantic_highlighting.rb

@@ -27,7 +27,7 @@ class SemanticHighlighting
       def initialize(dispatcher, response_builder)
         @response_builder = response_builder
         @special_methods = T.let(nil, T.nilable(T::Array[String]))
-        @current_scope = T.let(ParameterScope.new, ParameterScope)
+        @current_scope = T.let(Scope.new, Scope)
         @inside_regex_capture = T.let(false, T::Boolean)
         @inside_implicit_node = T.let(false, T::Boolean)
 
@@ -102,7 +102,7 @@ def on_match_write_node_leave(node)
 
       sig { params(node: Prism::DefNode).void }
       def on_def_node_enter(node)
-        @current_scope = ParameterScope.new(@current_scope)
+        @current_scope = Scope.new(@current_scope)
       end
 
       sig { params(node: Prism::DefNode).void }
@@ -112,7 +112,7 @@ def on_def_node_leave(node)
 
       sig { params(node: Prism::BlockNode).void }
       def on_block_node_enter(node)
-        @current_scope = ParameterScope.new(@current_scope)
+        @current_scope = Scope.new(@current_scope)
       end
 
       sig { params(node: Prism::BlockNode).void }
@@ -128,39 +128,39 @@ def on_block_local_variable_node_enter(node)
       sig { params(node: Prism::BlockParameterNode).void }
       def on_block_parameter_node_enter(node)
         name = node.name
-        @current_scope << name.to_sym if name
+        @current_scope.add(name.to_sym, :parameter) if name
       end
 
       sig { params(node: Prism::RequiredKeywordParameterNode).void }
       def on_required_keyword_parameter_node_enter(node)
-        @current_scope << node.name
+        @current_scope.add(node.name, :parameter)
       end
 
       sig { params(node: Prism::OptionalKeywordParameterNode).void }
       def on_optional_keyword_parameter_node_enter(node)
-        @current_scope << node.name
+        @current_scope.add(node.name, :parameter)
       end
 
       sig { params(node: Prism::KeywordRestParameterNode).void }
       def on_keyword_rest_parameter_node_enter(node)
         name = node.name
-        @current_scope << name.to_sym if name
+        @current_scope.add(name.to_sym, :parameter) if name
       end
 
       sig { params(node: Prism::OptionalParameterNode).void }
       def on_optional_parameter_node_enter(node)
-        @current_scope << node.name
+        @current_scope.add(node.name, :parameter)
       end
 
       sig { params(node: Prism::RequiredParameterNode).void }
       def on_required_parameter_node_enter(node)
-        @current_scope << node.name
+        @current_scope.add(node.name, :parameter)
       end
 
       sig { params(node: Prism::RestParameterNode).void }
       def on_rest_parameter_node_enter(node)
         name = node.name
-        @current_scope << name.to_sym if name
+        @current_scope.add(name.to_sym, :parameter) if name
       end
 
       sig { params(node: Prism::SelfNode).void }
@@ -170,8 +170,8 @@ def on_self_node_enter(node)
 
       sig { params(node: Prism::LocalVariableWriteNode).void }
       def on_local_variable_write_node_enter(node)
-        type = @current_scope.type_for(node.name)
-        @response_builder.add_token(node.name_loc, type) if type == :parameter
+        local = @current_scope.lookup(node.name)
+        @response_builder.add_token(node.name_loc, :parameter) if local&.type == :parameter
       end
 
       sig { params(node: Prism::LocalVariableReadNode).void }
@@ -184,25 +184,26 @@ def on_local_variable_read_node_enter(node)
           return
         end
 
-        @response_builder.add_token(node.location, @current_scope.type_for(node.name))
+        local = @current_scope.lookup(node.name)
+        @response_builder.add_token(node.location, local&.type || :variable)
       end
 
       sig { params(node: Prism::LocalVariableAndWriteNode).void }
       def on_local_variable_and_write_node_enter(node)
-        type = @current_scope.type_for(node.name)
-        @response_builder.add_token(node.name_loc, type) if type == :parameter
+        local = @current_scope.lookup(node.name)
+        @response_builder.add_token(node.name_loc, :parameter) if local&.type == :parameter
       end
 
       sig { params(node: Prism::LocalVariableOperatorWriteNode).void }
       def on_local_variable_operator_write_node_enter(node)
-        type = @current_scope.type_for(node.name)
-        @response_builder.add_token(node.name_loc, type) if type == :parameter
+        local = @current_scope.lookup(node.name)
+        @response_builder.add_token(node.name_loc, :parameter) if local&.type == :parameter
       end
 
       sig { params(node: Prism::LocalVariableOrWriteNode).void }
       def on_local_variable_or_write_node_enter(node)
-        type = @current_scope.type_for(node.name)
-        @response_builder.add_token(node.name_loc, type) if type == :parameter
+        local = @current_scope.lookup(node.name)
+        @response_builder.add_token(node.name_loc, :parameter) if local&.type == :parameter
       end
 
       sig { params(node: Prism::LocalVariableTargetNode).void }
@@ -213,7 +214,8 @@ def on_local_variable_target_node_enter(node)
         # prevent pushing local variable target tokens. See https://github.com/ruby/prism/issues/1912
         return if @inside_regex_capture
 
-        @response_builder.add_token(node.location, @current_scope.type_for(node.name))
+        local = @current_scope.lookup(node.name)
+        @response_builder.add_token(node.location, local&.type || :variable)
       end
 
       sig { params(node: Prism::ClassNode).void }
@@ -311,7 +313,8 @@ def process_regexp_locals(node)
           capture_name_offset = T.must(content.index("(?<#{name}>")) + 3
           local_var_loc = loc.copy(start_offset: loc.start_offset + capture_name_offset, length: name.length)
 
-          @response_builder.add_token(local_var_loc, @current_scope.type_for(name))
+          local = @current_scope.lookup(name)
+          @response_builder.add_token(local_var_loc, local&.type || :variable)
         end
       end
     end

lib/ruby_lsp/scope.rb

@@ -0,0 +1,47 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyLsp
+  class Scope
+    extend T::Sig
+
+    sig { returns(T.nilable(Scope)) }
+    attr_reader :parent
+
+    sig { params(parent: T.nilable(Scope)).void }
+    def initialize(parent = nil)
+      @parent = parent
+
+      # A hash of name => type
+      @locals = T.let({}, T::Hash[Symbol, Local])
+    end
+
+    # Add a new local to this scope. The types should only be `:parameter` or `:variable`
+    sig { params(name: T.any(String, Symbol), type: Symbol).void }
+    def add(name, type)
+      @locals[name.to_sym] = Local.new(type)
+    end
+
+    sig { params(name: T.any(String, Symbol)).returns(T.nilable(Local)) }
+    def lookup(name)
+      sym = name.to_sym
+      entry = @locals[sym]
+      return entry if entry
+      return unless @parent
+
+      @parent.lookup(sym)
+    end
+
+    class Local
+      extend T::Sig
+
+      sig { returns(Symbol) }
+      attr_reader :type
+
+      sig { params(type: Symbol).void }
+      def initialize(type)
+        @type = type
+      end
+    end
+  end
+end

test/scope_test.rb

@@ -0,0 +1,26 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+class ScopeTest < Minitest::Test
+  def test_finding_parameter_in_immediate_scope
+    scope = RubyLsp::Scope.new
+    scope.add("foo", :parameter)
+
+    assert_equal(:parameter, T.must(scope.lookup("foo")).type)
+  end
+
+  def test_finding_parameter_in_parent_scope
+    parent = RubyLsp::Scope.new
+    parent.add("foo", :parameter)
+
+    scope = RubyLsp::Scope.new(parent)
+    assert_equal(:parameter, T.must(scope.lookup("foo")).type)
+  end
+
+  def test_not_finding_parameter
+    scope = RubyLsp::Scope.new
+    refute(scope.lookup("foo"))
+  end
+end