Support Hover, Completion, Go to definition for Class Variables

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -80,6 +80,11 @@ def initialize(index, dispatcher, parse_result, uri, collect_comments: false)
         :on_instance_variable_or_write_node_enter,
         :on_instance_variable_target_node_enter,
         :on_alias_method_node_enter,
+        :on_class_variable_and_write_node_enter,
+        :on_class_variable_operator_write_node_enter,
+        :on_class_variable_or_write_node_enter,
+        :on_class_variable_target_node_enter,
+        :on_class_variable_write_node_enter,
       )
     end
 
@@ -434,6 +439,31 @@ def on_alias_method_node_enter(node)
       )
     end
 
+    sig { params(node: Prism::ClassVariableAndWriteNode).void }
+    def on_class_variable_and_write_node_enter(node)
+      handle_class_variable(node, node.name_loc)
+    end
+
+    sig { params(node: Prism::ClassVariableOperatorWriteNode).void }
+    def on_class_variable_operator_write_node_enter(node)
+      handle_class_variable(node, node.name_loc)
+    end
+
+    sig { params(node: Prism::ClassVariableOrWriteNode).void }
+    def on_class_variable_or_write_node_enter(node)
+      handle_class_variable(node, node.name_loc)
+    end
+
+    sig { params(node: Prism::ClassVariableTargetNode).void }
+    def on_class_variable_target_node_enter(node)
+      handle_class_variable(node, node.location)
+    end
+
+    sig { params(node: Prism::ClassVariableWriteNode).void }
+    def on_class_variable_write_node_enter(node)
+      handle_class_variable(node, node.name_loc)
+    end
+
     sig do
       params(
         name: String,
@@ -552,6 +582,40 @@ def handle_global_variable(node, loc)
       ))
     end
 
+    sig do
+      params(
+        node: T.any(
+          Prism::ClassVariableAndWriteNode,
+          Prism::ClassVariableOperatorWriteNode,
+          Prism::ClassVariableOrWriteNode,
+          Prism::ClassVariableTargetNode,
+          Prism::ClassVariableWriteNode,
+        ),
+        loc: Prism::Location,
+      ).void
+    end
+    def handle_class_variable(node, loc)
+      name = node.name.to_s
+      return if name == "@@"
+
+      comments = collect_comments(node)
+
+      owner = @owner_stack.last
+
+      # set the class variable's owner to the attached context when defined within a singleton scope.
+      if owner.is_a?(Entry::SingletonClass)
+        owner = @owner_stack.reverse.find { |entry| !entry.name.include?("<Class:") }
+      end
+
+      @index.add(Entry::ClassVariable.new(
+        name,
+        @uri,
+        Location.from_prism_location(loc, @code_units_cache),
+        comments,
+        owner,
+      ))
+    end
+
     sig do
       params(
         node: T.any(

lib/ruby_indexer/lib/ruby_indexer/entry.rb

@@ -486,6 +486,26 @@ def initialize(target, unresolved_alias)
     # Represents a global variable e.g.: $DEBUG
     class GlobalVariable < Entry; end
 
+    # Represents a class variable e.g.: @@a = 1
+    class ClassVariable < Entry
+      sig { returns(T.nilable(Entry::Namespace)) }
+      attr_reader :owner
+
+      sig do
+        params(
+          name: String,
+          uri: URI::Generic,
+          location: Location,
+          comments: T.nilable(String),
+          owner: T.nilable(Entry::Namespace),
+        ).void
+      end
+      def initialize(name, uri, location, comments, owner)
+        super(name, uri, location, comments)
+        @owner = owner
+      end
+    end
+
     # Represents an instance variable e.g.: @a = 1
     class InstanceVariable < Entry
       sig { returns(T.nilable(Entry::Namespace)) }

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 = self[variable_name]&.grep(Entry::ClassVariable)
+      return unless entries&.any?
+
+      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_indexer/test/class_variables_test.rb

@@ -0,0 +1,140 @@
+# typed: true
+# frozen_string_literal: true
+
+require_relative "test_case"
+
+module RubyIndexer
+  class ClassVariableTest < TestCase
+    def test_class_variable_and_write
+      index(<<~RUBY)
+        class Foo
+          @@bar &&= 1
+        end
+      RUBY
+
+      assert_entry("@@bar", Entry::ClassVariable, "/fake/path/foo.rb:1-2:1-7")
+
+      entry = T.must(@index["@@bar"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo", owner.name)
+    end
+
+    def test_class_variable_operator_write
+      index(<<~RUBY)
+        class Foo
+          @@bar += 1
+        end
+      RUBY
+
+      assert_entry("@@bar", Entry::ClassVariable, "/fake/path/foo.rb:1-2:1-7")
+    end
+
+    def test_class_variable_or_write
+      index(<<~RUBY)
+        class Foo
+          @@bar ||= 1
+        end
+      RUBY
+
+      assert_entry("@@bar", Entry::ClassVariable, "/fake/path/foo.rb:1-2:1-7")
+    end
+
+    def test_class_variable_target_node
+      index(<<~RUBY)
+        class Foo
+          @@foo, @@bar = 1
+        end
+      RUBY
+
+      assert_entry("@@foo", Entry::ClassVariable, "/fake/path/foo.rb:1-2:1-7")
+      assert_entry("@@bar", Entry::ClassVariable, "/fake/path/foo.rb:1-9:1-14")
+
+      entry = T.must(@index["@@foo"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo", owner.name)
+
+      entry = T.must(@index["@@bar"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo", owner.name)
+    end
+
+    def test_class_variable_write
+      index(<<~RUBY)
+        class Foo
+          @@bar = 1
+        end
+      RUBY
+
+      assert_entry("@@bar", Entry::ClassVariable, "/fake/path/foo.rb:1-2:1-7")
+    end
+
+    def test_empty_name_class_variable
+      index(<<~RUBY)
+        module Foo
+          @@ = 1
+        end
+      RUBY
+
+      refute_entry("@@")
+    end
+
+    def test_top_level_class_variable
+      index(<<~RUBY)
+        @foo = 123
+      RUBY
+
+      entry = T.must(@index["@foo"]&.first)
+      assert_nil(entry.owner)
+    end
+
+    def test_class_variable_inside_self_method
+      index(<<~RUBY)
+        class Foo
+          def self.bar
+            @@bar = 123
+          end
+        end
+      RUBY
+
+      entry = T.must(@index["@@bar"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo", owner.name)
+    end
+
+    def test_class_variable_inside_singleton_class
+      index(<<~RUBY)
+        class Foo
+          class << self
+            @@bar = 123
+          end
+        end
+      RUBY
+
+      entry = T.must(@index["@@bar"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo", owner.name)
+    end
+
+    def test_class_variable_in_singleton_class_method
+      index(<<~RUBY)
+        class Foo
+          class << self
+            def self.bar
+              @@bar = 123
+            end
+          end
+        end
+      RUBY
+
+      entry = T.must(@index["@@bar"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo", owner.name)
+    end
+  end
+end

lib/ruby_lsp/listeners/completion.rb

@@ -97,6 +97,12 @@ def initialize( # rubocop:disable Metrics/ParameterLists
           :on_instance_variable_operator_write_node_enter,
           :on_instance_variable_or_write_node_enter,
           :on_instance_variable_target_node_enter,
+          :on_class_variable_and_write_node_enter,
+          :on_class_variable_operator_write_node_enter,
+          :on_class_variable_or_write_node_enter,
+          :on_class_variable_read_node_enter,
+          :on_class_variable_target_node_enter,
+          :on_class_variable_write_node_enter,
         )
       end
 
@@ -246,6 +252,36 @@ def on_instance_variable_target_node_enter(node)
         handle_instance_variable_completion(node.name.to_s, node.location)
       end
 
+      sig { params(node: Prism::ClassVariableAndWriteNode).void }
+      def on_class_variable_and_write_node_enter(node)
+        handle_class_variable_completion(node.name.to_s, node.name_loc)
+      end
+
+      sig { params(node: Prism::ClassVariableOperatorWriteNode).void }
+      def on_class_variable_operator_write_node_enter(node)
+        handle_class_variable_completion(node.name.to_s, node.name_loc)
+      end
+
+      sig { params(node: Prism::ClassVariableOrWriteNode).void }
+      def on_class_variable_or_write_node_enter(node)
+        handle_class_variable_completion(node.name.to_s, node.name_loc)
+      end
+
+      sig { params(node: Prism::ClassVariableTargetNode).void }
+      def on_class_variable_target_node_enter(node)
+        handle_class_variable_completion(node.name.to_s, node.location)
+      end
+
+      sig { params(node: Prism::ClassVariableReadNode).void }
+      def on_class_variable_read_node_enter(node)
+        handle_class_variable_completion(node.name.to_s, node.location)
+      end
+
+      sig { params(node: Prism::ClassVariableWriteNode).void }
+      def on_class_variable_write_node_enter(node)
+        handle_class_variable_completion(node.name.to_s, node.name_loc)
+      end
+
       private
 
       sig { params(name: String, range: Interface::Range).void }
@@ -326,6 +362,37 @@ def handle_global_variable_completion(name, location)
         end
       end
 
+      sig { params(name: String, location: Prism::Location).void }
+      def handle_class_variable_completion(name, location)
+        type = @type_inferrer.infer_receiver_type(@node_context)
+        return unless type
+
+        range = range_from_location(location)
+
+        @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: variable_name,
+            label_details: label_details,
+            text_edit: Interface::TextEdit.new(
+              range: range,
+              new_text: variable_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 }
       def handle_instance_variable_completion(name, location)
         # Sorbet enforces that all instance variables be declared on typed strict or higher, which means it will be able