Merge pull request #2950 from ksss/buffer-character-offset-cache

Speed up RBS::InlineParser on large sources with non-ASCII characters

Author: soutaro

lib/rbs/ast/ruby/comment_block.rb

@@ -15,6 +15,7 @@ def initialize(source_buffer, comments)
           prefix_str = "# "
 
           ranges = [] #: Array[Range[Integer]]
+          byte_ranges = [] #: Array[Range[Integer]]
 
           comments.each do |comment|
             tuple = [comment, 2] #: [Prism::Comment, Integer]
@@ -25,12 +26,13 @@ def initialize(source_buffer, comments)
 
             offsets << tuple
 
-            start_char = comment.location.start_character_offset + tuple[1]
-            end_char = comment.location.end_character_offset
+            start_char = source_buffer.character_offset(comment.location.start_offset) + tuple[1]
+            end_char = source_buffer.character_offset(comment.location.end_offset)
             ranges << (start_char ... end_char)
+            byte_ranges << ((comment.location.start_offset + tuple[1]) ... comment.location.end_offset)
           end
 
-          @comment_buffer = source_buffer.sub_buffer(lines: ranges)
+          @comment_buffer = source_buffer.sub_buffer(lines: ranges, byte_lines_hint: byte_ranges)
         end
 
         def leading?
@@ -53,7 +55,7 @@ def end_line
 
         def line_starts
           offsets.map do |comment, prefix_size|
-            comment.location.start_character_offset + prefix_size
+            comment_buffer.character_offset(comment.location.start_offset) + prefix_size
           end
         end
 

lib/rbs/ast/ruby/helpers/location_helper.rb

@@ -6,7 +6,7 @@ module Ruby
       module Helpers
         module LocationHelper
           def rbs_location(location)
-            Location.new(buffer, location.start_character_offset, location.end_character_offset)
+            Location.new(buffer, buffer.character_offset(location.start_offset), buffer.character_offset(location.end_offset))
           end
         end
       end

lib/rbs/buffer.rb

@@ -87,25 +87,62 @@ def inspect
       "#<RBS::Buffer:#{__id__} @name=#{name}, @content=#{content.bytesize} bytes, @lines=#{ranges.size} lines,>"
     end
 
+    def character_offset(byte_offset)
+      top = top_buffer
+      return top.character_offset(byte_offset) unless top.equal?(self)
+
+      keys, vals = (@character_offset_cache ||= [[0], [0]])
+
+      idx = keys.bsearch_index { |k| k > byte_offset }
+      lo = idx ? idx - 1 : keys.size - 1
+
+      base_byte = keys[lo]
+      base_char = vals[lo]
+      delta = byte_offset - base_byte
+      return base_char if delta == 0
+
+      result = base_char + (content.byteslice(base_byte, delta) or raise).length
+
+      if base_byte == keys[-1]
+        keys << byte_offset
+        vals << result
+      end
+
+      result
+    end
+
     def rbs_location(location, loc2=nil)
+      top = top_buffer
       if loc2
-        Location.new(self.top_buffer, location.start_character_offset, loc2.end_character_offset)
+        Location.new(top, character_offset(location.start_offset), character_offset(loc2.end_offset))
       else
-        Location.new(self.top_buffer, location.start_character_offset, location.end_character_offset)
+        Location.new(top, character_offset(location.start_offset), character_offset(location.end_offset))
       end
     end
 
-    def sub_buffer(lines:)
+    def sub_buffer(lines:, byte_lines_hint: nil)
       buf = +""
-      lines.each_with_index do |range, index|
-        start_pos = range.begin
-        end_pos = range.end
-        slice = content[start_pos...end_pos] or raise
-        if slice.include?("\n")
-          raise "Line #{index + 1} cannot contain newline character."
+
+      if byte_lines_hint
+        byte_lines_hint.each_with_index do |range, index|
+          slice = content.byteslice(range.begin, range.end - range.begin) or raise
+          if slice.include?("\n")
+            raise "Line #{index + 1} cannot contain newline character."
+          end
+          buf << slice
+          buf << "\n"
+        end
+      else
+        lines.each_with_index do |range, index|
+          start_pos = range.begin
+          end_pos = range.end
+          slice = content[start_pos...end_pos] or raise
+          if slice.include?("\n")
+            raise "Line #{index + 1} cannot contain newline character."
+          end
+          buf << slice
+          buf << "\n"
         end
-        buf << slice
-        buf << "\n"
       end
 
       buf.chomp!

sig/buffer.rbs

@@ -60,6 +60,19 @@ module RBS
     def rbs_location: (Prism::Location) -> Location
                     | (Prism::Location, Prism::Location) -> Location
 
+    # Translate a byte offset (into the top buffer's source) to a character offset.
+    #
+    # Resolution is delegated to the top buffer, which keeps a sparse cache of
+    # resolved (byte, char) pairs so successive calls scan only the delta from the
+    # nearest cached pair. Amortizes to O(content_size) across all calls; a single
+    # call is one byteslice + length.
+    #
+    def character_offset: (Integer byte_offset) -> Integer
+
+    # Sparse cache backing `#character_offset`: a pair `[byte_keys, char_values]`
+    # kept in ascending byte order to support binary search.
+    @character_offset_cache: [Array[Integer], Array[Integer]]?
+
     # Construct a buffer from substrings of this buffer.
     #
     # The returned buffer contains lines from given ranges.
@@ -75,7 +88,12 @@ module RBS
     # buffer.sub_buffer(lines: [5..7])          # => Raises an error because the range contains newline
     # ```
     #
-    %a{pure} def sub_buffer: (lines: Array[Range[Integer]]) -> Buffer
+    # `byte_lines_hint:` is an optional performance hint: byte ranges corresponding
+    # to `lines:`. When provided, slicing uses `byteslice` (O(slice_size) per line)
+    # instead of `content[char_range]`, which on a multi-byte string is O(content_size)
+    # per call. Result is identical either way.
+    #
+    %a{pure} def sub_buffer: (lines: Array[Range[Integer]], ?byte_lines_hint: Array[Range[Integer]]?) -> Buffer
 
     %a{pure} def parent_buffer: () -> Buffer?