Handle concern inheritance in indexing enhancement

Co-authored-by: Andy Waite <[email protected]>

lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb

@@ -6,11 +6,14 @@ module Rails
     class IndexingEnhancement < RubyIndexer::Enhancement
       extend T::Sig
 
-      sig do
-        override.params(
-          call_node: Prism::CallNode,
-        ).void
+      sig { params(listener: RubyIndexer::DeclarationListener).void }
+      def initialize(listener)
+        super
+
+        @discovered_concerns = T.let([], T::Array[String])
       end
+
+      sig { override.params(call_node: Prism::CallNode).void }
       def on_call_node_enter(call_node)
         owner = @listener.current_owner
         return unless owner
@@ -26,11 +29,7 @@ def on_call_node_enter(call_node)
         end
       end
 
-      sig do
-        override.params(
-          call_node: Prism::CallNode,
-        ).void
-      end
+      sig { override.params(call_node: Prism::CallNode).void }
       def on_call_node_leave(call_node)
         if call_node.name == :class_methods && call_node.block
           @listener.pop_namespace_stack
@@ -39,12 +38,7 @@ def on_call_node_leave(call_node)
 
       private
 
-      sig do
-        params(
-          owner: RubyIndexer::Entry::Namespace,
-          call_node: Prism::CallNode,
-        ).void
-      end
+      sig { params(owner: RubyIndexer::Entry::Namespace, call_node: Prism::CallNode).void }
       def handle_association(owner, call_node)
         arguments = call_node.arguments&.arguments
         return unless arguments
@@ -84,13 +78,34 @@ def handle_concern_extend(owner, call_node)
           module_name = node.full_name
           next unless module_name == "ActiveSupport::Concern"
 
+          @discovered_concerns << owner.name
+
           @listener.register_included_hook do |index, base|
             class_methods_name = "#{owner.name}::ClassMethods"
 
+            singleton = index.existing_or_new_singleton_class(base.name)
+
             if index.indexed?(class_methods_name)
-              singleton = index.existing_or_new_singleton_class(base.name)
               singleton.mixin_operations << RubyIndexer::Entry::Include.new(class_methods_name)
             end
+
+            if @discovered_concerns.include?(owner.name)
+              owner.mixin_operations.each do |operation|
+                resolved_module = index.resolve(operation.module_name, base.nesting)
+                next unless resolved_module
+
+                name = T.must(resolved_module.first).name
+                module_name = "#{name}::ClassMethods"
+                next unless @discovered_concerns.include?(name) && index.indexed?(module_name)
+
+                case operation
+                when RubyIndexer::Entry::Include
+                  singleton.mixin_operations << RubyIndexer::Entry::Include.new(module_name)
+                when RubyIndexer::Entry::Prepend
+                  singleton.mixin_operations.unshift(RubyIndexer::Entry::Include.new(module_name))
+                end
+              end
+            end
           end
         rescue Prism::ConstantPathNode::DynamicPartsInConstantPathError,
                Prism::ConstantPathNode::MissingNodesInConstantPathError

test/ruby_lsp_rails/indexing_enhancement_test.rb

@@ -112,6 +112,72 @@ class Post < ActiveRecord::Base
         assert_declaration_on_line("tags=", "Post", 5)
       end
 
+      test "inherited class_methods" do
+        @index.index_single(@indexable_path, <<~RUBY)
+          module TheConcern
+            extend ActiveSupport::Concern
+
+            class_methods do
+              def found_me; end
+            end
+          end
+
+          module OtherConcern
+            extend ActiveSupport::Concern
+            include TheConcern
+          end
+
+          class Foo
+            include OtherConcern
+          end
+        RUBY
+
+        ancestors = @index.linearized_ancestors_of("Foo::<Class:Foo>")
+
+        assert_includes(ancestors, "TheConcern::ClassMethods")
+        refute_nil(@index.resolve_method("found_me", "Foo::<Class:Foo>"))
+      end
+
+      test "prepended and inherited class_methods" do
+        @index.index_single(@indexable_path, <<~RUBY)
+          module TheConcern
+            extend ActiveSupport::Concern
+
+            class_methods do
+              def found_me; end
+            end
+          end
+
+          module OtherConcern
+            extend ActiveSupport::Concern
+            prepend TheConcern
+
+            module ClassMethods
+              def other_found_me; end
+            end
+          end
+
+          class Foo
+            include OtherConcern
+          end
+        RUBY
+
+        ancestors = @index.linearized_ancestors_of("Foo::<Class:Foo>")
+        relevant_ancestors = ancestors[0..ancestors.index("BasicObject::<Class:BasicObject>")]
+
+        assert_equal(
+          [
+            "Foo::<Class:Foo>",
+            "OtherConcern::ClassMethods",
+            "TheConcern::ClassMethods",
+            "Object::<Class:Object>",
+            "BasicObject::<Class:BasicObject>",
+          ],
+          relevant_ancestors,
+        )
+        refute_nil(@index.resolve_method("other_found_me", "Foo::<Class:Foo>"))
+      end
+
       private
 
       def assert_declaration_on_line(method_name, class_name, line)