Search for class variables completion candidates in attached namespace (

#3100)

### Motivation

Class variables can be referenced both from singleton and attached contexts. To save memory, we decided to always associate class variables with only the attached version of namespaces, but we didn't account for this in either resolving class variables or finding completion candidates.

### Implementation

We need to always search for class variables within the attached namespace's ancestor chain.

### Automated Tests

Added tests that reproduce the bugs.

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -601,29 +601,52 @@ 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)
+      ancestors = linearized_attached_ancestors(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]) }
+    sig do
+      params(name: String, owner_name: String).returns(T::Array[T.any(Entry::InstanceVariable, Entry::ClassVariable)])
+    end
     def instance_variable_completion_candidates(name, owner_name)
-      entries = T.cast(prefix_search(name).flatten, T::Array[Entry::InstanceVariable])
+      entries = T.cast(prefix_search(name).flatten, T::Array[T.any(Entry::InstanceVariable, Entry::ClassVariable)])
+      # Avoid wasting time linearizing ancestors if we didn't find anything
+      return entries if entries.empty?
+
       ancestors = linearized_ancestors_of(owner_name)
 
-      variables = entries.select { |e| ancestors.any?(e.owner&.name) }
+      instance_variables, class_variables = entries.partition { |e| e.is_a?(Entry::InstanceVariable) }
+      variables = instance_variables.select { |e| ancestors.any?(e.owner&.name) }
+
+      # Class variables are only owned by the attached class in our representation. If the owner is in a singleton
+      # context, we have to search for ancestors of the attached class
+      if class_variables.any?
+        name_parts = owner_name.split("::")
+
+        if name_parts.last&.start_with?("<Class:")
+          attached_name = T.must(name_parts[0..-2]).join("::")
+          attached_ancestors = linearized_ancestors_of(attached_name)
+          variables.concat(class_variables.select { |e| attached_ancestors.any?(e.owner&.name) })
+        else
+          variables.concat(class_variables.select { |e| ancestors.any?(e.owner&.name) })
+        end
+      end
+
       variables.uniq!(&: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)
+      # Avoid wasting time linearizing ancestors if we didn't find anything
+      return entries if entries.empty?
 
+      ancestors = linearized_attached_ancestors(owner_name)
       variables = entries.select { |e| ancestors.any?(e.owner&.name) }
       variables.uniq!(&:name)
       variables
@@ -725,6 +748,20 @@ def entries_for(uri, type = nil)
 
     private
 
+    # Always returns the linearized ancestors for the attached class, regardless of whether `name` refers to a singleton
+    # or attached namespace
+    sig { params(name: String).returns(T::Array[String]) }
+    def linearized_attached_ancestors(name)
+      name_parts = name.split("::")
+
+      if name_parts.last&.start_with?("<Class:")
+        attached_name = T.must(name_parts[0..-2]).join("::")
+        linearized_ancestors_of(attached_name)
+      else
+        linearized_ancestors_of(name)
+      end
+    end
+
     # Runs the registered included hooks
     sig { params(fully_qualified_name: String, nesting: T::Array[String]).void }
     def run_included_hooks(fully_qualified_name, nesting)

lib/ruby_indexer/test/index_test.rb

@@ -2109,5 +2109,57 @@ class Foo
       refute_nil(entry, "Expected indexer to be able to handle unsaved URIs")
       assert_equal("I added this comment!", entry.comments)
     end
+
+    def test_instance_variable_completion_returns_class_variables_too
+      index(<<~RUBY)
+        class Parent
+          @@abc = 123
+        end
+
+        class Child < Parent
+          @@adf = 123
+
+          def self.do
+          end
+        end
+      RUBY
+
+      abc, adf = @index.instance_variable_completion_candidates("@", "Child::<Class:Child>")
+
+      refute_nil(abc)
+      refute_nil(adf)
+
+      assert_equal("@@abc", abc.name)
+      assert_equal("@@adf", adf.name)
+    end
+
+    def test_class_variable_completion_from_singleton_context
+      index(<<~RUBY)
+        class Foo
+          @@hello = 123
+
+          def self.do
+          end
+        end
+      RUBY
+
+      candidates = @index.class_variable_completion_candidates("@@", "Foo::<Class:Foo>")
+      refute_empty(candidates)
+
+      assert_equal("@@hello", candidates.first&.name)
+    end
+
+    def test_resolve_class_variable_in_singleton_context
+      index(<<~RUBY)
+        class Foo
+          @@hello = 123
+        end
+      RUBY
+
+      candidates = @index.resolve_class_variable("@@hello", "Foo::<Class:Foo>")
+      refute_empty(candidates)
+
+      assert_equal("@@hello", candidates.first&.name)
+    end
   end
 end

test/requests/completion_test.rb

@@ -1117,6 +1117,31 @@ def baz
     end
   end
 
+  def test_completion_for_class_variables_node
+    source = <<~RUBY
+      class Foo
+        def set_variables
+          @@foo = 1
+          @@foobar = 2
+        end
+
+        def bar
+          @@
+        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: 7, character: 6 },
+      })
+
+      result = server.pop_response.response
+      assert_equal(["@@foo", "@@foobar"], result.map(&:label))
+    end
+  end
+
   def test_completion_for_inherited_class_variables
     source = <<~RUBY
       module Foo