Allow for indexing augmentations and included hooks

lib/ruby_indexer/lib/ruby_indexer/augmentation.rb

@@ -0,0 +1,24 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyIndexer
+  module Augmentation
+    extend T::Sig
+    extend T::Helpers
+
+    interface!
+
+    # The `on_extend` indexing augmentation is invoked whenever an extend is encountered in the code. It can be used to
+    # register for an included callback, similar to what `ActiveSupport::Concern` does in order to auto-extend the
+    # `ClassMethods` modules
+    sig do
+      abstract.params(
+        index: Index,
+        owner: T.nilable(Entry::Namespace),
+        node: Prism::CallNode,
+        file_path: String,
+      ).void
+    end
+    def on_call_node(index, owner, node, file_path); end
+  end
+end

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -9,11 +9,18 @@ class DeclarationListener
     BASIC_OBJECT_NESTING = T.let(["BasicObject"].freeze, T::Array[String])
 
     sig do
-      params(index: Index, dispatcher: Prism::Dispatcher, parse_result: Prism::ParseResult, file_path: String).void
+      params(
+        index: Index,
+        dispatcher: Prism::Dispatcher,
+        parse_result: Prism::ParseResult,
+        file_path: String,
+        augmentations: T::Array[Augmentation],
+      ).void
     end
-    def initialize(index, dispatcher, parse_result, file_path)
+    def initialize(index, dispatcher, parse_result, file_path, augmentations = [])
       @index = index
       @file_path = file_path
+      @augmentations = augmentations
       @visibility_stack = T.let([Entry::Visibility::PUBLIC], T::Array[Entry::Visibility])
       @comments_by_line = T.let(
         parse_result.comments.to_h do |c|
@@ -279,6 +286,8 @@ def on_call_node_enter(node)
       when :private
         @visibility_stack.push(Entry::Visibility::PRIVATE)
       end
+
+      @augmentations.each { |aug| aug.on_call_node(@index, @owner_stack.last, node, @file_path) }
     end
 
     sig { params(node: Prism::CallNode).void }

lib/ruby_indexer/lib/ruby_indexer/index.rb

@@ -35,6 +35,27 @@ def initialize
 
       # Holds the linearized ancestors list for every namespace
       @ancestors = T.let({}, T::Hash[String, T::Array[String]])
+
+      # List of classes that are augmenting the index
+      @augmentations = T.let([], T::Array[Augmentation])
+
+      # Map of module name to included hooks that have to be executed when we include the given module
+      @included_hooks = T.let(
+        {},
+        T::Hash[String, T::Array[T.proc.params(index: Index, base: Entry::Namespace).void]],
+      )
+    end
+
+    # Register an indexing augmentation to the index. Augmentations must conform to the `Augmentation` interface
+    sig { params(augmentation: Augmentation).void }
+    def register_augmentation(augmentation)
+      @augmentations << augmentation
+    end
+
+    # Register an included `hook` that will be executed when `module_name` is included into any namespace
+    sig { params(module_name: String, hook: T.proc.params(index: Index, base: Entry::Namespace).void).void }
+    def register_included_hook(module_name, &hook)
+      (@included_hooks[module_name] ||= []) << hook
     end
 
     sig { params(indexable: IndexablePath).void }
@@ -296,7 +317,7 @@ def index_single(indexable_path, source = nil)
       dispatcher = Prism::Dispatcher.new
 
       result = Prism.parse(content)
-      DeclarationListener.new(self, dispatcher, result, indexable_path.full_path)
+      DeclarationListener.new(self, dispatcher, result, indexable_path.full_path, @augmentations)
       dispatcher.dispatch(result.value)
 
       require_path = indexable_path.require_path
@@ -457,6 +478,12 @@ def linearized_ancestors_of(fully_qualified_name)
         end
       end
 
+      # We only need to run included hooks when linearizing singleton classes. Included hooks are typically used to add
+      # new singleton methods or to extend a module through an include. There's no need to support instance methods, the
+      # inclusion of another module or the prepending of another module, because those features are already a part of
+      # Ruby and can be used directly without any metaprogramming
+      run_included_hooks(attached_class_name, nesting) if singleton_levels > 0
+
       linearize_mixins(ancestors, namespaces, nesting)
       linearize_superclass(
         ancestors,
@@ -570,6 +597,34 @@ def existing_or_new_singleton_class(name)
 
     private
 
+    # Runs the registered included hooks
+    sig { params(fully_qualified_name: String, nesting: T::Array[String]).void }
+    def run_included_hooks(fully_qualified_name, nesting)
+      return if @included_hooks.empty?
+
+      namespaces = self[fully_qualified_name]&.grep(Entry::Namespace)
+      return unless namespaces
+
+      namespaces.each do |namespace|
+        namespace.mixin_operations.each do |operation|
+          next unless operation.is_a?(Entry::Include)
+
+          # First we resolve the include name, so that we know the actual module being referred to in the include
+          resolved_modules = resolve(operation.module_name, nesting)
+          next unless resolved_modules
+
+          module_name = T.must(resolved_modules.first).name
+
+          # Then we grab any hooks registered for that module
+          hooks = @included_hooks[module_name]
+          next unless hooks
+
+          # We invoke the hooks with the index and the namespace that included the module
+          hooks.each { |hook| hook.call(self, namespace) }
+        end
+      end
+    end
+
     # Linearize mixins for an array of namespace entries. This method will mutate the `ancestors` array with the
     # linearized ancestors of the mixins
     sig do

lib/ruby_indexer/ruby_indexer.rb

@@ -6,6 +6,7 @@
 
 require "ruby_indexer/lib/ruby_indexer/indexable_path"
 require "ruby_indexer/lib/ruby_indexer/declaration_listener"
+require "ruby_indexer/lib/ruby_indexer/augmentation"
 require "ruby_indexer/lib/ruby_indexer/index"
 require "ruby_indexer/lib/ruby_indexer/entry"
 require "ruby_indexer/lib/ruby_indexer/configuration"

lib/ruby_indexer/test/augmentations_test.rb

@@ -0,0 +1,163 @@
+# typed: true
+# frozen_string_literal: true
+
+require_relative "test_case"
+
+module RubyIndexer
+  class AugmentationsTest < TestCase
+    def test_augmenting_indexing_included_hook
+      aug_class = Class.new do
+        include Augmentation
+
+        def on_call_node(index, owner, node, file_path)
+          return unless owner
+          return unless node.name == :extend
+
+          arguments = node.arguments&.arguments
+          return unless arguments
+
+          location = node.location
+
+          arguments.each do |node|
+            next unless node.is_a?(Prism::ConstantReadNode) || node.is_a?(Prism::ConstantPathNode)
+
+            module_name = node.full_name
+            next unless module_name == "ActiveSupport::Concern"
+
+            index.register_included_hook(owner.name) do |index, base|
+              class_methods_name = "#{owner.name}::ClassMethods"
+
+              if index.indexed?(class_methods_name)
+                singleton = index.existing_or_new_singleton_class(base.name)
+                singleton.mixin_operations << Entry::Include.new(class_methods_name)
+              end
+            end
+
+            index.add(Entry::Method.new(
+              "new_method",
+              file_path,
+              location,
+              location,
+              [],
+              [Entry::Signature.new([Entry::RequiredParameter.new(name: :a)])],
+              Entry::Visibility::PUBLIC,
+              owner,
+            ))
+          rescue Prism::ConstantPathNode::DynamicPartsInConstantPathError,
+                 Prism::ConstantPathNode::MissingNodesInConstantPathError
+            # Do nothing
+          end
+        end
+      end
+
+      @index.register_augmentation(aug_class.new)
+      index(<<~RUBY)
+        module ActiveSupport
+          module Concern
+            def self.extended(base)
+              base.class_eval("def new_method(a); end")
+            end
+          end
+        end
+
+        module ActiveRecord
+          module Associations
+            extend ActiveSupport::Concern
+
+            module ClassMethods
+              def belongs_to(something); end
+            end
+          end
+
+          class Base
+            include Associations
+          end
+        end
+
+        class User < ActiveRecord::Base
+        end
+      RUBY
+
+      assert_equal(
+        [
+          "User::<Class:User>",
+          "ActiveRecord::Base::<Class:Base>",
+          "ActiveRecord::Associations::ClassMethods",
+          "Object::<Class:Object>",
+          "BasicObject::<Class:BasicObject>",
+          "Class",
+          "Module",
+          "Object",
+          "Kernel",
+          "BasicObject",
+        ],
+        @index.linearized_ancestors_of("User::<Class:User>"),
+      )
+
+      assert_entry("new_method", Entry::Method, "/fake/path/foo.rb:10-4:10-33")
+    end
+
+    def test_augmenting_indexing_configuration_dsl
+      aug_class = Class.new do
+        include Augmentation
+
+        def on_call_node(index, owner, node, file_path)
+          return unless owner
+
+          name = node.name
+          return unless name == :has_many
+
+          arguments = node.arguments&.arguments
+          return unless arguments
+
+          association_name = arguments.first
+          return unless association_name.is_a?(Prism::SymbolNode)
+
+          location = association_name.location
+
+          index.add(Entry::Method.new(
+            T.must(association_name.value),
+            file_path,
+            location,
+            location,
+            [],
+            [],
+            Entry::Visibility::PUBLIC,
+            owner,
+          ))
+        end
+      end
+
+      @index.register_augmentation(aug_class.new)
+      index(<<~RUBY)
+        module ActiveSupport
+          module Concern
+            def self.extended(base)
+              base.class_eval("def new_method(a); end")
+            end
+          end
+        end
+
+        module ActiveRecord
+          module Associations
+            extend ActiveSupport::Concern
+
+            module ClassMethods
+              def belongs_to(something); end
+            end
+          end
+
+          class Base
+            include Associations
+          end
+        end
+
+        class User < ActiveRecord::Base
+          has_many :posts
+        end
+      RUBY
+
+      assert_entry("posts", Entry::Method, "/fake/path/foo.rb:23-11:23-17")
+    end
+  end
+end