Update assign_score during capturing

To stop long captures before they grow the heap more then they should.

lib/liquid/block_body.rb

@@ -170,7 +170,7 @@ module Liquid
         end
         idx += 1
 
-        context.resource_limits.check_render_length(output.bytesize)
+        context.resource_limits.increment_write_score(output)
       end
 
       output

lib/liquid/resource_limits.rb

@@ -22,8 +22,16 @@ module Liquid
       raise_limits_reached if @assign_score_limit && @assign_score > @assign_score_limit
     end
 
-    def check_render_length(output_byte_size)
-      raise_limits_reached if @render_length_limit && output_byte_size > @render_length_limit
+    # update either render_length or assign_score based on whether or not the writes are captured
+    def increment_write_score(output)
+      if (last_captured = @last_capture_length)
+        captured = output.bytesize
+        increment = captured - last_captured
+        @last_capture_length = captured
+        increment_assign_score(increment)
+      elsif @render_length_limit && output.bytesize > @render_length_limit
+        raise_limits_reached
+      end
     end
 
     def raise_limits_reached
@@ -37,7 +45,18 @@ module Liquid
 
     def reset
       @reached_limit = false
+      @last_capture_length = nil
       @render_score = @assign_score = 0
     end
+
+    def with_capture
+      old_capture_length = @last_capture_length
+      begin
+        @last_capture_length = 0
+        yield
+      ensure
+        @last_capture_length = old_capture_length
+      end
+    end
   end
 end

lib/liquid/tags/capture.rb

@@ -25,9 +25,10 @@ module Liquid
     end
 
     def render_to_output_buffer(context, output)
-      capture_output = render(context)
-      context.scopes.last[@to] = capture_output
-      context.resource_limits.increment_assign_score(capture_output.bytesize)
+      context.resource_limits.with_capture do
+        capture_output = render(context)
+        context.scopes.last[@to] = capture_output
+      end
       output
     end