Restructure DocumentSymbol request (#1311)

Expose DocumentSymbolStack and convert restructure DocumentSymbol request

lib/ruby_lsp/addon.rb

@@ -128,10 +128,11 @@ def create_hover_listener(nesting, index, dispatcher); end
     # Creates a new DocumentSymbol listener. This method is invoked on every DocumentSymbol request
     sig do
       overridable.params(
+        stack: Response::DocumentSymbolStack,
         dispatcher: Prism::Dispatcher,
-      ).returns(T.nilable(Listener[T::Array[Interface::DocumentSymbol]]))
+      ).void
     end
-    def create_document_symbol_listener(dispatcher); end
+    def create_document_symbol_listener(stack, dispatcher); end
 
     # Creates a new Definition listener. This method is invoked on every Definition request
     sig do

lib/ruby_lsp/internal.rb

@@ -23,6 +23,7 @@
 require "ruby_lsp/server"
 require "ruby_lsp/executor"
 require "ruby_lsp/requests"
+require "ruby_lsp/response"
 require "ruby_lsp/listener"
 require "ruby_lsp/document"
 require "ruby_lsp/ruby_document"

lib/ruby_lsp/listeners/document_symbol.rb

@@ -3,39 +3,21 @@
 
 module RubyLsp
   module Listeners
-    class DocumentSymbol < Listener
+    class DocumentSymbol
       extend T::Sig
       extend T::Generic
-
-      ResponseType = type_member { { fixed: T::Array[Interface::DocumentSymbol] } }
+      include Requests::Support::Common
 
       ATTR_ACCESSORS = T.let([:attr_reader, :attr_writer, :attr_accessor].freeze, T::Array[Symbol])
 
-      class SymbolHierarchyRoot
-        extend T::Sig
-
-        sig { returns(T::Array[Interface::DocumentSymbol]) }
-        attr_reader :children
-
-        sig { void }
-        def initialize
-          @children = T.let([], T::Array[Interface::DocumentSymbol])
-        end
+      sig do
+        params(
+          stack: Response::DocumentSymbolStack,
+          dispatcher: Prism::Dispatcher,
+        ).void
       end
-
-      sig { override.returns(T::Array[Interface::DocumentSymbol]) }
-      attr_reader :_response
-
-      sig { params(dispatcher: Prism::Dispatcher).void }
-      def initialize(dispatcher)
-        @root = T.let(SymbolHierarchyRoot.new, SymbolHierarchyRoot)
-        @_response = T.let(@root.children, T::Array[Interface::DocumentSymbol])
-        @stack = T.let(
-          [@root],
-          T::Array[T.any(SymbolHierarchyRoot, Interface::DocumentSymbol)],
-        )
-
-        super
+      def initialize(stack, dispatcher)
+        @stack = stack
 
         dispatcher.register(
           self,
@@ -214,7 +196,7 @@ def create_document_symbol(name:, kind:, range_location:, selection_range_locati
           children: [],
         )
 
-        T.must(@stack.last).children << symbol
+        @stack.last.children << symbol
 
         symbol
       end

lib/ruby_lsp/requests/document_symbol.rb

@@ -46,25 +46,20 @@ def provider
         end
       end
 
-      ResponseType = type_member { { fixed: T::Array[Interface::DocumentSymbol] } }
-
       sig { params(dispatcher: Prism::Dispatcher).void }
       def initialize(dispatcher)
         super()
-        @listeners = T.let(
-          [Listeners::DocumentSymbol.new(dispatcher)],
-          T::Array[Listener[ResponseType]],
-        )
+        @stack = T.let(Response::DocumentSymbolStack.new, Response::DocumentSymbolStack)
+        Listeners::DocumentSymbol.new(@stack, dispatcher)
 
         Addon.addons.each do |addon|
-          addon_listener = addon.create_document_symbol_listener(dispatcher)
-          @listeners << addon_listener if addon_listener
+          addon.create_document_symbol_listener(@stack, dispatcher)
         end
       end
 
-      sig { override.returns(ResponseType) }
+      sig { override.returns(T::Array[Interface::DocumentSymbol]) }
       def perform
-        @listeners.flat_map(&:response).compact
+        @stack.result
       end
     end
   end

lib/ruby_lsp/response.rb

@@ -0,0 +1,67 @@
+# typed: strict
+# frozen_string_literal: true
+
+module RubyLsp
+  class Response
+    extend T::Sig
+    extend T::Generic
+
+    abstract!
+
+    ResponseType = type_member
+
+    sig { abstract.returns(ResponseType) }
+    def result; end
+
+    class DocumentSymbolStack < Response
+      ResponseType = type_member { { fixed: T::Array[Interface::DocumentSymbol] } }
+
+      class SymbolHierarchyRoot
+        extend T::Sig
+
+        sig { returns(T::Array[Interface::DocumentSymbol]) }
+        attr_reader :children
+
+        sig { void }
+        def initialize
+          @children = T.let([], T::Array[Interface::DocumentSymbol])
+        end
+      end
+
+      extend T::Sig
+
+      sig { void }
+      def initialize
+        super
+        @stack = T.let(
+          [SymbolHierarchyRoot.new],
+          T::Array[T.any(SymbolHierarchyRoot, Interface::DocumentSymbol)],
+        )
+      end
+
+      sig { params(symbol: Interface::DocumentSymbol).void }
+      def push(symbol)
+        @stack << symbol
+      end
+
+      alias_method(:<<, :push)
+
+      sig { returns(T.nilable(Interface::DocumentSymbol)) }
+      def pop
+        if @stack.size > 1
+          T.cast(@stack.pop, Interface::DocumentSymbol)
+        end
+      end
+
+      sig { returns(T.any(SymbolHierarchyRoot, Interface::DocumentSymbol)) }
+      def last
+        T.must(@stack.last)
+      end
+
+      sig { override.returns(T::Array[Interface::DocumentSymbol]) }
+      def result
+        T.must(@stack.first).children
+      end
+    end
+  end
+end

test/requests/document_symbol_expectations_test.rb

@@ -9,18 +9,27 @@ class DocumentSymbolExpectationsTest < ExpectationsTestRunner
 
   def test_document_symbol_addons
     source = <<~RUBY
-      test "foo" do
+      class Foo
+        test "foo" do
+        end
       end
     RUBY
 
     test_addon(:create_document_symbol_addon, source: source) do |executor|
       response = executor.execute({
         method: "textDocument/documentSymbol",
         params: { textDocument: { uri: "file:///fake.rb" } },
-      }).response
+      })
+
+      assert_nil(response.error, response.error&.full_message)
+
+      response = response.response
+
+      assert_equal(1, response.count)
+      assert_equal("Foo", response.first.name)
 
-      assert_equal("foo", response.first.name)
-      assert_equal(LanguageServer::Protocol::Constant::SymbolKind::METHOD, response.first.kind)
+      test_symbol = response.first.children.first
+      assert_equal(LanguageServer::Protocol::Constant::SymbolKind::METHOD, test_symbol.kind)
     end
   end
 
@@ -34,31 +43,32 @@ def name
         "Document SymbolsAddon"
       end
 
-      def create_document_symbol_listener(dispatcher)
-        klass = Class.new(RubyLsp::Listener) do
-          attr_reader :_response
+      def create_document_symbol_listener(stack, dispatcher)
+        klass = Class.new do
+          include RubyLsp::Requests::Support::Common
 
-          def initialize(dispatcher)
-            super
+          def initialize(stack, dispatcher)
+            @stack = stack
             dispatcher.register(self, :on_call_node_enter)
           end
 
           def on_call_node_enter(node)
-            T.bind(self, RubyLsp::Listener[T.untyped])
+            parent = @stack.last
+            T.bind(self, RubyLsp::Requests::Support::Common)
             message_value = node.message
             arguments = node.arguments&.arguments
             return unless message_value == "test" && arguments&.any?
 
-            @_response = [RubyLsp::Interface::DocumentSymbol.new(
+            parent.children << RubyLsp::Interface::DocumentSymbol.new(
               name: arguments.first.content,
               kind: LanguageServer::Protocol::Constant::SymbolKind::METHOD,
               selection_range: range_from_node(node),
               range: range_from_node(node),
-            )]
+            )
           end
         end
 
-        klass.new(dispatcher)
+        T.unsafe(klass).new(stack, dispatcher)
       end
     end
   end