Linearize singleton ancestors

Co-authored-by: Alexandre Terrasa <Morriar@users.noreply.github.com>

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -314,20 +314,22 @@ def on_def_node_enter(node)
           @owner_stack.last,
         ))
       when Prism::SelfNode
-        singleton = singleton_klass
+        owner = @owner_stack.last
 
-        @index.add(Entry::Method.new(
-          method_name,
-          @file_path,
-          node.location,
-          node.name_loc,
-          comments,
-          list_params(node.parameters),
-          current_visibility,
-          singleton,
-        ))
+        if owner
+          singleton = @index.existing_or_new_singleton_class(owner.name)
+
+          @index.add(Entry::Method.new(
+            method_name,
+            @file_path,
+            node.location,
+            node.name_loc,
+            comments,
+            list_params(node.parameters),
+            current_visibility,
+            singleton,
+          ))
 
-        if singleton
           @owner_stack << singleton
           @stack << "<Class:#{@stack.last}>"
         end
@@ -405,7 +407,12 @@ def handle_instance_variable(node, loc)
 
       # 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
+      owner = @owner_stack.last
+
+      if owner && !@inside_def
+        owner = @index.existing_or_new_singleton_class(owner.name)
+      end
+
       @index.add(Entry::InstanceVariable.new(name, @file_path, loc, collect_comments(node), owner))
     end
 
@@ -605,7 +612,8 @@ def handle_module_operation(node, operation)
         when :prepend
           owner.mixin_operations << Entry::Prepend.new(node.full_name)
         when :extend
-          owner.mixin_operations << Entry::Extend.new(node.full_name)
+          singleton = @index.existing_or_new_singleton_class(owner.name)
+          singleton.mixin_operations << Entry::Include.new(node.full_name)
         end
       rescue Prism::ConstantPathNode::DynamicPartsInConstantPathError,
              Prism::ConstantPathNode::MissingNodesInConstantPathError
@@ -701,31 +709,5 @@ def parameter_name(node)
         :"(#{names_with_commas})"
       end
     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,
-        attached_class.name_location,
-        [],
-        nil,
-      )
-      @index.add(entry, skip_prefix_tree: true)
-      entry
-    end
   end
 end

lib/ruby_indexer/lib/ruby_indexer/entry.rb

@@ -86,7 +86,6 @@ def initialize(module_name)
 
     class Include < ModuleOperation; end
     class Prepend < ModuleOperation; end
-    class Extend < ModuleOperation; end
 
     class Namespace < Entry
       extend T::Sig

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -374,8 +374,25 @@ def linearized_ancestors_of(fully_qualified_name)
       cached_ancestors = @ancestors[fully_qualified_name]
       return cached_ancestors if cached_ancestors
 
+      parts = fully_qualified_name.split("::")
+      singleton_levels = 0
+
+      parts.reverse_each do |part|
+        break unless part.include?("<Class:")
+
+        singleton_levels += 1
+        parts.pop
+      end
+
+      attached_class_name = parts.join("::")
+
       # If we don't have an entry for `name`, raise
       entries = self[fully_qualified_name]
+
+      if singleton_levels > 0 && !entries
+        entries = [existing_or_new_singleton_class(attached_class_name)]
+      end
+
       raise NonExistingNamespaceError, "No entry found for #{fully_qualified_name}" unless entries
 
       ancestors = [fully_qualified_name]
@@ -405,6 +422,12 @@ def linearized_ancestors_of(fully_qualified_name)
       # included/prepended/extended modules and parent classes
       nesting = T.must(namespaces.first).nesting
 
+      if nesting.any?
+        singleton_levels.times do
+          nesting << "<Class:#{T.must(nesting.last)}>"
+        end
+      end
+
       mixin_operations.each do |operation|
         resolved_module = resolve(operation.module_name, nesting)
         next unless resolved_module
@@ -440,7 +463,14 @@ def linearized_ancestors_of(fully_qualified_name)
 
       # Find the first class entry that has a parent class. Notice that if the developer makes a mistake and inherits
       # from two diffent classes in different files, we simply ignore it
-      superclass = T.cast(namespaces.find { |n| n.is_a?(Entry::Class) && n.parent_class }, T.nilable(Entry::Class))
+      superclass = T.cast(
+        if singleton_levels > 0
+          self[attached_class_name]&.find { |n| n.is_a?(Entry::Class) && n.parent_class }
+        else
+          namespaces.find { |n| n.is_a?(Entry::Class) && n.parent_class }
+        end,
+        T.nilable(Entry::Class),
+      )
 
       if superclass
         # If the user makes a mistake and creates a class that inherits from itself, this method would throw a stack
@@ -451,7 +481,39 @@ def linearized_ancestors_of(fully_qualified_name)
         parent_class_name = resolved_parent_class&.first&.name
 
         if parent_class_name && fully_qualified_name != parent_class_name
-          ancestors.concat(linearized_ancestors_of(parent_class_name))
+
+          parent_name_parts = [parent_class_name]
+          singleton_levels.times do
+            parent_name_parts << "<Class:#{parent_name_parts.last}>"
+          end
+
+          ancestors.concat(linearized_ancestors_of(parent_name_parts.join("::")))
+        end
+
+        # When computing the linearization for a class's singleton class, it inherits from the linearized ancestors of
+        # the `Class` class
+        if parent_class_name&.start_with?("BasicObject") && singleton_levels > 0
+          class_class_name_parts = ["Class"]
+
+          (singleton_levels - 1).times do
+            class_class_name_parts << "<Class:#{class_class_name_parts.last}>"
+          end
+
+          ancestors.concat(linearized_ancestors_of(class_class_name_parts.join("::")))
+        end
+      elsif singleton_levels > 0
+        # When computing the linearization for a module's singleton class, it inherits from the linearized ancestors of
+        # the `Module` class
+        mod = T.cast(self[attached_class_name]&.find { |n| n.is_a?(Entry::Module) }, T.nilable(Entry::Module))
+
+        if mod
+          module_class_name_parts = ["Module"]
+
+          (singleton_levels - 1).times do
+            module_class_name_parts << "<Class:#{module_class_name_parts.last}>"
+          end
+
+          ancestors.concat(linearized_ancestors_of(module_class_name_parts.join("::")))
         end
       end
 
@@ -513,6 +575,29 @@ def handle_change(indexable)
       @ancestors.clear if original_map.any? { |name, hash| updated_map[name] != hash }
     end
 
+    sig { params(name: String).returns(Entry::SingletonClass) }
+    def existing_or_new_singleton_class(name)
+      *_namespace, unqualified_name = name.split("::")
+      full_singleton_name = "#{name}::<Class:#{unqualified_name}>"
+      singleton = T.cast(self[full_singleton_name]&.first, T.nilable(Entry::SingletonClass))
+
+      unless singleton
+        attached_ancestor = T.must(self[name]&.first)
+
+        singleton = Entry::SingletonClass.new(
+          [full_singleton_name],
+          attached_ancestor.file_path,
+          attached_ancestor.location,
+          attached_ancestor.name_location,
+          [],
+          nil,
+        )
+        add(singleton, skip_prefix_tree: true)
+      end
+
+      singleton
+    end
+
     private
 
     # Attempts to resolve an UnresolvedAlias into a resolved Alias. If the unresolved alias is pointing to a constant

lib/ruby_indexer/lib/ruby_indexer/rbs_indexer.rb

@@ -93,16 +93,15 @@ def to_ruby_indexer_location(rbs_location)
     def add_declaration_mixins_to_entry(declaration, entry)
       declaration.each_mixin do |mixin|
         name = mixin.name.name.to_s
-        mixin_operation =
-          case mixin
-          when RBS::AST::Members::Include
-            Entry::Include.new(name)
-          when RBS::AST::Members::Extend
-            Entry::Extend.new(name)
-          when RBS::AST::Members::Prepend
-            Entry::Prepend.new(name)
-          end
-        entry.mixin_operations << mixin_operation if mixin_operation
+        case mixin
+        when RBS::AST::Members::Include
+          entry.mixin_operations << Entry::Include.new(name)
+        when RBS::AST::Members::Prepend
+          entry.mixin_operations << Entry::Prepend.new(name)
+        when RBS::AST::Members::Extend
+          singleton = @index.existing_or_new_singleton_class(entry.name)
+          singleton.mixin_operations << Entry::Include.new(name)
+        end
       end
     end
 
@@ -122,26 +121,8 @@ def handle_method(member, owner)
         Entry::Visibility::PUBLIC
       end
 
-      real_owner = member.singleton? ? existing_or_new_singleton_klass(owner) : owner
+      real_owner = member.singleton? ? @index.existing_or_new_singleton_class(owner.name) : owner
       @index.add(Entry::Method.new(name, file_path, location, location, comments, [], visibility, real_owner))
     end
-
-    sig { params(owner: Entry::Namespace).returns(T.nilable(Entry::Class)) }
-    def existing_or_new_singleton_klass(owner)
-      *_parts, name = owner.name.split("::")
-
-      # Return the existing singleton class if available
-      singleton_entries = T.cast(
-        @index["#{owner.name}::<Class:#{name}>"],
-        T.nilable(T::Array[Entry::SingletonClass]),
-      )
-      return singleton_entries.first if singleton_entries
-
-      # If not available, create the singleton class lazily
-      nesting = owner.nesting + ["<Class:#{name}>"]
-      entry = Entry::SingletonClass.new(nesting, owner.file_path, owner.location, owner.name_location, [], nil)
-      @index.add(entry, skip_prefix_tree: true)
-      entry
-    end
   end
 end

lib/ruby_indexer/test/classes_and_modules_test.rb

@@ -461,13 +461,13 @@ class ConstantPathReferences
         end
       RUBY
 
-      foo = T.must(@index["Foo"][0])
+      foo = T.must(@index["Foo::<Class:Foo>"][0])
       assert_equal(["A1", "A2", "A3", "A4", "A5", "A6"], foo.mixin_operation_module_names)
 
-      qux = T.must(@index["Foo::Qux"][0])
+      qux = T.must(@index["Foo::Qux::<Class:Qux>"][0])
       assert_equal(["Corge", "Corge", "Baz"], qux.mixin_operation_module_names)
 
-      constant_path_references = T.must(@index["ConstantPathReferences"][0])
+      constant_path_references = T.must(@index["ConstantPathReferences::<Class:ConstantPathReferences>"][0])
       assert_equal(["Foo::Bar", "Foo::Bar2"], constant_path_references.mixin_operation_module_names)
     end