From 3604aa15e70002bd5e1b9bc34a2e3b40041fbf3c Mon Sep 17 00:00:00 2001
From: Kevin Newton <kddnewton@gmail.com>
Date: Tue, 23 Apr 2024 15:19:48 -0400
Subject: [PATCH] srange_find should only look on current line

---
 lib/prism/translation/parser/compiler.rb | 20 ++++++++++++--------
 1 file changed, 12 insertions(+), 8 deletions(-)

diff --git a/lib/prism/translation/parser/compiler.rb b/lib/prism/translation/parser/compiler.rb
index 22299d29ce7..a4aaa41d6f7 100644
--- a/lib/prism/translation/parser/compiler.rb
+++ b/lib/prism/translation/parser/compiler.rb
@@ -839,7 +839,7 @@ def visit_in_node(node)
             token(node.in_loc),
             pattern,
             guard,
-            srange_find(node.pattern.location.end_offset, node.statements&.location&.start_offset || node.location.end_offset, [";", "then"]),
+            srange_find(node.pattern.location.end_offset, node.statements&.location&.start_offset, [";", "then"]),
             visit(node.statements)
           )
         end
@@ -1679,7 +1679,7 @@ def visit_unless_node(node)
         end
 
         # until foo; bar end
-        # ^^^^^^^^^^^^^^^^^
+        # ^^^^^^^^^^^^^^^^^^
         #
         # bar until foo
         # ^^^^^^^^^^^^^
@@ -1712,7 +1712,7 @@ def visit_when_node(node)
             if node.then_keyword_loc
               token(node.then_keyword_loc)
             else
-              srange_find(node.conditions.last.location.end_offset, node.statements&.location&.start_offset || (node.conditions.last.location.end_offset + 1), [";"])
+              srange_find(node.conditions.last.location.end_offset, node.statements&.location&.start_offset, [";"])
             end,
             visit(node.statements)
           )
@@ -1871,12 +1871,16 @@ def srange_offsets(start_offset, end_offset)
 
         # Constructs a new source range by finding the given tokens between the
         # given start offset and end offset. If the needle is not found, it
-        # returns nil.
+        # returns nil. Importantly it does not search past newlines or comments.
+        #
+        # Note that end_offset is allowed to be nil, in which case this will
+        # search until the end of the string.
         def srange_find(start_offset, end_offset, tokens)
-          tokens.find do |token|
-            next unless (index = source_buffer.source.byteslice(start_offset...end_offset).index(token))
-            offset = start_offset + index
-            return [token, Range.new(source_buffer, offset_cache[offset], offset_cache[offset + token.length])]
+          if (match = source_buffer.source.byteslice(start_offset...end_offset).match(/(\s*)(#{tokens.join("|")})/))
+            _, whitespace, token = *match
+            token_offset = start_offset + whitespace.bytesize
+
+            [token, Range.new(source_buffer, offset_cache[token_offset], offset_cache[token_offset + token.bytesize])]
           end
         end