Compile assign tag to VM code

Shopify · Oct 20, 2020 · 0ed4a1d · 0ed4a1d
1 parent eb3e1e2
commit 0ed4a1d
Show file tree

Hide file tree

Showing 8 changed files with 103 additions and 2 deletions.
diff --git a/ext/liquid_c/context.c b/ext/liquid_c/context.c
@@ -131,6 +131,20 @@ VALUE context_filtering_p(VALUE self)
     return liquid_vm_filtering(self) ? Qtrue : Qfalse;
 }
 
+void context_assign(VALUE context, VALUE name, VALUE value)
+{
+    VALUE scopes = rb_ivar_get(context, id_ivar_scopes);
+    Check_Type(scopes, T_ARRAY);
+    long scopes_size = RARRAY_LEN(scopes);
+    if (RB_UNLIKELY(scopes_size == 0))
+        rb_raise(rb_eRuntimeError, "Liquid::Context#scopes is empty, missing the required outer scope");
+
+    VALUE last_scope = RARRAY_AREF(scopes, scopes_size - 1);
+    Check_Type(last_scope, T_HASH);
+
+    rb_hash_aset(last_scope, name, value);
+}
+
 void init_liquid_context()
 {
     id_has_key = rb_intern("key?");

diff --git a/ext/liquid_c/context.h b/ext/liquid_c/context.h
@@ -4,6 +4,7 @@
 void init_liquid_context();
 VALUE context_find_variable(VALUE self, VALUE key, VALUE raise_on_not_found);
 void context_maybe_raise_undefined_variable(VALUE self, VALUE key);
+void context_assign(VALUE context, VALUE name, VALUE value);
 
 extern ID id_aset, id_set_context;
 

diff --git a/ext/liquid_c/resource_limits.c b/ext/liquid_c/resource_limits.c
@@ -157,7 +157,7 @@ static VALUE resource_limits_increment_render_score_method(VALUE self, VALUE amo
     return Qnil;
 }
 
-static void resource_limits_increment_assign_score(resource_limits_t *resource_limits, long amount)
+void resource_limits_increment_assign_score(resource_limits_t *resource_limits, long amount)
 {
     resource_limits->assign_score = resource_limits->assign_score + amount;
 

diff --git a/ext/liquid_c/resource_limits.h b/ext/liquid_c/resource_limits.h
@@ -18,6 +18,7 @@ extern const rb_data_type_t resource_limits_data_type;
 void init_liquid_resource_limits();
 void resource_limits_raise_limits_reached(resource_limits_t *resource_limit);
 void resource_limits_increment_render_score(resource_limits_t *resource_limits, long amount);
+void resource_limits_increment_assign_score(resource_limits_t *resource_limits, long amount);
 void resource_limits_increment_write_score(resource_limits_t *resource_limits, VALUE output);
 
 #endif
diff --git a/ext/liquid_c/tag.c b/ext/liquid_c/tag.c
@@ -6,7 +6,8 @@
 
 static ID id_parse;
 
-static VALUE cLiquidTag;
+static VALUE cLiquidTag, cLiquidAssign;
+static VALUE liquid_assign_syntax;
 
 static VALUE tag_class_compile(VALUE klass, VALUE tag_name, VALUE markup,
         VALUE tokenizer_obj, VALUE parse_context_obj, VALUE block_body_obj)
@@ -52,6 +53,29 @@ static VALUE echo_class_compile(VALUE klass, VALUE tag_name, VALUE markup,
     return Qnil;
 }
 
+static VALUE assign_class_compile(VALUE klass, VALUE tag_name, VALUE markup,
+        VALUE tokenizer_obj, VALUE parse_context_obj, VALUE block_body_obj)
+{
+    block_body_t *body;
+    BlockBody_Get_Struct(block_body_obj, body);
+
+    if (rb_reg_match(liquid_assign_syntax, markup) == Qnil)
+        rb_funcall(cLiquidAssign, rb_intern("raise_syntax_error"), 1, parse_context_obj);
+    VALUE last_match = rb_backref_get();
+    VALUE var_name = rb_reg_nth_match(1, last_match);
+    VALUE var_markup = rb_reg_nth_match(2, last_match);
+
+    variable_parse_args_t parse_args = {
+        .markup = RSTRING_PTR(var_markup),
+        .markup_end = RSTRING_END(var_markup),
+        .body = body,
+        .parse_context = parse_context_obj,
+    };
+    internal_variable_compile_evaluate(&parse_args);
+    vm_assembler_add_assign(&body->code, var_name);
+    return Qnil;
+}
+
 void init_liquid_tag()
 {
     id_parse = rb_intern("parse");
@@ -64,4 +88,12 @@ void init_liquid_tag()
 
     VALUE cLiquidEcho = rb_const_get(mLiquid, rb_intern("Echo"));
     rb_define_singleton_method(cLiquidEcho, "compile", echo_class_compile, 5);
+
+    cLiquidAssign = rb_const_get(mLiquid, rb_intern("Assign"));
+    rb_global_variable(&cLiquidAssign);
+    rb_define_singleton_method(cLiquidAssign, "compile", assign_class_compile, 5);
+
+    liquid_assign_syntax = rb_const_get(cLiquidAssign, rb_intern("Syntax"));
+    rb_global_variable(&liquid_assign_syntax);
+    Check_Type(liquid_assign_syntax, T_REGEXP);
 }
diff --git a/ext/liquid_c/vm.c b/ext/liquid_c/vm.c
@@ -261,6 +261,42 @@ static void hash_bulk_insert(long argc, const VALUE *argv, VALUE hash)
 }
 #endif
 
+static long assign_score_of(VALUE value);
+
+static int assign_score_of_each_hash_value(VALUE key, VALUE value, VALUE func_arg)
+{
+    long *sum = (long *)func_arg;
+    *sum += assign_score_of(value);
+    return ST_CONTINUE;
+}
+
+static long assign_score_of(VALUE value)
+{
+    if (RB_SPECIAL_CONST_P(value))
+        return 1;
+
+    switch (RB_BUILTIN_TYPE(value)) {
+        case T_STRING:
+            return RSTRING_LEN(value);
+        case T_HASH:
+        {
+            long sum = 1;
+            rb_hash_foreach(value, assign_score_of_each_hash_value, (VALUE)&sum);
+            return RHASH_SIZE(value);
+        }
+        case T_ARRAY:
+        {
+            long sum = 1;
+            for (long i = 0; i < RARRAY_LEN(value); i++) {
+                sum += assign_score_of(RARRAY_AREF(value, i));
+            }
+            return sum;
+        }
+        default:
+            return 1;
+    }
+}
+
 // Actually returns a bool resume_rendering value
 static VALUE vm_render_until_error(VALUE uncast_args)
 {
@@ -389,6 +425,14 @@ static VALUE vm_render_until_error(VALUE uncast_args)
                 resource_limits_increment_write_score(vm->resource_limits, output);
                 break;
             }
+            case OP_ASSIGN:
+            {
+                VALUE var_name = (VALUE)*const_ptr++;
+                VALUE value = vm_stack_pop(vm);
+                context_assign(args->context, var_name, value);
+                resource_limits_increment_assign_score(vm->resource_limits, assign_score_of(value));
+                break;
+            }
 
             default:
                 rb_bug("invalid opcode: %u", ip[-1]);
@@ -447,6 +491,7 @@ void liquid_vm_next_instruction(const uint8_t **ip_ptr, const size_t **const_ptr
         case OP_FIND_STATIC_VAR:
         case OP_LOOKUP_CONST_KEY:
         case OP_LOOKUP_COMMAND:
+        case OP_ASSIGN:
             (*const_ptr_ptr)++;
             break;
 

diff --git a/ext/liquid_c/vm_assembler.c b/ext/liquid_c/vm_assembler.c
@@ -59,6 +59,7 @@ void vm_assembler_gc_mark(vm_assembler_t *code)
             case OP_FIND_STATIC_VAR:
             case OP_LOOKUP_CONST_KEY:
             case OP_LOOKUP_COMMAND:
+            case OP_ASSIGN:
                 rb_gc_mark(*const_ptr++);
                 break;
 

diff --git a/ext/liquid_c/vm_assembler.h b/ext/liquid_c/vm_assembler.h
@@ -10,6 +10,7 @@ enum opcode {
     OP_WRITE_RAW = 1,
     OP_WRITE_NODE = 2,
     OP_POP_WRITE,
+    OP_ASSIGN,
     OP_PUSH_CONST,
     OP_PUSH_NIL,
     OP_PUSH_TRUE,
@@ -212,4 +213,10 @@ static inline void vm_assembler_add_render_variable_rescue(vm_assembler_t *code,
     c_buffer_write(&code->instructions, &instructions, sizeof(instructions));
 }
 
+static inline void vm_assembler_add_assign(vm_assembler_t *code, VALUE variable_name)
+{
+    vm_assembler_write_ruby_constant(code, variable_name);
+    vm_assembler_write_opcode(code, OP_ASSIGN);
+}
+
 #endif