Merge pull request #1323 from Shopify/assign-score-hash

Avoid allocating arrays of key value pairs in assign_score_of

lib/liquid/tags/assign.rb

@@ -45,11 +45,18 @@ module Liquid
     def assign_score_of(val)
       if val.instance_of?(String)
         val.bytesize
-      elsif val.instance_of?(Array) || val.instance_of?(Hash)
+      elsif val.instance_of?(Array)
         sum = 1
         # Uses #each to avoid extra allocations.
         val.each { |child| sum += assign_score_of(child) }
         sum
+      elsif val.instance_of?(Hash)
+        sum = 1
+        val.each do |key, entry_value|
+          sum += assign_score_of(key)
+          sum += assign_score_of(entry_value)
+        end
+        sum
       else
         1
       end

test/integration/assign_test.rb

@@ -47,4 +47,66 @@ class AssignTest < Minitest::Test
       assert Template.parse("{% assign foo = ('X' | downcase) %}")
     end
   end
-end # AssignTest
+
+  def test_assign_score_exceeding_resource_limit
+    t = Template.parse("{% assign foo = 42 %}{% assign bar = 23 %}")
+    t.resource_limits.assign_score_limit = 1
+    assert_equal("Liquid error: Memory limits exceeded", t.render)
+    assert(t.resource_limits.reached?)
+
+    t.resource_limits.assign_score_limit = 2
+    assert_equal("", t.render!)
+    refute_nil(t.resource_limits.assign_score)
+  end
+
+  def test_assign_score_exceeding_limit_from_composite_object
+    t = Template.parse("{% assign foo = 'aaaa' | reverse %}")
+
+    t.resource_limits.assign_score_limit = 3
+    assert_equal("Liquid error: Memory limits exceeded", t.render)
+    assert(t.resource_limits.reached?)
+
+    t.resource_limits.assign_score_limit = 5
+    assert_equal("", t.render!)
+  end
+
+  def test_assign_score_of_int
+    assert_equal(1, assign_score_of(123))
+  end
+
+  def test_assign_score_of_string_counts_bytes
+    assert_equal(3, assign_score_of('123'))
+    assert_equal(5, assign_score_of('12345'))
+    assert_equal(9, assign_score_of('すごい'))
+  end
+
+  def test_assign_score_of_array
+    assert_equal(1, assign_score_of([]))
+    assert_equal(2, assign_score_of([123]))
+    assert_equal(6, assign_score_of([123, 'abcd']))
+  end
+
+  def test_assign_score_of_hash
+    assert_equal(1, assign_score_of({}))
+    assert_equal(5, assign_score_of('int' => 123))
+    assert_equal(12, assign_score_of('int' => 123, 'str' => 'abcd'))
+  end
+
+  private
+
+  class ObjectWrapperDrop < Liquid::Drop
+    def initialize(obj)
+      @obj = obj
+    end
+
+    def value
+      @obj
+    end
+  end
+
+  def assign_score_of(obj)
+    context = Liquid::Context.new('drop' => ObjectWrapperDrop.new(obj))
+    Liquid::Template.parse('{% assign obj = drop.value %}').render!(context)
+    context.resource_limits.assign_score
+  end
+end

test/integration/capture_test.rb

@@ -49,4 +49,10 @@ class CaptureTest < Minitest::Test
     rendered        = template.render!
     assert_equal("3-3", rendered.gsub(/\s/, ''))
   end
-end # CaptureTest
+
+  def test_increment_assign_score_by_bytes_not_characters
+    t = Template.parse("{% capture foo %}すごい{% endcapture %}")
+    t.render!
+    assert_equal(9, t.resource_limits.assign_score)
+  end
+end

test/integration/template_test.rb

@@ -135,38 +135,6 @@ class TemplateTest < Minitest::Test
     refute_nil(t.resource_limits.render_score)
   end
 
-  def test_resource_limits_assign_score
-    t = Template.parse("{% assign foo = 42 %}{% assign bar = 23 %}")
-    t.resource_limits.assign_score_limit = 1
-    assert_equal("Liquid error: Memory limits exceeded", t.render)
-    assert(t.resource_limits.reached?)
-
-    t.resource_limits.assign_score_limit = 2
-    assert_equal("", t.render!)
-    refute_nil(t.resource_limits.assign_score)
-  end
-
-  def test_resource_limits_assign_score_counts_bytes_not_characters
-    t = Template.parse("{% assign foo = 'すごい' %}")
-    t.render
-    assert_equal(9, t.resource_limits.assign_score)
-
-    t = Template.parse("{% capture foo %}すごい{% endcapture %}")
-    t.render
-    assert_equal(9, t.resource_limits.assign_score)
-  end
-
-  def test_resource_limits_assign_score_nested
-    t = Template.parse("{% assign foo = 'aaaa' | reverse %}")
-
-    t.resource_limits.assign_score_limit = 3
-    assert_equal("Liquid error: Memory limits exceeded", t.render)
-    assert(t.resource_limits.reached?)
-
-    t.resource_limits.assign_score_limit = 5
-    assert_equal("", t.render!)
-  end
-
   def test_resource_limits_aborts_rendering_after_first_error
     t = Template.parse("{% for a in (1..100) %} foo1 {% endfor %} bar {% for a in (1..100) %} foo2 {% endfor %}")
     t.resource_limits.render_score_limit = 50