Support Hover, Completion, Go to definition for Class Variables

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -584,6 +584,17 @@ def resolve_instance_variable(variable_name, owner_name)
       entries.select { |e| ancestors.include?(e.owner&.name) }
     end
 
+    sig { params(variable_name: String, owner_name: String).returns(T.nilable(T::Array[Entry::ClassVariable])) }
+    def resolve_class_variable(variable_name, owner_name)
+      entries = T.cast(self[variable_name], T.nilable(T::Array[Entry::ClassVariable]))
+      return unless entries
+
+      ancestors = linearized_ancestors_of(owner_name)
+      return if ancestors.empty?
+
+      entries.select { |e| ancestors.include?(e.owner&.name) }
+    end
+
     # Returns a list of possible candidates for completion of instance variables for a given owner name. The name must
     # include the `@` prefix
     sig { params(name: String, owner_name: String).returns(T::Array[Entry::InstanceVariable]) }
@@ -596,6 +607,16 @@ def instance_variable_completion_candidates(name, owner_name)
       variables
     end
 
+    sig { params(name: String, owner_name: String).returns(T::Array[Entry::ClassVariable]) }
+    def class_variable_completion_candidates(name, owner_name)
+      entries = T.cast(prefix_search(name).flatten, T::Array[Entry::ClassVariable])
+      ancestors = linearized_ancestors_of(owner_name)
+
+      variables = entries.select { |e| ancestors.any?(e.owner&.name) }
+      variables.uniq!(&:name)
+      variables
+    end
+
     # Synchronizes a change made to the given URI. This method will ensure that new declarations are indexed, removed
     # declarations removed and that the ancestor linearization cache is cleared if necessary
     sig { params(uri: URI::Generic, source: String).void }

lib/ruby_lsp/listeners/completion.rb

@@ -364,25 +364,33 @@ def handle_global_variable_completion(name, location)
 
       sig { params(name: String, location: Prism::Location).void }
       def handle_class_variable_completion(name, location)
-        candidates = @index.prefix_search(name)
-
-        return if candidates.none?
+        type = @type_inferrer.infer_receiver_type(@node_context)
+        return unless type
 
         range = range_from_location(location)
 
-        candidates.flatten.uniq(&:name).each do |entry|
-          entry_name = entry.name
+        @index.class_variable_completion_candidates(name, type.name).each do |entry|
+          variable_name = entry.name
+
+          label_details = Interface::CompletionItemLabelDetails.new(
+            description: entry.file_name,
+          )
 
           @response_builder << Interface::CompletionItem.new(
-            label: entry_name,
-            filter_text: entry_name,
-            label_details: Interface::CompletionItemLabelDetails.new(
-              description: entry.file_name,
+            label: variable_name,
+            label_details: label_details,
+            text_edit: Interface::TextEdit.new(
+              range: range,
+              new_text: variable_name,
             ),
-            text_edit: Interface::TextEdit.new(range: range, new_text: entry_name),
             kind: Constant::CompletionItemKind::FIELD,
+            data: {
+              owner_name: entry.owner&.name,
+            },
           )
         end
+      rescue RubyIndexer::Index::NonExistingNamespaceError
+        # If by any chance we haven't indexed the owner, then there's no way to find the right declaration
       end
 
       sig { params(name: String, location: Prism::Location).void }

lib/ruby_lsp/listeners/definition.rb

@@ -270,8 +270,10 @@ def handle_global_variable_definition(name)
 
       sig { params(name: String).void }
       def handle_class_variable_definition(name)
-        entries = @index[name]
+        type = @type_inferrer.infer_receiver_type(@node_context)
+        return unless type
 
+        entries = @index.resolve_class_variable(name, type.name)
         return unless entries
 
         entries.each do |entry|
@@ -285,6 +287,8 @@ def handle_class_variable_definition(name)
             ),
           )
         end
+      rescue RubyIndexer::Index::NonExistingNamespaceError
+        # If by any chance we haven't indexed the owner, then there's no way to find the right declaration
       end
 
       sig { params(name: String).void }

lib/ruby_lsp/listeners/hover.rb

@@ -361,12 +361,17 @@ def handle_global_variable_hover(name)
 
       sig { params(name: String).void }
       def handle_class_variable_hover(name)
-        entries = @index[name]
+        type = @type_inferrer.infer_receiver_type(@node_context)
+        return unless type
+
+        entries = @index.resolve_class_variable(name, type.name)
         return unless entries
 
         categorized_markdown_from_index_entries(name, entries).each do |category, content|
           @response_builder.push(content, category: category)
         end
+      rescue RubyIndexer::Index::NonExistingNamespaceError
+        # If by any chance we haven't indexed the owner, then there's no way to find the right declaration
       end
 
       sig { params(name: String, location: Prism::Location).void }

lib/ruby_lsp/type_inferrer.rb

@@ -21,7 +21,9 @@ def infer_receiver_type(node_context)
         infer_receiver_for_call_node(node, node_context)
       when Prism::InstanceVariableReadNode, Prism::InstanceVariableAndWriteNode, Prism::InstanceVariableWriteNode,
         Prism::InstanceVariableOperatorWriteNode, Prism::InstanceVariableOrWriteNode, Prism::InstanceVariableTargetNode,
-        Prism::SuperNode, Prism::ForwardingSuperNode
+        Prism::SuperNode, Prism::ForwardingSuperNode, Prism::ClassVariableAndWriteNode, Prism::ClassVariableWriteNode,
+        Prism::ClassVariableOperatorWriteNode, Prism::ClassVariableOrWriteNode, Prism::ClassVariableReadNode,
+        Prism::ClassVariableTargetNode
         self_receiver_handling(node_context)
       end
     end

test/requests/completion_test.rb

@@ -1117,6 +1117,40 @@ def baz
     end
   end
 
+  def test_completion_for_inherited_class_variables
+    source = <<~RUBY
+      module Foo
+        def set_variable
+          @@bar = 9
+        end
+      end
+
+      class Parent
+        def set_variable
+          @@baz = 5
+        end
+      end
+
+      class Child < Parent
+        include Foo
+
+        def do_something
+          @
+        end
+      end
+    RUBY
+
+    with_server(source, stub_no_typechecker: true) do |server, uri|
+      server.process_message(id: 1, method: "textDocument/completion", params: {
+        textDocument: { uri: uri },
+        position: { line: 16, character: 5 },
+      })
+
+      result = server.pop_response.response
+      assert_equal(["@@bar", "@@baz"], result.map(&:label))
+    end
+  end
+
   def test_completion_for_class_variables_show_only_uniq_entries
     source = <<~RUBY
       class Foo

test/requests/definition_expectations_test.rb

@@ -774,6 +774,42 @@ def baz
     end
   end
 
+  def test_definition_for_inherited_class_variables
+    source = <<~RUBY
+      module Foo
+        def set_variable
+          @@bar = 1
+        end
+      end
+
+      class Parent
+        def set_variable
+          @@bar = 5
+        end
+      end
+
+      class Child < Parent
+        include Foo
+
+        def do_something
+          @@bar
+        end
+      end
+    RUBY
+
+    with_server(source) do |server, uri|
+      server.process_message(
+        id: 1,
+        method: "textDocument/definition",
+        params: { textDocument: { uri: uri }, position: { character: 4, line: 16 } },
+      )
+      response = server.pop_response.response
+
+      assert_equal(2, response[0].range.start.line)
+      assert_equal(8, response[1].range.start.line)
+    end
+  end
+
   def test_definition_for_instance_variables
     source = <<~RUBY
       class Foo

test/requests/hover_expectations_test.rb

@@ -462,6 +462,44 @@ def baz
     end
   end
 
+  def test_hovering_for_inherited_class_variables
+    source = <<~RUBY
+      module Foo
+        def set_variable
+          # Foo
+          @@bar = 1
+        end
+      end
+
+      class Parent
+        def set_variable
+          # Parent
+          @@bar = 5
+        end
+      end
+
+      class Child < Parent
+        include Foo
+
+        def do_something
+          @@bar
+        end
+      end
+    RUBY
+
+    with_server(source) do |server, uri|
+      server.process_message(
+        id: 1,
+        method: "textDocument/hover",
+        params: { textDocument: { uri: uri }, position: { character: 4, line: 18 } },
+      )
+
+      contents = server.pop_response.response.contents.value
+      assert_match("Foo", contents)
+      assert_match("Parent", contents)
+    end
+  end
+
   def test_hovering_over_inherited_methods
     source = <<~RUBY
       module Foo