Lazily build offsets from a packed array

Author: kddnewton

ext/prism/extension.c

@@ -802,16 +802,14 @@ parse_lex_input(const uint8_t *input, size_t input_length, const pm_options_t *o
 
     pm_node_t *node = pm_parse(parser);
 
-    // Here we need to update the Source object to have the correct
-    // encoding for the source string and the correct newline offsets.
-    // We do it here because we've already created the Source object and given
-    // it over to all of the tokens, and both of these are only set after pm_parse().
+    /* Update the Source object with the correct encoding and line offsets,
+     * which are only available after pm_parse() completes. */
     rb_encoding *encoding = rb_enc_find(pm_parser_encoding_name(parser));
     rb_enc_associate(source_string, encoding);
 
     const pm_line_offset_list_t *line_offsets = pm_parser_line_offsets(parser);
     for (size_t index = 0; index < line_offsets->size; index++) {
-        rb_ary_push(offsets, ULONG2NUM(line_offsets->offsets[index]));
+        rb_ary_store(offsets, (long) index, ULONG2NUM(line_offsets->offsets[index]));
     }
 
     if (pm_options_freeze(options)) {

lib/prism/parse_result.rb

@@ -58,16 +58,26 @@ def self.for(source, start_line, offsets)
     # The line number where this source starts.
     attr_reader :start_line #: Integer
 
-    # The list of newline byte offsets in the source code.
-    attr_reader :offsets #: Array[Integer]
+    # The list of newline byte offsets in the source code. When initialized from
+    # the C extension, this may be a packed binary string of uint32_t values
+    # that is lazily unpacked on first access.
+    #--
+    #: () -> Array[Integer]
+    def offsets
+      offsets = @offsets
+      return offsets if offsets.is_a?(Array)
+      @offsets = offsets.unpack("L*")
+    end
 
-    # Create a new source object with the given source code.
+    # Create a new source object with the given source code. The offsets
+    # parameter can be either an Array of Integer byte offsets or a packed
+    # binary string of uint32_t values (from the C extension).
     #--
-    #: (String source, Integer start_line, Array[Integer] offsets) -> void
+    #: (String source, Integer start_line, Array[Integer] | String offsets) -> void
     def initialize(source, start_line, offsets)
       @source = source
-      @start_line = start_line # set after parsing is done
-      @offsets = offsets # set after parsing is done
+      @start_line = start_line
+      @offsets = offsets
     end
 
     # Replace the value of start_line with the given value.
@@ -81,7 +91,7 @@ def replace_start_line(start_line)
     #--
     #: (Array[Integer] offsets) -> void
     def replace_offsets(offsets)
-      @offsets.replace(offsets)
+      @offsets = offsets
     end
 
     # Returns the encoding of the source code, which is set by parameters to the

rbi/generated/prism/parse_result.rbi

@@ -81,11 +81,7 @@ pm_source_new(const pm_parser_t *parser, rb_encoding *encoding, bool freeze) {
     VALUE source_string = rb_enc_str_new((const char *) start, pm_parser_end(parser) - start, encoding);
 
     const pm_line_offset_list_t *line_offsets = pm_parser_line_offsets(parser);
-    VALUE offsets = rb_ary_new_capa(line_offsets->size);
-
-    for (size_t index = 0; index < line_offsets->size; index++) {
-        rb_ary_push(offsets, ULONG2NUM(line_offsets->offsets[index]));
-    }
+    VALUE offsets = rb_str_new((const char *) line_offsets->offsets, line_offsets->size * sizeof(uint32_t));
 
     if (freeze) {
         rb_obj_freeze(source_string);