Add ability to linearize ancestors

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -65,8 +65,10 @@ def on_class_node_enter(node)
         superclass.slice
       end
 
+      nesting = name.start_with?("::") ? [name.delete_prefix("::")] : @stack + [name.delete_prefix("::")]
+
       entry = Entry::Class.new(
-        fully_qualify_name(name),
+        nesting,
         @file_path,
         node.location,
         comments,
@@ -89,7 +91,9 @@ def on_module_node_enter(node)
       name = node.constant_path.location.slice
 
       comments = collect_comments(node)
-      entry = Entry::Module.new(fully_qualify_name(name), @file_path, node.location, comments)
+
+      nesting = name.start_with?("::") ? [name.delete_prefix("::")] : @stack + [name.delete_prefix("::")]
+      entry = Entry::Module.new(nesting, @file_path, node.location, comments)
 
       @owner_stack << entry
       @index << entry
@@ -199,10 +203,8 @@ def on_call_node_enter(node)
         handle_attribute(node, reader: false, writer: true)
       when :attr_accessor
         handle_attribute(node, reader: true, writer: true)
-      when :include
-        handle_module_operation(node, :included_modules)
-      when :prepend
-        handle_module_operation(node, :prepended_modules)
+      when :include, :prepend, :extend
+        handle_module_operation(node, message)
       end
     end
 
@@ -447,14 +449,14 @@ def handle_module_operation(node, operation)
 
       names = arguments.filter_map do |node|
         if node.is_a?(Prism::ConstantReadNode) || node.is_a?(Prism::ConstantPathNode)
-          node.full_name
+          [operation, node.full_name]
         end
       rescue Prism::ConstantPathNode::DynamicPartsInConstantPathError,
              Prism::ConstantPathNode::MissingNodesInConstantPathError
         # Do nothing
       end
-      collection = operation == :included_modules ? owner.included_modules : owner.prepended_modules
-      collection.concat(names)
+
+      owner.modules.concat(names)
     end
   end
 end

lib/ruby_indexer/lib/ruby_indexer/entry.rb

@@ -61,13 +61,30 @@ class Namespace < Entry
       abstract!
 
       sig { returns(T::Array[String]) }
-      def included_modules
-        @included_modules ||= T.let([], T.nilable(T::Array[String]))
+      attr_reader :nesting
+
+      sig do
+        params(
+          nesting: T::Array[String],
+          file_path: String,
+          location: T.any(Prism::Location, RubyIndexer::Location),
+          comments: T::Array[String],
+        ).void
       end
+      def initialize(nesting, file_path, location, comments)
+        @name = T.let(nesting.join("::"), String)
+        # The original nesting where this namespace was discovered
+        @nesting = nesting
 
-      sig { returns(T::Array[String]) }
-      def prepended_modules
-        @prepended_modules ||= T.let([], T.nilable(T::Array[String]))
+        super(@name, file_path, location, comments)
+      end
+
+      # Stores all prepend, include and extend operations in the exact order they were discovered in the source code.
+      # Maintaining the order is essential to linearize ancestors the right way when a module is both included and
+      # prepended
+      sig { returns(T::Array[[Symbol, String]]) }
+      def modules
+        @modules ||= T.let([], T.nilable(T::Array[[Symbol, String]]))
       end
     end
 
@@ -84,15 +101,16 @@ class Class < Namespace
 
       sig do
         params(
-          name: String,
+          nesting: T::Array[String],
           file_path: String,
           location: T.any(Prism::Location, RubyIndexer::Location),
           comments: T::Array[String],
           parent_class: T.nilable(String),
         ).void
       end
-      def initialize(name, file_path, location, comments, parent_class)
-        super(name, file_path, location, comments)
+      def initialize(nesting, file_path, location, comments, parent_class)
+        super(nesting, file_path, location, comments)
+
         @parent_class = T.let(parent_class, T.nilable(String))
       end
     end

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -6,6 +6,7 @@ class Index
     extend T::Sig
 
     class UnresolvableAliasError < StandardError; end
+    class NonExistingNamespaceError < StandardError; end
 
     # The minimum Jaro-Winkler similarity score for an entry to be considered a match for a given fuzzy search query
     ENTRY_SIMILARITY_THRESHOLD = 0.7
@@ -31,6 +32,9 @@ def initialize
 
       # Holds all require paths for every indexed item so that we can provide autocomplete for requires
       @require_paths_tree = T.let(PrefixTree[IndexablePath].new, PrefixTree[IndexablePath])
+
+      # Holds the linearized ancestors list for every namespace
+      @ancestors = T.let({}, T::Hash[String, T::Array[String]])
     end
 
     sig { params(indexable: IndexablePath).void }
@@ -255,6 +259,100 @@ def resolve_method(method_name, receiver_name)
       end
     end
 
+    # Linearizes the ancestors for a given name, returning the order of namespaces in which Ruby will search for method
+    # or constant declarations.
+    #
+    # When we add an ancestor in Ruby, that namespace might have ancestors of its own. Therefore, we need to linearize
+    # everything recursively to ensure that we are placing ancestors in the right order. For example, if you include a
+    # module that prepends another module, then the prepend module appears before the included module.
+    #
+    # The order of ancestors is [linearized_prepends, self, linearized_includes, linearized_superclass]
+    sig { params(name: String).returns(T::Array[String]) }
+    def linearized_ancestors_of(name)
+      # If we already computed the ancestors for this namespace, return it straight away
+      cached_ancestors = @ancestors[name]
+      return cached_ancestors if cached_ancestors
+
+      ancestors = [name]
+
+      # Cache the linearized ancestors array eagerly. This is important because we might have circular dependencies and
+      # this will prevent us from falling into an infinite recursion loop. Because we mutate the ancestors array later,
+      # the cache will reflect the final result
+      @ancestors[name] = ancestors
+
+      # If we don't have an entry for `name`, raise
+      entries = self[name]
+      raise NonExistingNamespaceError, "No entry found for #{name}" unless entries
+
+      # If none of the entries for `name` are namespaces, return an empty array
+      namespaces = T.cast(entries.select { |e| e.is_a?(Entry::Namespace) }, T::Array[Entry::Namespace])
+      raise NonExistingNamespaceError, "None of the entries for #{name} are modules or classes" if namespaces.empty?
+
+      modules = namespaces.flat_map(&:modules)
+      prepended_modules_count = 0
+      included_modules_count = 0
+
+      # The original nesting where we discovered this namespace, so that we resolve the correct names of the
+      # included/prepended/extended modules and parent classes
+      nesting = T.must(namespaces.first).nesting
+
+      modules.each do |operation, module_name|
+        resolved_module = resolve(module_name, nesting)
+        next unless resolved_module
+
+        fully_qualified_name = T.must(resolved_module.first).name
+
+        case operation
+        when :prepend
+          # When a module is prepended, Ruby checks if it hasn't been prepended already to prevent adding it in front of
+          # the actual namespace twice. However, it does not check if it has been included because you are allowed to
+          # prepend the same module after it has already been included
+          linearized_prepends = linearized_ancestors_of(fully_qualified_name)
+
+          # When there are duplicate prepended modules, we have to insert the new prepends after the existing ones. For
+          # example, if the current ancestors are `["A", "Foo"]` and we try to prepend `["A", "B"]`, then `"B"` has to
+          # be inserted after `"A`
+          uniq_prepends = linearized_prepends - T.must(ancestors[0...prepended_modules_count])
+          insert_position = linearized_prepends.length - uniq_prepends.length
+
+          T.unsafe(ancestors).insert(
+            insert_position,
+            *(linearized_prepends - T.must(ancestors[0...prepended_modules_count])),
+          )
+
+          prepended_modules_count += linearized_prepends.length
+        when :include
+          # When including a module, Ruby will always prevent duplicate entries in case the module has already been
+          # prepended or included
+          linearized_includes = linearized_ancestors_of(fully_qualified_name)
+
+          T.unsafe(ancestors).insert(
+            ancestors.length - included_modules_count,
+            *(linearized_includes - ancestors),
+          )
+
+          included_modules_count += linearized_includes.length
+        end
+      end
+
+      # 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))
+
+      if superclass
+        # If the user makes a mistake and creates a class that inherits from itself, this method would throw a stack
+        # error. We need to ensure that this isn't the case
+        parent_class = T.must(superclass.parent_class)
+
+        resolved_parent_class = resolve(parent_class, nesting)
+        parent_class_name = resolved_parent_class&.first&.name
+
+        ancestors.concat(linearized_ancestors_of(parent_class_name)) if parent_class_name && name != parent_class_name
+      end
+
+      ancestors
+    end
+
     private
 
     # Attempts to resolve an UnresolvedAlias into a resolved Alias. If the unresolved alias is pointing to a constant

lib/ruby_indexer/test/classes_and_modules_test.rb

@@ -369,13 +369,13 @@ class ConstantPathReferences
       RUBY
 
       foo = T.must(@index["Foo"][0])
-      assert_equal(["A1", "A2", "A3", "A4", "A5", "A6"], foo.included_modules)
+      assert_equal(["A1", "A2", "A3", "A4", "A5", "A6"], foo.modules.flat_map(&:last))
 
       qux = T.must(@index["Foo::Qux"][0])
-      assert_equal(["Corge", "Corge", "Baz"], qux.included_modules)
+      assert_equal(["Corge", "Corge", "Baz"], qux.modules.flat_map(&:last))
 
       constant_path_references = T.must(@index["ConstantPathReferences"][0])
-      assert_equal(["Foo::Bar", "Foo::Bar2"], constant_path_references.included_modules)
+      assert_equal(["Foo::Bar", "Foo::Bar2"], constant_path_references.modules.flat_map(&:last))
     end
 
     def test_keeping_track_of_prepended_modules
@@ -415,13 +415,59 @@ class ConstantPathReferences
       RUBY
 
       foo = T.must(@index["Foo"][0])
-      assert_equal(["A1", "A2", "A3", "A4", "A5", "A6"], foo.prepended_modules)
+      assert_equal(["A1", "A2", "A3", "A4", "A5", "A6"], foo.modules.flat_map(&:last))
 
       qux = T.must(@index["Foo::Qux"][0])
-      assert_equal(["Corge", "Corge", "Baz"], qux.prepended_modules)
+      assert_equal(["Corge", "Corge", "Baz"], qux.modules.flat_map(&:last))
 
       constant_path_references = T.must(@index["ConstantPathReferences"][0])
-      assert_equal(["Foo::Bar", "Foo::Bar2"], constant_path_references.prepended_modules)
+      assert_equal(["Foo::Bar", "Foo::Bar2"], constant_path_references.modules.flat_map(&:last))
+    end
+
+    def test_keeping_track_of_extended_modules
+      index(<<~RUBY)
+        class Foo
+          # valid syntaxes that we can index
+          extend A1
+          self.extend A2
+          extend A3, A4
+          self.extend A5, A6
+
+          # valid syntaxes that we cannot index because of their dynamic nature
+          extend some_variable_or_method_call
+          self.extend some_variable_or_method_call
+
+          def something
+            extend A7 # We should not index this because of this dynamic nature
+          end
+
+          # Valid inner class syntax definition with its own modules prepended
+          class Qux
+            extend Corge
+            self.extend Corge
+            extend Baz
+
+            extend some_variable_or_method_call
+          end
+        end
+
+        class ConstantPathReferences
+          extend Foo::Bar
+          self.extend Foo::Bar2
+
+          extend dynamic::Bar
+          extend Foo::
+        end
+      RUBY
+
+      foo = T.must(@index["Foo"][0])
+      assert_equal(["A1", "A2", "A3", "A4", "A5", "A6"], foo.modules.flat_map(&:last))
+
+      qux = T.must(@index["Foo::Qux"][0])
+      assert_equal(["Corge", "Corge", "Baz"], qux.modules.flat_map(&:last))
+
+      constant_path_references = T.must(@index["ConstantPathReferences"][0])
+      assert_equal(["Foo::Bar", "Foo::Bar2"], constant_path_references.modules.flat_map(&:last))
     end
   end
 end