Add concatstrings to yjit codegen (#58)

* Add ETYPE_TRUE and ETYPE_FALSE

* Implement checktype

* Implement concatstrings

* Update deps

4 files changed, 139 insertions, 1 deletions

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 51b4a9b85f..5bf415889e 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -977,3 +977,47 @@ assert_equal '{:foo=>:bar}', %q{
   build_hash(:bar)
   build_hash(:bar)
 }
+
+# test string interpolation with known types
+assert_equal 'foobar', %q{
+  def make_str
+    foo = -"foo"
+    bar = -"bar"
+    "#{foo}#{bar}"
+  end
+
+  make_str
+  make_str
+}
+
+# test string interpolation with unknown types
+assert_equal 'foobar', %q{
+  def make_str(foo, bar)
+    "#{foo}#{bar}"
+  end
+
+  make_str("foo", "bar")
+  make_str("foo", "bar")
+}
+
+# test string interpolation with known non-strings
+assert_equal 'foo123', %q{
+  def make_str
+    foo = -"foo"
+    bar = 123
+    "#{foo}#{bar}"
+  end
+
+  make_str
+  make_str
+}
+
+# test string interpolation with unknown non-strings
+assert_equal 'foo123', %q{
+  def make_str(foo, bar)
+    "#{foo}#{bar}"
+  end
+
+  make_str("foo", 123)
+  make_str("foo", 123)
+}
diff --git a/common.mk b/common.mk
index 01773deb21..10e12a3740 100644
--- a/common.mk
+++ b/common.mk
@@ -16712,6 +16712,7 @@ yjit_codegen.$(OBJEXT): $(top_srcdir)/internal/imemo.h
 yjit_codegen.$(OBJEXT): $(top_srcdir)/internal/object.h
 yjit_codegen.$(OBJEXT): $(top_srcdir)/internal/serial.h
 yjit_codegen.$(OBJEXT): $(top_srcdir)/internal/static_assert.h
+yjit_codegen.$(OBJEXT): $(top_srcdir)/internal/string.h
 yjit_codegen.$(OBJEXT): $(top_srcdir)/internal/vm.h
 yjit_codegen.$(OBJEXT): $(top_srcdir)/internal/warnings.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}assert.h
@@ -16730,6 +16731,7 @@ yjit_codegen.$(OBJEXT): {$(VPATH)}config.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}darray.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}debug_counter.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}defines.h
+yjit_codegen.$(OBJEXT): {$(VPATH)}encoding.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}id.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}id_table.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}insns.def
@@ -16880,6 +16882,8 @@ yjit_codegen.$(OBJEXT): {$(VPATH)}iseq.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}method.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}missing.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}node.h
+yjit_codegen.$(OBJEXT): {$(VPATH)}onigmo.h
+yjit_codegen.$(OBJEXT): {$(VPATH)}oniguruma.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}ruby_assert.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}ruby_atomic.h
 yjit_codegen.$(OBJEXT): {$(VPATH)}st.h
diff --git a/yjit_codegen.c b/yjit_codegen.c
index b1b427ab75..d2b4e3ddbb 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -8,6 +8,7 @@
 #include "internal/compile.h"
 #include "internal/class.h"
 #include "internal/object.h"
+#include "internal/string.h"
 #include "insns_info.inc"
 #include "yjit.h"
 #include "yjit_iface.h"
@@ -1317,6 +1318,89 @@ gen_defined(jitstate_t* jit, ctx_t* ctx)
     return YJIT_KEEP_COMPILING;
 }
 
+static codegen_status_t
+gen_checktype(jitstate_t* jit, ctx_t* ctx)
+{
+    // TODO: could we specialize on the type we detect
+    uint8_t* side_exit = yjit_side_exit(jit, ctx);
+
+    enum ruby_value_type type_val = (enum ruby_value_type)jit_get_arg(jit, 0);
+    // Only three types are emitted by compile.c
+    if (type_val == T_STRING || type_val == T_ARRAY || type_val == T_HASH) {
+        val_type_t val_type = ctx_get_opnd_type(ctx, OPND_STACK(0));
+        x86opnd_t val = ctx_stack_pop(ctx, 1);
+
+        x86opnd_t stack_ret;
+
+        // Check if we know from type information
+        if ((type_val == T_STRING && val_type.type == ETYPE_STRING) ||
+                (type_val == T_ARRAY && val_type.type == ETYPE_ARRAY) ||
+                (type_val == T_HASH && val_type.type == ETYPE_HASH)) {
+            // guaranteed type match
+            stack_ret = ctx_stack_push(ctx, TYPE_TRUE);
+            mov(cb, stack_ret, imm_opnd(Qtrue));
+            return YJIT_KEEP_COMPILING;
+        } else if (val_type.is_imm || val_type.type != ETYPE_UNKNOWN) {
+            // guaranteed not to match T_STRING/T_ARRAY/T_HASH
+            stack_ret = ctx_stack_push(ctx, TYPE_FALSE);
+            mov(cb, stack_ret, imm_opnd(Qfalse));
+            return YJIT_KEEP_COMPILING;
+        }
+
+        mov(cb, REG0, val);
+
+        if (!val_type.is_heap) {
+            // if (SPECIAL_CONST_P(val)) {
+            // Bail if receiver is not a heap object
+            test(cb, REG0, imm_opnd(RUBY_IMMEDIATE_MASK));
+            jnz_ptr(cb, side_exit);
+            cmp(cb, REG0, imm_opnd(Qfalse));
+            je_ptr(cb, side_exit);
+            cmp(cb, REG0, imm_opnd(Qnil));
+            je_ptr(cb, side_exit);
+        }
+
+        // Check type on object
+        mov(cb, REG0, mem_opnd(64, REG0, offsetof(struct RBasic, flags)));
+        and(cb, REG0, imm_opnd(RUBY_T_MASK));
+        cmp(cb, REG0, imm_opnd(type_val));
+        mov(cb, REG1, imm_opnd(Qfalse));
+        cmovne(cb, REG0, REG1);
+
+        stack_ret = ctx_stack_push(ctx, TYPE_IMM);
+        mov(cb, stack_ret, REG0);
+
+        return YJIT_KEEP_COMPILING;
+    } else {
+        return YJIT_CANT_COMPILE;
+    }
+}
+
+static codegen_status_t
+gen_concatstrings(jitstate_t* jit, ctx_t* ctx)
+{
+    rb_num_t n = (rb_num_t)jit_get_arg(jit, 0);
+
+    // Save the PC and SP because we are allocating
+    jit_save_pc(jit, REG0);
+    jit_save_sp(jit, ctx);
+
+    x86opnd_t values_ptr = ctx_sp_opnd(ctx, -(sizeof(VALUE) * (uint32_t)n));
+
+    // call rb_str_concat_literals(long n, const VALUE *strings);
+    yjit_save_regs(cb);
+    mov(cb, C_ARG_REGS[0], imm_opnd(n));
+    lea(cb, C_ARG_REGS[1], values_ptr);
+    call_ptr(cb, REG0, (void *)rb_str_concat_literals);
+    yjit_load_regs(cb);
+
+    ctx_stack_pop(ctx, n);
+    x86opnd_t stack_ret = ctx_stack_push(ctx, TYPE_STRING);
+    mov(cb, stack_ret, RAX);
+
+    return YJIT_KEEP_COMPILING;
+}
+
 static void
 guard_two_fixnums(ctx_t* ctx, uint8_t* side_exit)
 {
@@ -2902,6 +2986,7 @@ yjit_init_codegen(void)
     yjit_reg_op(BIN(adjuststack), gen_adjuststack);
     yjit_reg_op(BIN(newarray), gen_newarray);
     yjit_reg_op(BIN(newhash), gen_newhash);
+    yjit_reg_op(BIN(concatstrings), gen_concatstrings);
     yjit_reg_op(BIN(putnil), gen_putnil);
     yjit_reg_op(BIN(putobject), gen_putobject);
     yjit_reg_op(BIN(putobject_INT2FIX_0_), gen_putobject_int2fix);
@@ -2913,6 +2998,7 @@ yjit_init_codegen(void)
     yjit_reg_op(BIN(getinstancevariable), gen_getinstancevariable);
     yjit_reg_op(BIN(setinstancevariable), gen_setinstancevariable);
     yjit_reg_op(BIN(defined), gen_defined);
+    yjit_reg_op(BIN(checktype), gen_checktype);
     yjit_reg_op(BIN(opt_lt), gen_opt_lt);
     yjit_reg_op(BIN(opt_le), gen_opt_le);
     yjit_reg_op(BIN(opt_ge), gen_opt_ge);
diff --git a/yjit_core.h b/yjit_core.h
index d9cce3fe56..a2436e6eb3 100644
--- a/yjit_core.h
+++ b/yjit_core.h
@@ -32,6 +32,8 @@ enum yjit_type_enum
 {
     ETYPE_UNKNOWN = 0,
     ETYPE_NIL,
+    ETYPE_TRUE,
+    ETYPE_FALSE,
     ETYPE_FIXNUM,
     ETYPE_ARRAY,
     ETYPE_HASH,
@@ -49,7 +51,7 @@ typedef struct yjit_type_struct
     uint8_t is_imm : 1;
 
     // Specific value type, if known
-    uint8_t type : 3;
+    uint8_t type : 4;
 
 } val_type_t;
 STATIC_ASSERT(val_type_size, sizeof(val_type_t) == 1);
@@ -64,6 +66,8 @@ STATIC_ASSERT(val_type_size, sizeof(val_type_t) == 1);
 #define TYPE_IMM ( (val_type_t){ .is_imm = 1 } )
 
 #define TYPE_NIL ( (val_type_t){ .is_imm = 1, .type = ETYPE_NIL } )
+#define TYPE_TRUE ( (val_type_t){ .is_imm = 1, .type = ETYPE_TRUE } )
+#define TYPE_FALSE ( (val_type_t){ .is_imm = 1, .type = ETYPE_FALSE } )
 #define TYPE_FIXNUM ( (val_type_t){ .is_imm = 1, .type = ETYPE_FIXNUM } )
 #define TYPE_ARRAY ( (val_type_t){ .is_heap = 1, .type = ETYPE_ARRAY } )
 #define TYPE_HASH ( (val_type_t){ .is_heap = 1, .type = ETYPE_HASH } )