From 2f1cc887afa5f80712d1e925d035176aad75c2b6 Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Fri, 14 Jun 2024 17:12:39 -0400
Subject: [PATCH] Push singleton context when inside singleton method

---
 .../lib/ruby_indexer/declaration_listener.rb  | 14 +++++++-
 .../test/instance_variables_test.rb           | 32 +++++++++++++++++++
 test/fixtures/prism                           |  2 +-
 3 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
index 996c12b8e..3393c636d 100644
--- a/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
+++ b/lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb
@@ -303,6 +303,8 @@ def on_def_node_enter(node)
           @owner_stack.last,
         ))
       when Prism::SelfNode
+        singleton = singleton_klass
+
         @index.add(Entry::Method.new(
           method_name,
           @file_path,
@@ -310,14 +312,24 @@ def on_def_node_enter(node)
           comments,
           list_params(node.parameters),
           current_visibility,
-          singleton_klass,
+          singleton,
         ))
+
+        if singleton
+          @owner_stack << singleton
+          @stack << "<Class:#{@stack.last}>"
+        end
       end
     end
 
     sig { params(node: Prism::DefNode).void }
     def on_def_node_leave(node)
       @inside_def = false
+
+      if node.receiver.is_a?(Prism::SelfNode)
+        @owner_stack.pop
+        @stack.pop
+      end
     end
 
     sig { params(node: Prism::InstanceVariableWriteNode).void }
diff --git a/lib/ruby_indexer/test/instance_variables_test.rb b/lib/ruby_indexer/test/instance_variables_test.rb
index 291c76b0b..1bf14a857 100644
--- a/lib/ruby_indexer/test/instance_variables_test.rb
+++ b/lib/ruby_indexer/test/instance_variables_test.rb
@@ -172,5 +172,37 @@ def test_top_level_instance_variables
       entry = T.must(@index["@a"]&.first)
       assert_nil(entry.owner)
     end
+
+    def test_class_instance_variables_inside_self_method
+      index(<<~RUBY)
+        class Foo
+          def self.bar
+            @a = 123
+          end
+        end
+      RUBY
+
+      entry = T.must(@index["@a"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::SingletonClass, owner)
+      assert_equal("Foo::<Class:Foo>", owner.name)
+    end
+
+    def test_instance_variable_inside_dynamic_method_declaration
+      index(<<~RUBY)
+        class Foo
+          def something.bar
+            @a = 123
+          end
+        end
+      RUBY
+
+      # If the surrounding method is beind defined on any dynamic value that isn't `self`, then we attribute the
+      # instance variable to the wrong owner since there's no way to understand that statically
+      entry = T.must(@index["@a"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo", owner.name)
+    end
   end
 end
diff --git a/test/fixtures/prism b/test/fixtures/prism
index b6a97a43c..8e7d126f3 160000
--- a/test/fixtures/prism
+++ b/test/fixtures/prism
@@ -1 +1 @@
-Subproject commit b6a97a43c2b5f8529e026125c7fe2779454692ca
+Subproject commit 8e7d126f3e64bf3cabe6305943bd4fca2ba5e8ca