Handle singleton class contexts

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -33,6 +33,8 @@ def initialize(index, dispatcher, parse_result, file_path)
         :on_class_node_leave,
         :on_module_node_enter,
         :on_module_node_leave,
+        :on_singleton_class_node_enter,
+        :on_singleton_class_node_leave,
         :on_def_node_enter,
         :on_def_node_leave,
         :on_call_node_enter,
@@ -115,6 +117,36 @@ def on_module_node_leave(node)
       @visibility_stack.pop
     end
 
+    sig { params(node: Prism::SingletonClassNode).void }
+    def on_singleton_class_node_enter(node)
+      @visibility_stack.push(Entry::Visibility::PUBLIC)
+
+      current_owner = @owner_stack.last
+
+      if current_owner
+        @stack << "<Class:#{@stack.last}>"
+
+        existing_entries = T.cast(@index[@stack.join("::")], T.nilable(T::Array[Entry::SingletonClass]))
+
+        if existing_entries
+          entry = T.must(existing_entries.first)
+          entry.update_singleton_information(node.location, collect_comments(node))
+        else
+          entry = Entry::SingletonClass.new(@stack, @file_path, node.location, collect_comments(node), nil)
+          @index << entry
+        end
+
+        @owner_stack << entry
+      end
+    end
+
+    sig { params(node: Prism::SingletonClassNode).void }
+    def on_singleton_class_node_leave(node)
+      @stack.pop
+      @owner_stack.pop
+      @visibility_stack.pop
+    end
+
     sig { params(node: Prism::MultiWriteNode).void }
     def on_multi_write_node_enter(node)
       value = node.value
@@ -246,7 +278,7 @@ def on_def_node_enter(node)
 
       case node.receiver
       when nil
-        @index << Entry::InstanceMethod.new(
+        @index << Entry::Method.new(
           method_name,
           @file_path,
           node.location,
@@ -256,14 +288,14 @@ def on_def_node_enter(node)
           @owner_stack.last,
         )
       when Prism::SelfNode
-        @index << Entry::SingletonMethod.new(
+        @index << Entry::Method.new(
           method_name,
           @file_path,
           node.location,
           comments,
           node.parameters,
           current_visibility,
-          @owner_stack.last,
+          singleton_klass,
         )
       end
     end
@@ -275,72 +307,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 }
@@ -359,6 +346,28 @@ 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 = @inside_def ? @owner_stack.last : singleton_klass
+      @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
@@ -558,5 +567,24 @@ def handle_module_operation(node, operation)
     def current_visibility
       T.must(@visibility_stack.last)
     end
+
+    sig { returns(T.nilable(Entry::Class)) }
+    def singleton_klass
+      attached_class = @owner_stack.last
+      return unless attached_class
+
+      # Return the existing singleton class if available
+      owner = T.cast(
+        @index["#{attached_class.name}::<Class:#{attached_class.name}>"],
+        T.nilable(T::Array[Entry::SingletonClass]),
+      )
+      return owner.first if owner
+
+      # If not available, create the singleton class lazily
+      nesting = @stack + ["<Class:#{@stack.last}>"]
+      entry = Entry::SingletonClass.new(nesting, @file_path, attached_class.location, [], nil)
+      @index << entry
+      entry
+    end
   end
 end

lib/ruby_indexer/lib/ruby_indexer/entry.rb

@@ -162,6 +162,21 @@ def ancestor_hash
       end
     end
 
+    class SingletonClass < Class
+      extend T::Sig
+
+      sig { params(location: Prism::Location, comments: T::Array[String]).void }
+      def update_singleton_information(location, comments)
+        @location = Location.new(
+          location.start_line,
+          location.end_line,
+          location.start_column,
+          location.end_column,
+        )
+        @comments.concat(comments)
+      end
+    end
+
     class Constant < Entry
     end
 
@@ -280,9 +295,6 @@ def parameters
 
     class Method < Member
       extend T::Sig
-      extend T::Helpers
-
-      abstract!
 
       sig { override.returns(T::Array[Parameter]) }
       attr_reader :parameters
@@ -391,12 +403,6 @@ def parameter_name(node)
       end
     end
 
-    class SingletonMethod < Method
-    end
-
-    class InstanceMethod < Method
-    end
-
     # An UnresolvedAlias points to a constant alias with a right hand side that has not yet been resolved. For
     # example, if we find
     #

lib/ruby_indexer/test/classes_and_modules_test.rb

@@ -470,5 +470,51 @@ class ConstantPathReferences
       constant_path_references = T.must(@index["ConstantPathReferences"][0])
       assert_equal(["Foo::Bar", "Foo::Bar2"], constant_path_references.mixin_operation_module_names)
     end
+
+    def test_tracking_singleton_classes
+      index(<<~RUBY)
+        class Foo; end
+        class Foo
+          # Some extra comments
+          class << self
+          end
+        end
+      RUBY
+
+      foo = T.must(@index["Foo::<Class:Foo>"].first)
+      assert_equal(4, foo.location.start_line)
+      assert_equal("Some extra comments", foo.comments.join("\n"))
+    end
+
+    def test_dynamic_singleton_class_blocks
+      index(<<~RUBY)
+        class Foo
+          # Some extra comments
+          class << bar
+          end
+        end
+      RUBY
+
+      singleton = T.must(@index["Foo::<Class:Foo>"].first)
+
+      # Even though this is not correct, we consider any dynamic singleton class block as a regular singleton class.
+      # That pattern cannot be properly analyzed statically and assuming that it's always a regular singleton simplifies
+      # the implementation considerably.
+      assert_equal(3, singleton.location.start_line)
+      assert_equal("Some extra comments", singleton.comments.join("\n"))
+    end
+
+    def test_namespaces_inside_singleton_blocks
+      index(<<~RUBY)
+        class Foo
+          class << self
+            class Bar
+            end
+          end
+        end
+      RUBY
+
+      assert_entry("Foo::<Class:Foo>::Bar", Entry::Class, "/fake/path/foo.rb:2-4:3-7")
+    end
   end
 end