Add ability to linearize ancestors

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -194,10 +194,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
 
@@ -371,15 +369,15 @@ 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
         # TO DO: add MissingNodesInConstantPathError when released in Prism
         # If a constant path reference is dynamic or missing parts, we can't
         # index it
       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

@@ -60,14 +60,12 @@ class Namespace < Entry
 
       abstract!
 
-      sig { returns(T::Array[String]) }
-      def included_modules
-        @included_modules ||= T.let([], T.nilable(T::Array[String]))
-      end
-
-      sig { returns(T::Array[String]) }
-      def prepended_modules
-        @prepended_modules ||= T.let([], T.nilable(T::Array[String]))
+      # 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
 

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -31,6 +31,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 ancestor list for every namespace
+      @ancestors = T.let({}, T::Hash[String, T::Array[String]])
     end
 
     sig { params(indexable: IndexablePath).void }
@@ -257,6 +260,85 @@ def resolve_method(method_name, receiver_name)
       )
     end
 
+    # Linearizes the ancestors for a given name, returning the order of namespace 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]
+
+      # If we don't have an entry for `name` return an empty array
+      entries = self[name]
+      return ancestors 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])
+      return ancestors if namespaces.empty?
+
+      modules = namespaces.flat_map(&:modules)
+      prepended_modules_count = 0
+      included_modules_count = 0
+
+      modules.each do |operation, module_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(module_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(module_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)
+        ancestors.concat(linearized_ancestors_of(parent_class)) if name != parent_class
+      end
+
+      # Cache the linearized ancestors so that we don't have to compute it again unless the user changes the ancestors
+      # of a namespace
+      @ancestors[name] = ancestors
+      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