Properly assign instance variable owners

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -294,72 +294,27 @@ def on_def_node_leave(node)
 
     sig { params(node: Prism::InstanceVariableWriteNode).void }
     def on_instance_variable_write_node_enter(node)
-      name = node.name.to_s
-      return if name == "@"
-
-      @index << Entry::InstanceVariable.new(
-        name,
-        @file_path,
-        node.name_loc,
-        collect_comments(node),
-        @owner_stack.last,
-      )
+      handle_instance_variable(node, node.name_loc)
     end
 
     sig { params(node: Prism::InstanceVariableAndWriteNode).void }
     def on_instance_variable_and_write_node_enter(node)
-      name = node.name.to_s
-      return if name == "@"
-
-      @index << Entry::InstanceVariable.new(
-        name,
-        @file_path,
-        node.name_loc,
-        collect_comments(node),
-        @owner_stack.last,
-      )
+      handle_instance_variable(node, node.name_loc)
     end
 
     sig { params(node: Prism::InstanceVariableOperatorWriteNode).void }
     def on_instance_variable_operator_write_node_enter(node)
-      name = node.name.to_s
-      return if name == "@"
-
-      @index << Entry::InstanceVariable.new(
-        name,
-        @file_path,
-        node.name_loc,
-        collect_comments(node),
-        @owner_stack.last,
-      )
+      handle_instance_variable(node, node.name_loc)
     end
 
     sig { params(node: Prism::InstanceVariableOrWriteNode).void }
     def on_instance_variable_or_write_node_enter(node)
-      name = node.name.to_s
-      return if name == "@"
-
-      @index << Entry::InstanceVariable.new(
-        name,
-        @file_path,
-        node.name_loc,
-        collect_comments(node),
-        @owner_stack.last,
-      )
+      handle_instance_variable(node, node.name_loc)
     end
 
     sig { params(node: Prism::InstanceVariableTargetNode).void }
     def on_instance_variable_target_node_enter(node)
-      name = node.name.to_s
-      return if name == "@"
-
-      @index << Entry::InstanceVariable.new(
-        name,
-        @file_path,
-        node.location,
-        collect_comments(node),
-        @owner_stack.last,
-      )
+      handle_instance_variable(node, node.location)
     end
 
     sig { params(node: Prism::AliasMethodNode).void }
@@ -378,6 +333,30 @@ def on_alias_method_node_enter(node)
 
     private
 
+    sig do
+      params(
+        node: T.any(
+          Prism::InstanceVariableAndWriteNode,
+          Prism::InstanceVariableOperatorWriteNode,
+          Prism::InstanceVariableOrWriteNode,
+          Prism::InstanceVariableTargetNode,
+          Prism::InstanceVariableWriteNode,
+        ),
+        loc: Prism::Location,
+      ).void
+    end
+    def handle_instance_variable(node, loc)
+      name = node.name.to_s
+      return if name == "@"
+
+      # When instance variables are declared inside the class body, they turn into class instance variables rather than
+      # regular instance variables
+      owner = @owner_stack.last
+      owner = owner&.singleton_klass unless @inside_def
+
+      @index << Entry::InstanceVariable.new(name, @file_path, loc, collect_comments(node), owner)
+    end
+
     sig { params(node: Prism::CallNode).void }
     def handle_private_constant(node)
       arguments = node.arguments&.arguments

lib/ruby_indexer/test/instance_variables_test.rb

@@ -20,7 +20,9 @@ def initialize
       assert_entry("@a", Entry::InstanceVariable, "/fake/path/foo.rb:4-6:4-8")
 
       entry = T.must(@index["@a"]&.first)
-      assert_equal("Foo::Bar", T.must(entry.owner).name)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo::Bar", owner.name)
     end
 
     def test_instance_variable_and_write
@@ -38,7 +40,9 @@ def initialize
       assert_entry("@a", Entry::InstanceVariable, "/fake/path/foo.rb:4-6:4-8")
 
       entry = T.must(@index["@a"]&.first)
-      assert_equal("Foo::Bar", T.must(entry.owner).name)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo::Bar", owner.name)
     end
 
     def test_instance_variable_operator_write
@@ -56,7 +60,9 @@ def initialize
       assert_entry("@a", Entry::InstanceVariable, "/fake/path/foo.rb:4-6:4-8")
 
       entry = T.must(@index["@a"]&.first)
-      assert_equal("Foo::Bar", T.must(entry.owner).name)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo::Bar", owner.name)
     end
 
     def test_instance_variable_or_write
@@ -74,7 +80,9 @@ def initialize
       assert_entry("@a", Entry::InstanceVariable, "/fake/path/foo.rb:4-6:4-8")
 
       entry = T.must(@index["@a"]&.first)
-      assert_equal("Foo::Bar", T.must(entry.owner).name)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo::Bar", owner.name)
     end
 
     def test_instance_variable_target
@@ -93,10 +101,14 @@ def initialize
       assert_entry("@b", Entry::InstanceVariable, "/fake/path/foo.rb:4-10:4-12")
 
       entry = T.must(@index["@a"]&.first)
-      assert_equal("Foo::Bar", T.must(entry.owner).name)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo::Bar", owner.name)
 
       entry = T.must(@index["@b"]&.first)
-      assert_equal("Foo::Bar", T.must(entry.owner).name)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::Class, owner)
+      assert_equal("Foo::Bar", owner.name)
     end
 
     def test_empty_name_instance_variables
@@ -118,14 +130,41 @@ def test_class_instance_variables
         module Foo
           class Bar
             @a = 123
+
+            class << self
+              def hello
+                @b = 123
+              end
+
+              @c = 123
+            end
           end
         end
       RUBY
 
       assert_entry("@a", Entry::InstanceVariable, "/fake/path/foo.rb:2-4:2-6")
 
       entry = T.must(@index["@a"]&.first)
-      assert_equal("Foo::Bar", T.must(entry.owner).name)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::SingletonClass, owner)
+      assert_equal("Foo::Bar", owner.name)
+
+      assert_entry("@b", Entry::InstanceVariable, "/fake/path/foo.rb:6-8:6-10")
+
+      entry = T.must(@index["@b"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::SingletonClass, owner)
+      assert_equal("Foo::Bar", owner.name)
+
+      assert_entry("@c", Entry::InstanceVariable, "/fake/path/foo.rb:9-6:9-8")
+
+      entry = T.must(@index["@c"]&.first)
+      owner = T.must(entry.owner)
+      assert_instance_of(Entry::SingletonClass, owner)
+      assert_equal("Foo::Bar", owner.name)
+
+      foo_bar = T.must(@index["Foo::Bar"]&.first)
+      assert_same(owner, foo_bar.singleton_klass.singleton_klass)
     end
   end
 end