Make document generic (#2436)

* Make Document generic

* Move locate to RubyDocument and ERBDocument

Author: vinistock

lib/ruby_lsp/document.rb

@@ -12,10 +12,13 @@ class LanguageId < T::Enum
 
     extend T::Sig
     extend T::Helpers
+    extend T::Generic
+
+    ParseResultType = type_member
 
     abstract!
 
-    sig { returns(Prism::ParseResult) }
+    sig { returns(ParseResultType) }
     attr_reader :parse_result
 
     sig { returns(String) }
@@ -38,10 +41,10 @@ def initialize(source:, version:, uri:, encoding: Encoding::UTF_8)
       @version = T.let(version, Integer)
       @uri = T.let(uri, URI::Generic)
       @needs_parsing = T.let(true, T::Boolean)
-      @parse_result = T.let(parse, Prism::ParseResult)
+      @parse_result = T.let(parse, ParseResultType)
     end
 
-    sig { params(other: Document).returns(T::Boolean) }
+    sig { params(other: Document[T.untyped]).returns(T::Boolean) }
     def ==(other)
       self.class == other.class && uri == other.uri && @source == other.source
     end
@@ -54,7 +57,7 @@ def language_id; end
       type_parameters(:T)
         .params(
           request_name: String,
-          block: T.proc.params(document: Document).returns(T.type_parameter(:T)),
+          block: T.proc.params(document: Document[ParseResultType]).returns(T.type_parameter(:T)),
         ).returns(T.type_parameter(:T))
     end
     def cache_fetch(request_name, &block)
@@ -93,7 +96,7 @@ def push_edits(edits, version:)
       @cache.clear
     end
 
-    sig { abstract.returns(Prism::ParseResult) }
+    sig { abstract.returns(ParseResultType) }
     def parse; end
 
     sig { abstract.returns(T::Boolean) }
@@ -104,115 +107,6 @@ def create_scanner
       Scanner.new(@source, @encoding)
     end
 
-    sig do
-      params(
-        position: T::Hash[Symbol, T.untyped],
-        node_types: T::Array[T.class_of(Prism::Node)],
-      ).returns(NodeContext)
-    end
-    def locate_node(position, node_types: [])
-      locate(@parse_result.value, create_scanner.find_char_position(position), node_types: node_types)
-    end
-
-    sig do
-      params(
-        node: Prism::Node,
-        char_position: Integer,
-        node_types: T::Array[T.class_of(Prism::Node)],
-      ).returns(NodeContext)
-    end
-    def locate(node, char_position, node_types: [])
-      queue = T.let(node.child_nodes.compact, T::Array[T.nilable(Prism::Node)])
-      closest = node
-      parent = T.let(nil, T.nilable(Prism::Node))
-      nesting_nodes = T.let(
-        [],
-        T::Array[T.any(
-          Prism::ClassNode,
-          Prism::ModuleNode,
-          Prism::SingletonClassNode,
-          Prism::DefNode,
-          Prism::BlockNode,
-          Prism::LambdaNode,
-          Prism::ProgramNode,
-        )],
-      )
-
-      nesting_nodes << node if node.is_a?(Prism::ProgramNode)
-      call_node = T.let(nil, T.nilable(Prism::CallNode))
-
-      until queue.empty?
-        candidate = queue.shift
-
-        # Skip nil child nodes
-        next if candidate.nil?
-
-        # Add the next child_nodes to the queue to be processed. The order here is important! We want to move in the
-        # same order as the visiting mechanism, which means searching the child nodes before moving on to the next
-        # sibling
-        T.unsafe(queue).unshift(*candidate.child_nodes)
-
-        # Skip if the current node doesn't cover the desired position
-        loc = candidate.location
-        next unless (loc.start_offset...loc.end_offset).cover?(char_position)
-
-        # If the node's start character is already past the position, then we should've found the closest node
-        # already
-        break if char_position < loc.start_offset
-
-        # If the candidate starts after the end of the previous nesting level, then we've exited that nesting level and
-        # need to pop the stack
-        previous_level = nesting_nodes.last
-        nesting_nodes.pop if previous_level && loc.start_offset > previous_level.location.end_offset
-
-        # Keep track of the nesting where we found the target. This is used to determine the fully qualified name of the
-        # target when it is a constant
-        case candidate
-        when Prism::ClassNode, Prism::ModuleNode, Prism::SingletonClassNode, Prism::DefNode, Prism::BlockNode,
-          Prism::LambdaNode
-          nesting_nodes << candidate
-        end
-
-        if candidate.is_a?(Prism::CallNode)
-          arg_loc = candidate.arguments&.location
-          blk_loc = candidate.block&.location
-          if (arg_loc && (arg_loc.start_offset...arg_loc.end_offset).cover?(char_position)) ||
-              (blk_loc && (blk_loc.start_offset...blk_loc.end_offset).cover?(char_position))
-            call_node = candidate
-          end
-        end
-
-        # If there are node types to filter by, and the current node is not one of those types, then skip it
-        next if node_types.any? && node_types.none? { |type| candidate.class == type }
-
-        # If the current node is narrower than or equal to the previous closest node, then it is more precise
-        closest_loc = closest.location
-        if loc.end_offset - loc.start_offset <= closest_loc.end_offset - closest_loc.start_offset
-          parent = closest
-          closest = candidate
-        end
-      end
-
-      # When targeting the constant part of a class/module definition, we do not want the nesting to be duplicated. That
-      # is, when targeting Bar in the following example:
-      #
-      # ```ruby
-      #   class Foo::Bar; end
-      # ```
-      # The correct target is `Foo::Bar` with an empty nesting. `Foo::Bar` should not appear in the nesting stack, even
-      # though the class/module node does indeed enclose the target, because it would lead to incorrect behavior
-      if closest.is_a?(Prism::ConstantReadNode) || closest.is_a?(Prism::ConstantPathNode)
-        last_level = nesting_nodes.last
-
-        if (last_level.is_a?(Prism::ModuleNode) || last_level.is_a?(Prism::ClassNode)) &&
-            last_level.constant_path == closest
-          nesting_nodes.pop
-        end
-      end
-
-      NodeContext.new(closest, parent, nesting_nodes, call_node)
-    end
-
     class Scanner
       extend T::Sig
 

lib/ruby_lsp/erb_document.rb

@@ -4,8 +4,11 @@
 module RubyLsp
   class ERBDocument < Document
     extend T::Sig
+    extend T::Generic
 
-    sig { override.returns(Prism::ParseResult) }
+    ParseResultType = type_member { { fixed: Prism::ParseResult } }
+
+    sig { override.returns(ParseResultType) }
     def parse
       return @parse_result unless @needs_parsing
 
@@ -26,6 +29,16 @@ def language_id
       LanguageId::ERB
     end
 
+    sig do
+      params(
+        position: T::Hash[Symbol, T.untyped],
+        node_types: T::Array[T.class_of(Prism::Node)],
+      ).returns(NodeContext)
+    end
+    def locate_node(position, node_types: [])
+      RubyDocument.locate(@parse_result.value, create_scanner.find_char_position(position), node_types: node_types)
+    end
+
     class ERBScanner
       extend T::Sig
 

lib/ruby_lsp/requests/code_action_resolve.rb

@@ -112,7 +112,7 @@ def refactor_variable
         extracted_source = T.must(@document.source[start_index...end_index])
 
         # Find the closest statements node, so that we place the refactor in a valid position
-        node_context = @document
+        node_context = RubyDocument
           .locate(@document.parse_result.value, start_index, node_types: [Prism::StatementsNode, Prism::BlockNode])
 
         closest_statements = node_context.node
@@ -206,7 +206,7 @@ def refactor_method
         extracted_source = T.must(@document.source[start_index...end_index])
 
         # Find the closest method declaration node, so that we place the refactor in a valid position
-        node_context = @document.locate(@document.parse_result.value, start_index, node_types: [Prism::DefNode])
+        node_context = RubyDocument.locate(@document.parse_result.value, start_index, node_types: [Prism::DefNode])
         closest_def = T.cast(node_context.node, Prism::DefNode)
         return Error::InvalidTargetRange if closest_def.nil?
 

lib/ruby_lsp/requests/code_actions.rb

@@ -37,7 +37,7 @@ def provider
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           range: T::Hash[Symbol, T.untyped],
           context: T::Hash[Symbol, T.untyped],
         ).void

lib/ruby_lsp/requests/completion.rb

@@ -46,7 +46,7 @@ def provider
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           global_state: GlobalState,
           params: T::Hash[Symbol, T.untyped],
           sorbet_level: RubyDocument::SorbetLevel,
@@ -60,7 +60,7 @@ def initialize(document, global_state, params, sorbet_level, dispatcher)
         # Completion always receives the position immediately after the character that was just typed. Here we adjust it
         # back by 1, so that we find the right node
         char_position = document.create_scanner.find_char_position(params[:position]) - 1
-        node_context = document.locate(
+        node_context = RubyDocument.locate(
           document.parse_result.value,
           char_position,
           node_types: [

lib/ruby_lsp/requests/definition.rb

@@ -38,7 +38,7 @@ class Definition < Request
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           global_state: GlobalState,
           position: T::Hash[Symbol, T.untyped],
           dispatcher: Prism::Dispatcher,

lib/ruby_lsp/requests/diagnostics.rb

@@ -32,7 +32,7 @@ def provider
         end
       end
 
-      sig { params(global_state: GlobalState, document: Document).void }
+      sig { params(global_state: GlobalState, document: RubyDocument).void }
       def initialize(global_state, document)
         super()
         @active_linters = T.let(global_state.active_linters, T::Array[Support::Formatter])

lib/ruby_lsp/requests/document_highlight.rb

@@ -29,7 +29,7 @@ class DocumentHighlight < Request
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           position: T::Hash[Symbol, T.untyped],
           dispatcher: Prism::Dispatcher,
         ).void

lib/ruby_lsp/requests/formatting.rb

@@ -40,7 +40,7 @@ def provider
         end
       end
 
-      sig { params(global_state: GlobalState, document: Document).void }
+      sig { params(global_state: GlobalState, document: RubyDocument).void }
       def initialize(global_state, document)
         super()
         @document = document

lib/ruby_lsp/requests/hover.rb

@@ -32,7 +32,7 @@ def provider
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           global_state: GlobalState,
           position: T::Hash[Symbol, T.untyped],
           dispatcher: Prism::Dispatcher,

lib/ruby_lsp/requests/inlay_hints.rb

@@ -52,7 +52,7 @@ def provider
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           range: T::Hash[Symbol, T.untyped],
           hints_configuration: RequestConfig,
           dispatcher: Prism::Dispatcher,

lib/ruby_lsp/requests/on_type_formatting.rb

@@ -42,7 +42,7 @@ def provider
 
       sig do
         params(
-          document: Document,
+          document: RubyDocument,
           position: T::Hash[Symbol, T.untyped],
           trigger_character: String,
           client_name: String,

lib/ruby_lsp/requests/prepare_type_hierarchy.rb

@@ -35,7 +35,7 @@ def provider
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           index: RubyIndexer::Index,
           position: T::Hash[Symbol, T.untyped],
         ).void

lib/ruby_lsp/requests/selection_ranges.rb

@@ -23,7 +23,8 @@ module Requests
     class SelectionRanges < Request
       extend T::Sig
       include Support::Common
-      sig { params(document: Document).void }
+
+      sig { params(document: T.any(RubyDocument, ERBDocument)).void }
       def initialize(document)
         super()
         @document = document

lib/ruby_lsp/requests/show_syntax_tree.rb

@@ -20,7 +20,7 @@ module Requests
     class ShowSyntaxTree < Request
       extend T::Sig
 
-      sig { params(document: Document, range: T.nilable(T::Hash[Symbol, T.untyped])).void }
+      sig { params(document: RubyDocument, range: T.nilable(T::Hash[Symbol, T.untyped])).void }
       def initialize(document, range)
         super()
         @document = document

lib/ruby_lsp/requests/signature_help.rb

@@ -41,7 +41,7 @@ def provider
 
       sig do
         params(
-          document: Document,
+          document: T.any(RubyDocument, ERBDocument),
           global_state: GlobalState,
           position: T::Hash[Symbol, T.untyped],
           context: T.nilable(T::Hash[Symbol, T.untyped]),

lib/ruby_lsp/requests/support/formatter.rb

@@ -10,13 +10,13 @@ module Formatter
 
         interface!
 
-        sig { abstract.params(uri: URI::Generic, document: Document).returns(T.nilable(String)) }
+        sig { abstract.params(uri: URI::Generic, document: RubyDocument).returns(T.nilable(String)) }
         def run_formatting(uri, document); end
 
         sig do
           abstract.params(
             uri: URI::Generic,
-            document: Document,
+            document: RubyDocument,
           ).returns(T.nilable(T::Array[Interface::Diagnostic]))
         end
         def run_diagnostic(uri, document); end

lib/ruby_lsp/requests/support/rubocop_diagnostic.rb

@@ -31,7 +31,7 @@ class RuboCopDiagnostic
 
         # TODO: avoid passing document once we have alternative ways to get at
         # encoding and file source
-        sig { params(document: Document, offense: RuboCop::Cop::Offense, uri: URI::Generic).void }
+        sig { params(document: RubyDocument, offense: RuboCop::Cop::Offense, uri: URI::Generic).void }
         def initialize(document, offense, uri)
           @document = document
           @offense = offense

lib/ruby_lsp/requests/support/rubocop_formatter.rb

@@ -17,7 +17,7 @@ def initialize
           @format_runner = T.let(RuboCopRunner.new("-a"), RuboCopRunner)
         end
 
-        sig { override.params(uri: URI::Generic, document: Document).returns(T.nilable(String)) }
+        sig { override.params(uri: URI::Generic, document: RubyDocument).returns(T.nilable(String)) }
         def run_formatting(uri, document)
           filename = T.must(uri.to_standardized_path || uri.opaque)
 
@@ -29,7 +29,7 @@ def run_formatting(uri, document)
         sig do
           override.params(
             uri: URI::Generic,
-            document: Document,
+            document: RubyDocument,
           ).returns(T.nilable(T::Array[Interface::Diagnostic]))
         end
         def run_diagnostic(uri, document)

lib/ruby_lsp/requests/support/syntax_tree_formatter.rb

@@ -29,7 +29,7 @@ def initialize
             )
         end
 
-        sig { override.params(uri: URI::Generic, document: Document).returns(T.nilable(String)) }
+        sig { override.params(uri: URI::Generic, document: RubyDocument).returns(T.nilable(String)) }
         def run_formatting(uri, document)
           path = uri.to_standardized_path
           return if path && @options.ignore_files.any? { |pattern| File.fnmatch?("*/#{pattern}", path) }
@@ -40,7 +40,7 @@ def run_formatting(uri, document)
         sig do
           override.params(
             uri: URI::Generic,
-            document: Document,
+            document: RubyDocument,
           ).returns(T.nilable(T::Array[Interface::Diagnostic]))
         end
         def run_diagnostic(uri, document)