Generalize scope to be used with any locals

lib/ruby_lsp/listeners/semantic_highlighting.rb

@@ -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.add(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.add(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.add(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.add(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.add(node.name)
+        @current_scope.add(node.name, :parameter)
       end
 
       sig { params(node: Prism::RequiredParameterNode).void }
       def on_required_parameter_node_enter(node)
-        @current_scope.add(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.add(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

@@ -11,23 +11,37 @@ class Scope
     sig { params(parent: T.nilable(Scope)).void }
     def initialize(parent = nil)
       @parent = parent
-      @parameters = T.let(Set.new, T::Set[Symbol])
-    end
 
-    sig { params(name: T.any(String, Symbol)).void }
-    def add(name)
-      @parameters << name.to_sym
+      # A hash of name => type
+      @locals = T.let({}, T::Hash[Symbol, Local])
     end
 
-    sig { params(name: T.any(Symbol, String)).returns(Symbol) }
-    def type_for(name)
-      parameter?(name) ? :parameter : :variable
+    # 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(Symbol, String)).returns(T::Boolean) }
-    def parameter?(name)
+    sig { params(name: T.any(String, Symbol)).returns(T.nilable(Local)) }
+    def lookup(name)
       sym = name.to_sym
-      @parameters.include?(sym) || (!@parent.nil? && @parent.parameter?(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

@@ -6,22 +6,21 @@
 class ScopeTest < Minitest::Test
   def test_finding_parameter_in_immediate_scope
     scope = RubyLsp::Scope.new
-    scope.add("foo")
+    scope.add("foo", :parameter)
 
-    assert(scope.parameter?("foo"))
+    assert_equal(:parameter, T.must(scope.lookup("foo")).type)
   end
 
   def test_finding_parameter_in_parent_scope
     parent = RubyLsp::Scope.new
-    parent.add("foo")
+    parent.add("foo", :parameter)
 
     scope = RubyLsp::Scope.new(parent)
-
-    assert(scope.parameter?("foo"))
+    assert_equal(:parameter, T.must(scope.lookup("foo")).type)
   end
 
   def test_not_finding_parameter
     scope = RubyLsp::Scope.new
-    refute(scope.parameter?("foo"))
+    refute(scope.lookup("foo"))
   end
 end