Update index enhancement due to ruby-lsp changes

Gemfile

@@ -19,6 +19,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: "main"
 
 platforms :mingw, :x64_mingw, :mswin, :jruby do
   gem "tzinfo"

Gemfile.lock

@@ -1,8 +1,19 @@
+GIT
+  remote: https://github.com/Shopify/ruby-lsp.git
+  revision: 0baf62000453877023764b8dcbb56fa252538789
+  branch: main
+  specs:
+    ruby-lsp (0.21.0)
+      language_server-protocol (~> 3.17.0)
+      prism (>= 1.2, < 2.0)
+      rbs (>= 3, < 4)
+      sorbet-runtime (>= 0.5.10782)
+
 PATH
   remote: .
   specs:
     ruby-lsp-rails (0.3.21)
-      ruby-lsp (>= 0.20.0, < 0.21.0)
+      ruby-lsp (>= 0.21.0, < 0.22.0)
 
 GEM
   remote: https://rubygems.org/
@@ -227,11 +238,6 @@ GEM
       rubocop (~> 1.51)
     rubocop-sorbet (0.8.3)
       rubocop (>= 0.90.0)
-    ruby-lsp (0.20.0)
-      language_server-protocol (~> 3.17.0)
-      prism (>= 1.2, < 2.0)
-      rbs (>= 3, < 4)
-      sorbet-runtime (>= 0.5.10782)
     ruby-progressbar (1.13.0)
     ruby2_keywords (0.0.5)
     sorbet (0.5.11577)
@@ -298,6 +304,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

@@ -53,7 +53,7 @@ def activate(global_state, outgoing_queue)
         @outgoing_queue << Notification.window_log_message("Activating Ruby LSP Rails add-on v#{VERSION}")
 
         register_additional_file_watchers(global_state: global_state, outgoing_queue: outgoing_queue)
-        @global_state.index.register_enhancement(IndexingEnhancement.new)
+        @global_state.index.register_enhancement(IndexingEnhancement.new(@global_state.index))
 
         # Start booting the real client in a background thread. Until this completes, the client will be a NullClient
         @client_mutex.unlock

lib/ruby_lsp/ruby_lsp_rails/indexing_enhancement.rb

@@ -3,13 +3,17 @@
 
 module RubyLsp
   module Rails
-    class IndexingEnhancement
+    class IndexingEnhancement < RubyIndexer::Enhancement
       extend T::Sig
-      include RubyIndexer::Enhancement
+
+      sig { params(index: RubyIndexer::Index).void }
+      def initialize(index)
+        super
+        @index = index
+      end
 
       sig do
         override.params(
-          index: RubyIndexer::Index,
           owner: T.nilable(RubyIndexer::Entry::Namespace),
           node: Prism::CallNode,
           file_path: String,
@@ -19,24 +23,23 @@ class IndexingEnhancement
           ),
         ).void
       end
-      def on_call_node(index, owner, node, file_path, code_units_cache)
+      def on_call_node_enter(owner, node, file_path, code_units_cache)
         return unless owner
 
         name = node.name
 
         case name
         when :extend
-          handle_concern_extend(index, owner, node)
+          handle_concern_extend(owner, node)
         when :has_one, :has_many, :belongs_to, :has_and_belongs_to_many
-          handle_association(index, owner, node, file_path, code_units_cache)
+          handle_association(owner, node, file_path, code_units_cache)
         end
       end
 
       private
 
       sig do
         params(
-          index: RubyIndexer::Index,
           owner: RubyIndexer::Entry::Namespace,
           node: Prism::CallNode,
           file_path: String,
@@ -46,7 +49,7 @@ def on_call_node(index, owner, node, file_path, code_units_cache)
           ),
         ).void
       end
-      def handle_association(index, owner, node, file_path, code_units_cache)
+      def handle_association(owner, node, file_path, code_units_cache)
         arguments = node.arguments&.arguments
         return unless arguments
 
@@ -64,7 +67,7 @@ def handle_association(index, owner, node, file_path, code_units_cache)
         loc = RubyIndexer::Location.from_prism_location(name_arg.location, code_units_cache)
 
         # Reader
-        index.add(RubyIndexer::Entry::Method.new(
+        @index.add(RubyIndexer::Entry::Method.new(
           name,
           file_path,
           loc,
@@ -76,7 +79,7 @@ def handle_association(index, owner, node, file_path, code_units_cache)
         ))
 
         # Writer
-        index.add(RubyIndexer::Entry::Method.new(
+        @index.add(RubyIndexer::Entry::Method.new(
           "#{name}=",
           file_path,
           loc,
@@ -90,12 +93,11 @@ def handle_association(index, owner, node, file_path, code_units_cache)
 
       sig do
         params(
-          index: RubyIndexer::Index,
           owner: RubyIndexer::Entry::Namespace,
           node: Prism::CallNode,
         ).void
       end
-      def handle_concern_extend(index, owner, node)
+      def handle_concern_extend(owner, node)
         arguments = node.arguments&.arguments
         return unless arguments
 
@@ -105,7 +107,7 @@ def handle_concern_extend(index, owner, node)
           module_name = node.full_name
           next unless module_name == "ActiveSupport::Concern"
 
-          index.register_included_hook(owner.name) do |index, base|
+          @index.register_included_hook(owner.name) do |index, base|
             class_methods_name = "#{owner.name}::ClassMethods"
 
             if index.indexed?(class_methods_name)

ruby-lsp-rails.gemspec

@@ -24,5 +24,5 @@ Gem::Specification.new do |spec|
     Dir["lib/**/*", "LICENSE.txt", "Rakefile", "README.md"]
   end
 
-  spec.add_dependency("ruby-lsp", ">= 0.20.0", "< 0.21.0")
+  spec.add_dependency("ruby-lsp", ">= 0.21.0", "< 0.22.0")
 end