From 44ffa06df777f26e19d808e56e5de9436b5915b7 Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Fri, 26 Jul 2024 09:44:15 -0400
Subject: [PATCH] Add the ability to locate the first node within a range

---
 lib/ruby_lsp/document.rb   | 42 ++++++++++++++++++++++++++++++++++++++
 test/ruby_document_test.rb | 23 +++++++++++++++++++++
 2 files changed, 65 insertions(+)

diff --git a/lib/ruby_lsp/document.rb b/lib/ruby_lsp/document.rb
index e57ad3de8..bbd19b055 100644
--- a/lib/ruby_lsp/document.rb
+++ b/lib/ruby_lsp/document.rb
@@ -223,6 +223,48 @@ 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)])
+      closest = T.let(nil, T.nilable(Prism::Node))
+
+      until queue.empty? || closest
+        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 desired_range.cover?(loc.start_offset...loc.end_offset)
+
+        # If the node's start character is already past the position, then we should've found the closest node
+        # already
+        break if end_position < loc.start_offset
+
+        # 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 }
+
+        closest = candidate
+      end
+
+      closest
+    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 1aca63b30..7571cd6a0 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)