Add Prism::Source#line_to_byte_offset and replace direct accesses to offsets

Author: eregon

lib/prism/parse_result.rb

@@ -76,6 +76,15 @@ def slice(byte_offset, length)
       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)
@@ -145,9 +154,10 @@ def code_units_column(byte_offset, encoding)
 
     # Returns the byte offset for a given line number and column number
     def line_and_character_column_to_byte_offset(line, column)
-      line_start = offsets[line - 1]
-      line_end = offsets[line]
-      byte_column = (@source.byteslice(line_start, line_end) or raise)[0...column]&.bytesize #: Integer
+      line_start = line_to_byte_offset(line)
+      line_end = offsets[line + 1 - @start_line] || source.bytesize
+      byteslice = @source.byteslice(line_start, line_end) or raise ArgumentError, "line #{line} is out of range"
+      byte_column = (byteslice[0...column] or raise).bytesize
       line_start + byte_column
     end
 
@@ -284,8 +294,7 @@ def code_units_column(byte_offset, encoding)
     # Specialized version of `line_and_character_column_to_byte_offset`
     # which does not need to access the source String
     def line_and_character_column_to_byte_offset(line, column)
-      line_start = offsets[line - 1]
-      line_start + column
+      line_to_byte_offset(line) + column
     end
   end
 

sig/prism/parse_result.rbs

@@ -14,6 +14,7 @@ module Prism
     def encoding: () -> Encoding
     def lines: () -> Array[String]
     def slice: (Integer byte_offset, Integer length) -> String
+    def line_to_byte_offset: (Integer line) -> Integer
     def line: (Integer byte_offset) -> Integer
     def line_start: (Integer byte_offset) -> Integer
     def line_end: (Integer byte_offset) -> Integer

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

test/prism/ruby/source_test.rb

@@ -0,0 +1,41 @@
+# 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)
+      assert_raise(ArgumentError.new("line 5 is out of range")) { source.line_to_byte_offset(5) }
+      assert_raise(ArgumentError.new("line 0 is out of range")) { source.line_to_byte_offset(0) }
+      assert_raise(ArgumentError.new("line -1 is out of range")) { source.line_to_byte_offset(-1) }
+    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)
+      assert_raise(ArgumentError.new("line 15 is out of range")) { source.line_to_byte_offset(15) }
+      assert_raise(ArgumentError.new("line 10 is out of range")) { source.line_to_byte_offset(10) }
+      assert_raise(ArgumentError.new("line 9 is out of range")) { source.line_to_byte_offset(9) }
+    end
+  end
+end