[ruby/prism] Add Prism::Source#line_to_byte_offset and replace direct accesses to offsetsHEADmaster

https://github.com/ruby/prism/commit/ff81a29ba5

3 files changed, 57 insertions, 2 deletions

diff --git a/lib/prism/parse_result.rb b/lib/prism/parse_result.rb
index 3570af136a..12d19da562 100644
--- a/lib/prism/parse_result.rb
+++ b/lib/prism/parse_result.rb
@@ -76,6 +76,15 @@ module Prism
       source.byteslice(byte_offset, length) or raise
     end
 
+    # Converts the line number to a byte offset corresponding to the start of that line
+    def line_to_byte_offset(line)
+      l = line - @start_line
+      if l < 0 || l >= offsets.size
+        raise ArgumentError, "line #{line} is out of range"
+      end
+      offsets[l]
+    end
+
     # Binary search through the offsets to find the line number for the given
     # byte offset.
     def line(byte_offset)
diff --git a/prism/templates/lib/prism/node.rb.erb b/prism/templates/lib/prism/node.rb.erb
index 181842e230..8225bfb328 100644
--- a/prism/templates/lib/prism/node.rb.erb
+++ b/prism/templates/lib/prism/node.rb.erb
@@ -184,8 +184,7 @@ module Prism
       queue = [self] #: Array[Prism::node]
       result = [] #: Array[Prism::node]
 
-      line_offset = source.offsets[line - 1] or raise
-      search_offset = line_offset + column
+      search_offset = source.line_to_byte_offset(line) + column
 
       while (node = queue.shift)
         result << node
diff --git a/test/prism/ruby/source_test.rb b/test/prism/ruby/source_test.rb
new file mode 100644
index 0000000000..afd2825765
--- /dev/null
+++ b/test/prism/ruby/source_test.rb
@@ -0,0 +1,47 @@
+# frozen_string_literal: true
+
+require_relative "../test_helper"
+
+module Prism
+  class SourceTest < TestCase
+    def test_line_to_byte_offset
+      parse_result = Prism.parse(<<~SRC)
+      abcd
+      efgh
+      ijkl
+      SRC
+      source = parse_result.source
+
+      assert_equal 0, source.line_to_byte_offset(1)
+      assert_equal 5, source.line_to_byte_offset(2)
+      assert_equal 10, source.line_to_byte_offset(3)
+      assert_equal 15, source.line_to_byte_offset(4)
+      e = assert_raise(ArgumentError) { source.line_to_byte_offset(5) }
+      assert_equal "line 5 is out of range", e.message
+      e = assert_raise(ArgumentError) { source.line_to_byte_offset(0) }
+      assert_equal "line 0 is out of range", e.message
+      e = assert_raise(ArgumentError) { source.line_to_byte_offset(-1) }
+      assert_equal "line -1 is out of range", e.message
+    end
+
+    def test_line_to_byte_offset_with_start_line
+      parse_result = Prism.parse(<<~SRC, line: 11)
+      abcd
+      efgh
+      ijkl
+      SRC
+      source = parse_result.source
+
+      assert_equal 0, source.line_to_byte_offset(11)
+      assert_equal 5, source.line_to_byte_offset(12)
+      assert_equal 10, source.line_to_byte_offset(13)
+      assert_equal 15, source.line_to_byte_offset(14)
+      e = assert_raise(ArgumentError) { source.line_to_byte_offset(15) }
+      assert_equal "line 15 is out of range", e.message
+      e = assert_raise(ArgumentError) { source.line_to_byte_offset(10) }
+      assert_equal "line 10 is out of range", e.message
+      e = assert_raise(ArgumentError) { source.line_to_byte_offset(9) }
+      assert_equal "line 9 is out of range", e.message
+    end
+  end
+end