From b185421405583b74801f147f1f98e43b9108e8bf Mon Sep 17 00:00:00 2001
From: Youssef Makboul <youssef.makboul@shopify.com>
Date: Tue, 26 Nov 2024 14:09:27 -0500
Subject: [PATCH] store scores in regs

---
 lib/liquid/block_body.rb  | 8 ++++++++
 lib/liquid/tags/assign.rb | 4 ++++
 2 files changed, 12 insertions(+)

diff --git a/lib/liquid/block_body.rb b/lib/liquid/block_body.rb
index e4ada7d1..80c33450 100644
--- a/lib/liquid/block_body.rb
+++ b/lib/liquid/block_body.rb
@@ -218,6 +218,9 @@ module Liquid
     def render_to_output_buffer(context, output)
       freeze unless frozen?
 
+      context.registers[:render_scores] ||= []
+      context.registers[:render_scores] << @nodelist.length
+
       context.resource_limits.increment_render_score(@nodelist.length)
 
       idx = 0
@@ -233,6 +236,11 @@ module Liquid
         end
         idx += 1
 
+        context.registers[:write_scores] ||= []
+
+        last_captured = context.registers[:write_scores].any? ? context.registers[:write_scores].last : 0
+        context.registers[:write_scores] << (output.bytesize - last_captured)
+
         context.resource_limits.increment_write_score(output)
       end
 
diff --git a/lib/liquid/tags/assign.rb b/lib/liquid/tags/assign.rb
index 70283178..ad978e30 100644
--- a/lib/liquid/tags/assign.rb
+++ b/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[:assign_scores] ||= []
+      context.registers[:assign_scores] << assign_score_of(val)
+
       context.resource_limits.increment_assign_score(assign_score_of(val))
       output
     end