From a3ed38099e2754f9e7d6b4d55822a6100e1d97ea Mon Sep 17 00:00:00 2001
From: Tyler Lemburg <trlemburg@gmail.com>
Date: Tue, 19 Nov 2024 09:33:55 -0600
Subject: [PATCH] Add support for `extend self` to `handle_module_operation`
 (#2855)

By adding the instance methods as class methods through the
singleton mixin operations, we are able to get all instance
methods "copied" over into the class methods implicitly.

I don't think it's necessary to test whether this extend self
applies when the module is opened after being closed again,
or re-opened in another file.

Closes #2782
---
 .../lib/ruby_indexer/declaration_listener.rb  | 24 ++++++++------
 lib/ruby_indexer/test/index_test.rb           | 32 +++++++++++++++++++
 2 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
index 2815ec1a8..0aa89ea66 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
+++ b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
@@ -751,16 +751,22 @@ def handle_module_operation(node, operation)
       return unless arguments
 
       arguments.each do |node|
-        next unless node.is_a?(Prism::ConstantReadNode) || node.is_a?(Prism::ConstantPathNode)
-
-        case operation
-        when :include
-          owner.mixin_operations << Entry::Include.new(node.full_name)
-        when :prepend
-          owner.mixin_operations << Entry::Prepend.new(node.full_name)
-        when :extend
+        next unless node.is_a?(Prism::ConstantReadNode) || node.is_a?(Prism::ConstantPathNode) ||
+          (node.is_a?(Prism::SelfNode) && operation == :extend)
+
+        if node.is_a?(Prism::SelfNode)
           singleton = @index.existing_or_new_singleton_class(owner.name)
-          singleton.mixin_operations << Entry::Include.new(node.full_name)
+          singleton.mixin_operations << Entry::Include.new(owner.name)
+        else
+          case operation
+          when :include
+            owner.mixin_operations << Entry::Include.new(node.full_name)
+          when :prepend
+            owner.mixin_operations << Entry::Prepend.new(node.full_name)
+          when :extend
+            singleton = @index.existing_or_new_singleton_class(owner.name)
+            singleton.mixin_operations << Entry::Include.new(node.full_name)
+          end
         end
       rescue Prism::ConstantPathNode::DynamicPartsInConstantPathError,
              Prism::ConstantPathNode::MissingNodesInConstantPathError
diff --git a/lib/ruby_indexer/test/index_test.rb b/lib/ruby_indexer/test/index_test.rb
index 2723edc0a..69b5b68af 100644
--- a/lib/ruby_indexer/test/index_test.rb
+++ b/lib/ruby_indexer/test/index_test.rb
@@ -1672,6 +1672,38 @@ def test_linearizing_singleton_object
       )
     end
 
+    def test_extend_self
+      @index.index_single(IndexablePath.new(nil, "/fake/path/foo.rb"), <<~RUBY)
+        module Foo
+          def bar
+          end
+
+          extend self
+
+          def baz
+          end
+        end
+      RUBY
+
+      ["bar", "baz"].product(["Foo", "Foo::<Class:Foo>"]).each do |method, receiver|
+        entry = @index.resolve_method(method, receiver)&.first
+        refute_nil(entry)
+        assert_equal(method, T.must(entry).name)
+      end
+
+      assert_equal(
+        [
+          "Foo::<Class:Foo>",
+          "Foo",
+          "Module",
+          "Object",
+          "Kernel",
+          "BasicObject",
+        ],
+        @index.linearized_ancestors_of("Foo::<Class:Foo>"),
+      )
+    end
+
     def test_linearizing_singleton_ancestors
       @index.index_single(IndexablePath.new(nil, "/fake/path/foo.rb"), <<~RUBY)
         module First