Compile assign tag to VM code

ext/liquid_c/context.c

@@ -194,6 +194,19 @@ VALUE context_filtering_p(VALUE self)
     return liquid_vm_filtering(self) ? Qtrue : Qfalse;
 }
 
+void context_assign(context_t *context, VALUE name, VALUE value)
+{
+    VALUE scopes = context->scopes;
+    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 liquid_define_context()
 {
     id_has_key = rb_intern("key?");

ext/liquid_c/context.h

@@ -23,6 +23,7 @@ void context_internal_init(VALUE context_obj, context_t *context);
 void context_mark(context_t *context);
 VALUE context_find_variable(context_t *context, VALUE key, VALUE raise_on_not_found);
 void context_maybe_raise_undefined_variable(VALUE self, VALUE key);
+void context_assign(context_t *context, VALUE name, VALUE value);
 
 extern ID id_aset, id_set_context;
 

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;
 

ext/liquid_c/resource_limits.h

@@ -18,6 +18,7 @@ extern const rb_data_type_t resource_limits_data_type;
 void liquid_define_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

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)
@@ -57,6 +58,31 @@ 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);
+    block_body_ensure_intermediate(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),
+        .code = body->as.intermediate.code,
+        .code_obj = block_body_obj,
+        .parse_context = parse_context_obj,
+    };
+    internal_variable_compile_evaluate(&parse_args);
+    vm_assembler_add_assign(body->as.intermediate.code, var_name);
+    return Qnil;
+}
+
 void liquid_define_tag()
 {
     id_parse = rb_intern("parse");
@@ -69,4 +95,12 @@ void liquid_define_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);
 }

ext/liquid_c/vm.c

@@ -227,6 +227,43 @@ 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(key);
+    *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 sum;
+        }
+        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)
 {
@@ -386,6 +423,14 @@ static VALUE vm_render_until_error(VALUE uncast_args)
                 resource_limits_increment_write_score(vm->context.resource_limits, output);
                 break;
             }
+            case OP_ASSIGN:
+            {
+                VALUE var_name = (VALUE)*const_ptr++;
+                VALUE value = vm_stack_pop(vm);
+                context_assign(&vm->context, var_name, value);
+                resource_limits_increment_assign_score(vm->context.resource_limits, assign_score_of(value));
+                break;
+            }
 
             default:
                 rb_bug("invalid opcode: %u", ip[-1]);
@@ -444,6 +489,7 @@ void liquid_vm_next_instruction(const uint8_t **ip_ptr, const VALUE **const_ptr_
         case OP_FIND_STATIC_VAR:
         case OP_LOOKUP_CONST_KEY:
         case OP_LOOKUP_COMMAND:
+        case OP_ASSIGN:
             (*const_ptr_ptr)++;
             break;
 

ext/liquid_c/vm_assembler.c

@@ -196,6 +196,10 @@ VALUE vm_assembler_disassemble(const uint8_t *start_ip, const uint8_t *end_ip, c
                 rb_str_catf(output, "lookup_command(%+"PRIsVALUE")\n", const_ptr[0]);
                 break;
 
+            case OP_ASSIGN:
+                rb_str_catf(output, "assign(%+"PRIsVALUE")\n", const_ptr[0]);
+                break;
+
             case OP_FILTER:
                 rb_str_catf(output, "filter(name: %+"PRIsVALUE", num_args: %u)\n", const_ptr[0], ip[1]);
                 break;

ext/liquid_c/vm_assembler.h

@@ -12,6 +12,7 @@ enum opcode {
     OP_WRITE_NODE = 2,
     OP_POP_WRITE,
     OP_WRITE_RAW_SKIP,
+    OP_ASSIGN,
     OP_PUSH_CONST,
     OP_PUSH_NIL,
     OP_PUSH_TRUE,
@@ -222,4 +223,11 @@ static inline void vm_assembler_add_render_variable_rescue(vm_assembler_t *code,
     uint24_to_bytes((unsigned int)node_line_number, &instructions[1]);
 }
 
+static inline void vm_assembler_add_assign(vm_assembler_t *code, VALUE variable_name)
+{
+    code->stack_size--;
+    vm_assembler_write_ruby_constant(code, variable_name);
+    vm_assembler_write_opcode(code, OP_ASSIGN);
+}
+
 #endif