clean up write recording

lib/liquid/block_body.rb

@@ -218,6 +218,9 @@ module Liquid
     def render_to_output_buffer(context, output)
       freeze unless frozen?
 
+      context.registers.static[:render_scores] ||= []
+      context.registers.static[:render_scores] << @nodelist.length
+
       context.resource_limits.increment_render_score(@nodelist.length)
 
       idx = 0
@@ -233,6 +236,18 @@ module Liquid
         end
         idx += 1
 
+        context.registers.static[:write_scores] ||= []
+
+        if (last_captured = context.resource_limits.last_capture_length)
+          captured = output.bytesize
+          increment = captured - last_captured
+          context.registers.static[:assign_scores] ||= []
+          context.registers.static[:assign_scores] << ['cap', increment]
+        else
+          context.registers.static[:write_scores] ||= []
+          context.registers.static[:write_scores] << (output.bytesize)
+        end
+
         context.resource_limits.increment_write_score(output)
       end
 

lib/liquid/resource_limits.rb

@@ -2,7 +2,7 @@
 
 module Liquid
   class ResourceLimits
-    attr_accessor :render_length_limit, :render_score_limit, :assign_score_limit
+    attr_accessor :render_length_limit, :render_score_limit, :assign_score_limit, :last_capture_length
     attr_reader :render_score, :assign_score
 
     def initialize(limits)

lib/liquid/tags/assign.rb

@@ -36,6 +36,10 @@ module Liquid
     def render_to_output_buffer(context, output)
       val = @from.render(context)
       context.scopes.last[@to] = val
+
+      context.registers.static[:assign_scores] ||= []
+      context.registers.static[:assign_scores] << [@to, assign_score_of(val)]
+
       context.resource_limits.increment_assign_score(assign_score_of(val))
       output
     end
@@ -47,14 +51,16 @@ module Liquid
     private
 
     def assign_score_of(val)
-      if val.instance_of?(String)
+      return 0 if val.nil?
+
+      if val.is_a?(String)
         val.bytesize
-      elsif val.instance_of?(Array)
+      elsif val.is_a?(Array)
         sum = 1
         # Uses #each to avoid extra allocations.
         val.each { |child| sum += assign_score_of(child) }
         sum
-      elsif val.instance_of?(Hash)
+      elsif val.is_a?(Hash)
         sum = 1
         val.each do |key, entry_value|
           sum += assign_score_of(key)