diff --git a/lib/ruby_lsp/document.rb b/lib/ruby_lsp/document.rb index e57ad3de83..f17d439a39 100644 --- a/lib/ruby_lsp/document.rb +++ b/lib/ruby_lsp/document.rb @@ -223,6 +223,40 @@ def locate(node, char_position, node_types: []) NodeContext.new(closest, parent, nesting_nodes, call_node) end + sig do + params( + range: T::Hash[Symbol, T.untyped], + node_types: T::Array[T.class_of(Prism::Node)], + ).returns(T.nilable(Prism::Node)) + end + def locate_first_within_range(range, node_types: []) + scanner = create_scanner + start_position = scanner.find_char_position(range[:start]) + end_position = scanner.find_char_position(range[:end]) + desired_range = (start_position...end_position) + queue = T.let(@parse_result.value.child_nodes.compact, T::Array[T.nilable(Prism::Node)]) + + 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 + + if desired_range.cover?(loc.start_offset...loc.end_offset) && + (node_types.empty? || node_types.any? { |type| candidate.class == type }) + return candidate + end + end + end + sig { returns(SorbetLevel) } def sorbet_level sigil = parse_result.magic_comments.find do |comment| diff --git a/test/ruby_document_test.rb b/test/ruby_document_test.rb index 1aca63b307..7571cd6a00 100644 --- a/test/ruby_document_test.rb +++ b/test/ruby_document_test.rb @@ -716,6 +716,29 @@ def qux assert_equal("qux", node_context.surrounding_method) end + def test_locate_first_within_range + document = RubyLsp::RubyDocument.new(source: +<<~RUBY, version: 1, uri: URI("file:///foo/bar.rb")) + method_call(other_call).each do |a| + nested_call(fourth_call).each do |b| + end + end + RUBY + + target = document.locate_first_within_range( + { start: { line: 0, character: 0 }, end: { line: 3, character: 3 } }, + node_types: [Prism::CallNode], + ) + + assert_equal("each", T.cast(target, Prism::CallNode).message) + + target = document.locate_first_within_range( + { start: { line: 1, character: 2 }, end: { line: 2, character: 5 } }, + node_types: [Prism::CallNode], + ) + + assert_equal("each", T.cast(target, Prism::CallNode).message) + end + private def assert_error_edit(actual, error_range)