Push singleton context when inside singleton method

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -303,21 +303,33 @@ def on_def_node_enter(node)
           @owner_stack.last,
         ))
       when Prism::SelfNode
+        singleton = singleton_klass
+
         @index.add(Entry::Method.new(
           method_name,
           @file_path,
           node.location,
           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 }

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