Add concern indexing augmentation

Author: vinistock

Gemfile

@@ -20,6 +20,7 @@ gem "tapioca", "~> 0.13", require: false, platforms: :ruby
 gem "psych", "~> 5.1", require: false
 gem "rails"
 gem "webmock"
+gem "ruby-lsp", github: "Shopify/ruby-lsp", branch: "vs-indexing-augmentations"
 
 platforms :mingw, :x64_mingw, :mswin, :jruby do
   gem "tzinfo"

Gemfile.lock

@@ -6,6 +6,17 @@ GIT
     rdoc (6.6.3.1)
       psych (>= 4.0.0)
 
+GIT
+  remote: https://github.com/Shopify/ruby-lsp.git
+  revision: c83da98b4a7286c74f976309a84293b05817affc
+  branch: vs-indexing-augmentations
+  specs:
+    ruby-lsp (0.17.10)
+      language_server-protocol (~> 3.17.0)
+      prism (>= 0.29.0, < 0.31)
+      rbs (>= 3, < 4)
+      sorbet-runtime (>= 0.5.10782)
+
 PATH
   remote: .
   specs:
@@ -146,8 +157,6 @@ GEM
     nio4r (2.7.3)
     nokogiri (1.16.5-arm64-darwin)
       racc (~> 1.4)
-    nokogiri (1.16.5-x64-mingw-ucrt)
-      racc (~> 1.4)
     nokogiri (1.16.5-x86_64-darwin)
       racc (~> 1.4)
     nokogiri (1.16.5-x86_64-linux)
@@ -234,11 +243,6 @@ GEM
       rubocop (~> 1.51)
     rubocop-sorbet (0.8.3)
       rubocop (>= 0.90.0)
-    ruby-lsp (0.17.9)
-      language_server-protocol (~> 3.17.0)
-      prism (>= 0.29.0, < 0.31)
-      rbs (>= 3, < 4)
-      sorbet-runtime (>= 0.5.10782)
     ruby-progressbar (1.13.0)
     ruby2_keywords (0.0.5)
     sorbet (0.5.11406)
@@ -255,7 +259,6 @@ GEM
       sorbet-static-and-runtime (>= 0.5.10187)
       thor (>= 0.19.2)
     sqlite3 (1.7.3-arm64-darwin)
-    sqlite3 (1.7.3-x64-mingw-ucrt)
     sqlite3 (1.7.3-x86_64-darwin)
     sqlite3 (1.7.3-x86_64-linux)
     stringio (3.1.0)
@@ -273,8 +276,6 @@ GEM
     timeout (0.4.1)
     tzinfo (2.0.6)
       concurrent-ruby (~> 1.0)
-    tzinfo-data (1.2024.1)
-      tzinfo (>= 1.0.0)
     unicode-display_width (2.5.0)
     webmock (3.23.1)
       addressable (>= 2.8.0)
@@ -292,7 +293,6 @@ GEM
 
 PLATFORMS
   arm64-darwin
-  x64-mingw-ucrt
   x86_64-darwin
   x86_64-linux
 
@@ -307,6 +307,7 @@ DEPENDENCIES
   rubocop-rake (~> 0.6.0)
   rubocop-shopify (~> 2.15)
   rubocop-sorbet (~> 0.8)
+  ruby-lsp!
   ruby-lsp-rails!
   sorbet-static-and-runtime
   sqlite3 (< 2)

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_augmentation"
 
 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_augmentation(IndexingAugmentation.new)
       end
 
       sig { override.void }
@@ -77,7 +80,10 @@ def create_document_symbol_listener(response_builder, dispatcher)
 
       sig do
         override.params(
-          response_builder: ResponseBuilders::CollectionResponseBuilder[Interface::Location],
+          response_builder: ResponseBuilders::CollectionResponseBuilder[T.any(
+            Interface::Location,
+            Interface::LocationLink,
+          )],
           uri: URI::Generic,
           node_context: NodeContext,
           dispatcher: Prism::Dispatcher,

lib/ruby_lsp/ruby_lsp_rails/definition.rb

@@ -34,7 +34,10 @@ class Definition
       sig do
         params(
           client: RunnerClient,
-          response_builder: ResponseBuilders::CollectionResponseBuilder[Interface::Location],
+          response_builder: ResponseBuilders::CollectionResponseBuilder[T.any(
+            Interface::Location,
+            Interface::LocationLink,
+          )],
           node_context: NodeContext,
           index: RubyIndexer::Index,
           dispatcher: Prism::Dispatcher,

lib/ruby_lsp/ruby_lsp_rails/indexing_augmentation.rb

@@ -0,0 +1,63 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyLsp
+  module Rails
+    class IndexingAugmentation
+      extend T::Sig
+      include RubyIndexer::Augmentation
+
+      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