From d3e085058b8789a3afb4ba78b422424f0576965a Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Thu, 8 Aug 2024 14:52:36 -0400
Subject: [PATCH] Add concern indexing enhancement

---
 lib/ruby_lsp/ruby_lsp_rails/addon.rb          |  3 +
 .../ruby_lsp_rails/indexing_enhancement.rb    | 63 +++++++++++++++++++
 .../indexing_enhancement_test.rb              | 39 ++++++++++++
 3 files changed, 105 insertions(+)
 create mode 100644 lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb
 create mode 100644 test/ruby_lsp_rails/indexing_enhancement_test.rb

diff --git a/lib/ruby_lsp/ruby_lsp_rails/addon.rb b/lib/ruby_lsp/ruby_lsp_rails/addon.rb
index ad82f387..7358aa56 100644
--- a/lib/ruby_lsp/ruby_lsp_rails/addon.rb
+++ b/lib/ruby_lsp/ruby_lsp_rails/addon.rb
@@ -13,6 +13,7 @@
 require_relative "code_lens"
 require_relative "document_symbol"
 require_relative "definition"
+require_relative "indexing_enhancement"
 
 module RubyLsp
   module Rails
@@ -35,6 +36,8 @@ def activate(global_state, message_queue)
         # Start booting the real client in a background thread. Until this completes, the client will be a NullClient
         Thread.new { @client = RunnerClient.create_client }
         register_additional_file_watchers(global_state: global_state, message_queue: message_queue)
+
+        T.must(@global_state).index.register_enhancement(IndexingEnhancement.new)
       end
 
       sig { override.void }
diff --git a/lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb b/lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb
new file mode 100644
index 00000000..cb16f59d
--- /dev/null
+++ b/lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb
@@ -0,0 +1,63 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyLsp
+  module Rails
+    class IndexingEnhancement
+      extend T::Sig
+      include RubyIndexer::Enhancement
+
+      sig do
+        override.params(
+          index: RubyIndexer::Index,
+          owner: T.nilable(RubyIndexer::Entry::Namespace),
+          node: Prism::CallNode,
+          file_path: String,
+        ).void
+      end
+      def on_call_node(index, owner, node, file_path)
+        return unless owner
+
+        name = node.name
+
+        case name
+        when :extend
+          handle_concern_extend(index, owner, node)
+        end
+      end
+
+      private
+
+      sig do
+        params(
+          index: RubyIndexer::Index,
+          owner: RubyIndexer::Entry::Namespace,
+          node: Prism::CallNode,
+        ).void
+      end
+      def handle_concern_extend(index, owner, node)
+        arguments = node.arguments&.arguments
+        return unless arguments
+
+        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 << RubyIndexer::Entry::Include.new(class_methods_name)
+            end
+          end
+        rescue Prism::ConstantPathNode::DynamicPartsInConstantPathError,
+               Prism::ConstantPathNode::MissingNodesInConstantPathError
+          # Do nothing
+        end
+      end
+    end
+  end
+end
diff --git a/test/ruby_lsp_rails/indexing_enhancement_test.rb b/test/ruby_lsp_rails/indexing_enhancement_test.rb
new file mode 100644
index 00000000..3eb2bb5b
--- /dev/null
+++ b/test/ruby_lsp_rails/indexing_enhancement_test.rb
@@ -0,0 +1,39 @@
+# typed: true
+# frozen_string_literal: true
+
+require "test_helper"
+
+module RubyLsp
+  module Rails
+    class IndexingEnhancementTest < ActiveSupport::TestCase
+      class << self
+        # For these tests, it's convenient to have the index fully populated with Rails information, but we don't have
+        # to reindex on every single example or that will be too slow
+        def populated_index
+          @index ||= begin
+            index = RubyIndexer::Index.new
+            index.register_enhancement(IndexingEnhancement.new)
+            index.index_all
+            index
+          end
+        end
+      end
+
+      def setup
+        @index = self.class.populated_index
+      end
+
+      test "ClassMethods module inside concerns are automatically extended" do
+        @index.index_single(RubyIndexer::IndexablePath.new(nil, "/fake.rb"), <<~RUBY)
+          class Post < ActiveRecord::Base
+          end
+        RUBY
+
+        ancestors = @index.linearized_ancestors_of("Post::<Class:Post>")
+        assert_includes(ancestors, "ActiveRecord::Associations::ClassMethods")
+        assert_includes(ancestors, "ActiveRecord::Store::ClassMethods")
+        assert_includes(ancestors, "ActiveRecord::AttributeMethods::ClassMethods")
+      end
+    end
+  end
+end