Support on_call_node_leave events in index enhancement

jekyll/add-ons.markdown

@@ -271,16 +271,14 @@ end
 This is how you could write an enhancement to teach the Ruby LSP to understand that DSL:
 
 ```ruby
-class MyIndexingEnhancement
-  include RubyLsp::Enhancement
-
+class MyIndexingEnhancement < RubyIndexer::Enhancement
   # This on call node handler is invoked any time during indexing when we find a method call. It can be used to insert
   # more entries into the index depending on the conditions
-  def on_call_node(index, owner, node, file_path)
+  def on_call_node_enter(owner, node, file_path, code_units_cache)
     return unless owner
 
     # Get the ancestors of the current class
-    ancestors = index.linearized_ancestors_of(owner.name)
+    ancestors = @index.linearized_ancestors_of(owner.name)
 
     # Return early unless the method call is the one we want to handle and the class invoking the DSL inherits from
     # our library's parent class
@@ -304,15 +302,19 @@ class MyIndexingEnhancement
       location,     # The Prism node location where the DSL call was found
       location,     # The Prism node location for the DSL name location. May or not be the same
       nil,          # The documentation for this DSL call. This should always be `nil` to ensure lazy fetching of docs
-      index.configuration.encoding, # The negotiated encoding. This should always be `indexing.configuration.encoding`
+      @index.configuration.encoding, # The negotiated encoding. This should always be `indexing.configuration.encoding`
       signatures,   # All signatures for this method (every way it can be invoked)
       Entry::Visibility::PUBLIC, # The method's visibility
       owner,        # The method's owner. This is almost always going to be the same owner received
     )
 
     # Push the new entry to the index
-    index.add(new_entry)
+    @index.add(new_entry)
   end
+
+  # This method is invoked when the parser has finished processing the method call node.
+  # It can be used to perform cleanups like popping a stack...etc.
+  def on_call_node_leave(owner, node, file_path, code_units_cache); end
 end
 ```
 
@@ -324,7 +326,7 @@ module RubyLsp
     class Addon < ::RubyLsp::Addon
       def activate(global_state, message_queue)
         # Register the enhancement as part of the indexing process
-        @index.register_enhancement(MyIndexingEnhancement.new)
+        @index.register_enhancement(MyIndexingEnhancement.new(@index))
       end
 
       def deactivate

lib/ruby_indexer/lib/ruby_indexer/declaration_listener.rb

@@ -315,9 +315,11 @@ def on_call_node_enter(node)
       end
 
       @enhancements.each do |enhancement|
-        enhancement.on_call_node(@index, @owner_stack.last, node, @file_path, @code_units_cache)
+        enhancement.on_call_node_enter(@owner_stack.last, node, @file_path, @code_units_cache)
       rescue StandardError => e
-        @indexing_errors << "Indexing error in #{@file_path} with '#{enhancement.class.name}' enhancement: #{e.message}"
+        @indexing_errors << <<~MSG
+          Indexing error in #{@file_path} with '#{enhancement.class.name}' on call node enter enhancement: #{e.message}
+        MSG
       end
     end
 
@@ -332,6 +334,14 @@ def on_call_node_leave(node)
           @visibility_stack.pop
         end
       end
+
+      @enhancements.each do |enhancement|
+        enhancement.on_call_node_leave(@owner_stack.last, node, @file_path, @code_units_cache)
+      rescue StandardError => e
+        @indexing_errors << <<~MSG
+          Indexing error in #{@file_path} with '#{enhancement.class.name}' on call node leave enhancement: #{e.message}
+        MSG
+      end
     end
 
     sig { params(node: Prism::DefNode).void }

lib/ruby_indexer/lib/ruby_indexer/enhancement.rb

@@ -2,20 +2,35 @@
 # frozen_string_literal: true
 
 module RubyIndexer
-  module Enhancement
+  class Enhancement
     extend T::Sig
     extend T::Helpers
 
-    interface!
+    abstract!
 
-    requires_ancestor { Object }
+    sig { params(index: Index).void }
+    def initialize(index)
+      @index = index
+    end
 
     # The `on_extend` indexing enhancement 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,
+      overridable.params(
+        owner: T.nilable(Entry::Namespace),
+        node: Prism::CallNode,
+        file_path: String,
+        code_units_cache: T.any(
+          T.proc.params(arg0: Integer).returns(Integer),
+          Prism::CodeUnitsCache,
+        ),
+      ).void
+    end
+    def on_call_node_enter(owner, node, file_path, code_units_cache); end
+
+    sig do
+      overridable.params(
         owner: T.nilable(Entry::Namespace),
         node: Prism::CallNode,
         file_path: String,
@@ -25,6 +40,6 @@ module Enhancement
         ),
       ).void
     end
-    def on_call_node(index, owner, node, file_path, code_units_cache); end
+    def on_call_node_leave(owner, node, file_path, code_units_cache); end
   end
 end

lib/ruby_indexer/test/enhancements_test.rb

@@ -6,10 +6,8 @@
 module RubyIndexer
   class EnhancementTest < TestCase
     def test_enhancing_indexing_included_hook
-      enhancement_class = Class.new do
-        include Enhancement
-
-        def on_call_node(index, owner, node, file_path, code_units_cache)
+      enhancement_class = Class.new(Enhancement) do
+        def on_call_node_enter(owner, node, file_path, code_units_cache)
           return unless owner
           return unless node.name == :extend
 
@@ -24,7 +22,7 @@ def on_call_node(index, owner, node, file_path, code_units_cache)
             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)
@@ -33,7 +31,7 @@ def on_call_node(index, owner, node, file_path, code_units_cache)
               end
             end
 
-            index.add(Entry::Method.new(
+            @index.add(Entry::Method.new(
               "new_method",
               file_path,
               location,
@@ -50,7 +48,7 @@ def on_call_node(index, owner, node, file_path, code_units_cache)
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
       index(<<~RUBY)
         module ActiveSupport
           module Concern
@@ -98,10 +96,8 @@ class User < ActiveRecord::Base
     end
 
     def test_enhancing_indexing_configuration_dsl
-      enhancement_class = Class.new do
-        include Enhancement
-
-        def on_call_node(index, owner, node, file_path, code_units_cache)
+      enhancement_class = Class.new(Enhancement) do
+        def on_call_node_enter(owner, node, file_path, code_units_cache)
           return unless owner
 
           name = node.name
@@ -115,7 +111,7 @@ def on_call_node(index, owner, node, file_path, code_units_cache)
 
           location = Location.from_prism_location(association_name.location, code_units_cache)
 
-          index.add(Entry::Method.new(
+          @index.add(Entry::Method.new(
             T.must(association_name.value),
             file_path,
             location,
@@ -128,7 +124,7 @@ def on_call_node(index, owner, node, file_path, code_units_cache)
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
       index(<<~RUBY)
         module ActiveSupport
           module Concern
@@ -160,11 +156,44 @@ class User < ActiveRecord::Base
       assert_entry("posts", Entry::Method, "/fake/path/foo.rb:23-11:23-17")
     end
 
-    def test_error_handling_in_enhancement
-      enhancement_class = Class.new do
-        include Enhancement
+    def test_error_handling_in_on_call_node_enter_enhancement
+      enhancement_class = Class.new(Enhancement) do
+        def on_call_node_enter(owner, node, file_path, code_units_cache)
+          raise "Error"
+        end
+
+        class << self
+          def name
+            "TestEnhancement"
+          end
+        end
+      end
+
+      @index.register_enhancement(enhancement_class.new(@index))
+
+      _stdout, stderr = capture_io do
+        index(<<~RUBY)
+          module ActiveSupport
+            module Concern
+              def self.extended(base)
+                base.class_eval("def new_method(a); end")
+              end
+            end
+          end
+        RUBY
+      end
+
+      assert_match(
+        %r{Indexing error in /fake/path/foo\.rb with 'TestEnhancement' on call node enter enhancement},
+        stderr,
+      )
+      # The module should still be indexed
+      assert_entry("ActiveSupport::Concern", Entry::Module, "/fake/path/foo.rb:1-2:5-5")
+    end
 
-        def on_call_node(index, owner, node, file_path, code_units_cache)
+    def test_error_handling_in_on_call_node_leave_enhancement
+      enhancement_class = Class.new(Enhancement) do
+        def on_call_node_leave(owner, node, file_path, code_units_cache)
           raise "Error"
         end
 
@@ -175,7 +204,7 @@ def name
         end
       end
 
-      @index.register_enhancement(enhancement_class.new)
+      @index.register_enhancement(enhancement_class.new(@index))
 
       _stdout, stderr = capture_io do
         index(<<~RUBY)
@@ -189,7 +218,10 @@ def self.extended(base)
         RUBY
       end
 
-      assert_match(%r{Indexing error in /fake/path/foo\.rb with 'TestEnhancement' enhancement}, stderr)
+      assert_match(
+        %r{Indexing error in /fake/path/foo\.rb with 'TestEnhancement' on call node leave enhancement},
+        stderr,
+      )
       # The module should still be indexed
       assert_entry("ActiveSupport::Concern", Entry::Module, "/fake/path/foo.rb:1-2:5-5")
     end