Bound resource use when iterating ranges (#2117)

lib/liquid.rb

@@ -83,6 +83,7 @@ require 'liquid/resource_limits'
 require 'liquid/expression'
 require 'liquid/template'
 require 'liquid/condition'
+require 'liquid/range_slice'
 require 'liquid/utils'
 require 'liquid/tokenizer'
 require 'liquid/parse_context'

lib/liquid/range_slice.rb

@@ -0,0 +1,44 @@
+# frozen_string_literal: true
+
+module Liquid
+  class RangeSlice
+    attr_reader :length
+
+    def initialize(range, from, to, resource_limits)
+      range_length = range.end - range.begin
+      range_length += 1 unless range.exclude_end?
+      range_length = 0 if range_length.negative?
+
+      start = [from, 0].max
+      finish = [to || range_length, range_length].min
+
+      @first = range.begin + start
+      @length = [finish - start, 0].max
+      @direction = 1
+      @resource_limits = resource_limits
+    end
+
+    def empty?
+      @length.zero?
+    end
+
+    def each
+      return enum_for(:each) unless block_given?
+
+      value = @first
+      @length.times do
+        @resource_limits.increment_render_score(1)
+        yield value
+        value += @direction
+      end
+    end
+
+    def reverse!
+      unless empty?
+        @first += @direction * (@length - 1)
+        @direction = -@direction
+      end
+      self
+    end
+  end
+end

lib/liquid/tags/for.rb

@@ -130,7 +130,6 @@ module Liquid
       end
 
       collection = context.evaluate(@collection_name)
-      collection = collection.to_a if collection.is_a?(Range)
 
       limit_value = context.evaluate(@limit)
       to = if limit_value.nil?
@@ -139,7 +138,9 @@ module Liquid
         Utils.to_integer(limit_value) + from
       end
 
-      segment = Utils.slice_collection(collection, from, to)
+      segment = Utils.slice_collection_for_iteration(
+        collection, from, to, context.resource_limits, use_range_to_a: true
+      )
       segment.reverse! if @reversed
 
       offsets[@name] = from + segment.length

lib/liquid/tags/table_row.rb

@@ -85,12 +85,13 @@ module Liquid
       from = @attributes.key?('offset') ? to_integer(context.evaluate(@attributes['offset'])) : 0
       to = @attributes.key?('limit') ? from + to_integer(context.evaluate(@attributes['limit'])) : nil
 
-      collection = Utils.slice_collection(collection, from, to)
+      collection = Utils.slice_collection_for_iteration(collection, from, to, context.resource_limits, allow_endless: true)
       length     = collection.length
 
       cols = @attributes.key?('cols') ? to_integer(context.evaluate(@attributes['cols'])) : length
 
       output << "<tr class=\"row1\">\n"
+      context.resource_limits.increment_write_score(output)
       context.stack do
         tablerowloop = Liquid::TablerowloopDrop.new(length, cols)
         context['tablerowloop'] = tablerowloop
@@ -101,6 +102,7 @@ module Liquid
           output << "<td class=\"col#{tablerowloop.col}\">"
           super
           output << '</td>'
+          context.resource_limits.increment_write_score(output)
 
           # Handle any interrupts if they exist.
           if context.interrupt?
@@ -110,6 +112,7 @@ module Liquid
 
           if tablerowloop.col_last && !tablerowloop.last
             output << "</tr>\n<tr class=\"row#{tablerowloop.row + 1}\">"
+            context.resource_limits.increment_write_score(output)
           end
 
           tablerowloop.send(:increment!)
@@ -117,6 +120,7 @@ module Liquid
       end
 
       output << "</tr>\n"
+      context.resource_limits.increment_write_score(output)
       output
     end
 

lib/liquid/utils.rb

@@ -13,6 +13,26 @@ module Liquid
       end
     end
 
+    # This is intentionally separate from slice_collection, whose Array-returning
+    # behavior is used outside of the iteration tags.
+    def self.slice_collection_for_iteration(
+      collection, from, to, resource_limits, allow_endless: false, use_range_to_a: false
+    )
+      if integer_range?(collection)
+        RangeSlice.new(collection, from, to, resource_limits)
+      elsif collection.is_a?(Range)
+        if use_range_to_a && range_method_overridden?(collection, :to_a)
+          # For historically honored custom Range#to_a. Charge the resulting
+          # selection before buffering it, just as for a custom #each.
+          slice_collection_for_iteration_using_each(collection.to_a, from, to, resource_limits)
+        else
+          slice_range_using_each(collection, from, to, resource_limits, allow_endless: allow_endless)
+        end
+      else
+        slice_collection(collection, from, to)
+      end
+    end
+
     def self.slice_collection_using_each(collection, from, to)
       segments = []
       index    = 0
@@ -38,6 +58,55 @@ module Liquid
       segments
     end
 
+    # Arithmetic slicing must not bypass a Range subclass's custom #each.
+    def self.integer_range?(collection)
+      collection.instance_of?(Range) && collection.begin.is_a?(Integer) && collection.end.is_a?(Integer)
+    end
+    private_class_method :integer_range?
+
+    # Preserve support for Ruby-supplied string ranges and custom Range#each.
+    # Their selected length cannot be inferred from integer bounds, but the tags
+    # need it before rendering for loop metadata, continuation offsets, and columns.
+    # Buffer the selection so we do not have to replay a potentially custom iterator.
+    def self.slice_range_using_each(collection, from, to, resource_limits, allow_endless:)
+      # TableRow historically accepted an endless subclass when its custom #each
+      # was finite, while For historically raised through Range#to_a.
+      if collection.end.nil? && !(allow_endless && (!to.nil? || range_method_overridden?(collection, :each)))
+        raise RangeError, "cannot convert endless range to an array"
+      end
+      if collection.begin.nil? && !range_method_overridden?(collection, :each)
+        raise TypeError, "can't iterate from NilClass"
+      end
+
+      slice_collection_for_iteration_using_each(collection, from, to, resource_limits)
+    end
+    private_class_method :slice_range_using_each
+
+    # Custom Range#each can make a nominally beginless range finite; standard
+    # beginless ranges were rejected before reaching this budgeted traversal.
+    def self.slice_collection_for_iteration_using_each(collection, from, to, resource_limits)
+      return [] if to && to <= from
+
+      segments = []
+      index = 0
+      collection.each do |item|
+        break if to && to <= index
+
+        # Charge preparation, including skipped offsets, before buffering; checking
+        # only while rendering the buffered values would leave this work unbudgeted.
+        resource_limits.increment_render_score(1)
+        segments << item if from <= index
+        index += 1
+      end
+      segments
+    end
+    private_class_method :slice_collection_for_iteration_using_each
+
+    def self.range_method_overridden?(collection, method_name)
+      collection.method(method_name).owner != Range.instance_method(method_name).owner
+    end
+    private_class_method :range_method_overridden?
+
     def self.to_integer(num)
       return num if num.is_a?(Integer)
       num = num.to_s