Integrate document symbol for callbacks

Author: Aryan Soni

lib/ruby_lsp/ruby_lsp_rails/document_symbol.rb

@@ -12,6 +12,61 @@ class DocumentSymbol
       include Requests::Support::Common
       include ActiveSupportTestCaseHelper
 
+      MODEL_CALLBACKS = T.let(
+        [
+          "before_validation",
+          "after_validation",
+          "before_save",
+          "around_save",
+          "before_create",
+          "around_create",
+          "after_create",
+          "after_save",
+          "after_commit",
+          "after_rollback",
+          "before_update",
+          "around_update",
+          "after_update",
+          "before_destroy",
+          "around_destroy",
+          "after_destroy",
+          "after_initialize",
+          "after_find",
+          "after_touch",
+        ].freeze,
+        T::Array[String],
+      )
+
+      CONTROLLER_CALLBACKS = T.let(
+        [
+          "after_action",
+          "append_after_action",
+          "append_around_action",
+          "append_before_action",
+          "around_action",
+          "before_action",
+          "prepend_after_action",
+          "prepend_around_action",
+          "prepend_before_action",
+          "skip_after_action",
+          "skip_around_action",
+          "skip_before_action",
+        ].freeze,
+        T::Array[String],
+      )
+
+      JOB_CALLBACKS = T.let(
+        [
+          "after_enqueue",
+          "after_perform",
+          "around_enqueue",
+          "around_perform",
+          "before_enqueue",
+          "before_perform",
+        ].freeze,
+        T::Array[String],
+      )
+
       sig do
         params(
           response_builder: ResponseBuilders::DocumentSymbol,
@@ -28,13 +83,79 @@ def initialize(response_builder, dispatcher)
       def on_call_node_enter(node)
         content = extract_test_case_name(node)
 
-        return unless content
+        if content
+          append_document_symbol(
+            name: content,
+            selection_range: range_from_node(node),
+            range: range_from_node(node),
+          )
+        end
+
+        extract_callbacks(node)
+      end
+
+      private
+
+      sig { params(node: Prism::CallNode).void }
+      def extract_callbacks(node)
+        receiver = node.receiver
+        return if receiver && !receiver.is_a?(Prism::SelfNode)
+
+        message_value = node.message
+        callbacks = [MODEL_CALLBACKS, CONTROLLER_CALLBACKS, JOB_CALLBACKS]
+
+        return unless callbacks.any? { |callback| callback.include?(message_value) }
+
+        block = node.block
 
+        if block
+          append_document_symbol(
+            name: "#{message_value}(<anonymous>)",
+            range: range_from_location(node.location),
+            selection_range: range_from_location(block.location),
+          )
+          return
+        end
+
+        arguments = node.arguments&.arguments
+        return unless arguments&.any?
+
+        arguments.each do |argument|
+          if argument.is_a?(Prism::SymbolNode)
+            name = argument.value
+            next unless name
+
+            append_document_symbol(
+              name: "#{message_value}(#{name})",
+              range: range_from_location(argument.location),
+              selection_range: range_from_location(T.must(argument.value_loc)),
+            )
+          elsif argument.is_a?(Prism::StringNode)
+            name = argument.content
+            next if name.empty?
+
+            append_document_symbol(
+              name: "#{message_value}(#{name})",
+              range: range_from_location(argument.location),
+              selection_range: range_from_location(argument.content_loc),
+            )
+          end
+        end
+      end
+
+      sig do
+        params(
+          name: String,
+          range: RubyLsp::Interface::Range,
+          selection_range: RubyLsp::Interface::Range,
+        ).void
+      end
+      def append_document_symbol(name:, range:, selection_range:)
         @response_builder.last.children << RubyLsp::Interface::DocumentSymbol.new(
-          name: content,
-          kind: LanguageServer::Protocol::Constant::SymbolKind::METHOD,
-          selection_range: range_from_node(node),
-          range: range_from_node(node),
+          name: name,
+          kind: RubyLsp::Constant::SymbolKind::METHOD,
+          range: range,
+          selection_range: selection_range,
         )
       end
     end

test/ruby_lsp_rails/document_symbol_test.rb

@@ -163,6 +163,60 @@ class NestedTest < ActiveSupport::TestCase
         assert_equal("back to the same level", response[0].children[2].name)
       end
 
+      test "correctly handles model callbacks with multiple Prism::StringNode arguments" do
+        response = generate_document_symbols_for_source(<<~RUBY)
+          class FooModel < ApplicationRecord
+            before_save "foo_method", "bar_method", on: :update
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+        assert_equal("FooModel", response[0].name)
+        assert_equal(2, response[0].children.size)
+        assert_equal("before_save(foo_method)", response[0].children[0].name)
+        assert_equal("before_save(bar_method)", response[0].children[1].name)
+      end
+
+      test "correctly handles controller callback with block" do
+        response = generate_document_symbols_for_source(<<~RUBY)
+          class FooController < ApplicationController
+            before_action do
+              # block body
+            end
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+        assert_equal("FooController", response[0].name)
+        assert_equal(1, response[0].children.size)
+        assert_equal("before_action(<anonymous>)", response[0].children[0].name)
+      end
+
+      test "correctly handles job callback with Prism::SymbolNode argument" do
+        response = generate_document_symbols_for_source(<<~RUBY)
+          class FooJob < ApplicationJob
+            before_perform :foo_method
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+        assert_equal("FooJob", response[0].name)
+        assert_equal(1, response[0].children.size)
+        assert_equal("before_perform(foo_method)", response[0].children[0].name)
+      end
+
+      test "ignore unrecognized callback" do
+        response = generate_document_symbols_for_source(<<~RUBY)
+          class FooJob < ApplicationJob
+            unrecognized_callback :foo_method
+          end
+        RUBY
+
+        assert_equal(1, response.size)
+        assert_equal("FooJob", response[0].name)
+        assert_empty(response[0].children)
+      end
+
       private
 
       def generate_document_symbols_for_source(source)