From be648e0940054a693105b1a3da6c2675ee724499 Mon Sep 17 00:00:00 2001
From: Kevin Newton <kddnewton@gmail.com>
Date: Wed, 14 Jul 2021 15:16:56 -0400
Subject: [PATCH] Implement splatarray

---
 bootstraptest/test_yjit.rb |  9 +++++++++
 vm_insnhelper.c            |  6 ++++++
 yjit_codegen.c             | 30 ++++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+)

diff --git a/bootstraptest/test_yjit.rb b/bootstraptest/test_yjit.rb
index 7fe637a86b..f55d63672f 100644
--- a/bootstraptest/test_yjit.rb
+++ b/bootstraptest/test_yjit.rb
@@ -1386,3 +1386,12 @@ assert_equal '[nil, nil, nil, 1]', %q{
     is_inf(1.0/0.0)
   ]
 }
+
+assert_equal '[1, 2, 3, 4, 5]', %q{
+  def splatarray
+    [*(1..5)]
+  end
+
+  splatarray
+  splatarray
+}
diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index 8babcca0a4..1b1958320d 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -4257,6 +4257,12 @@ vm_splat_array(VALUE flag, VALUE ary)
     }
 }
 
+VALUE
+rb_vm_splat_array(VALUE flag, VALUE ary)
+{
+    return vm_splat_array(flag, ary);
+}
+
 static VALUE
 vm_check_match(rb_execution_context_t *ec, VALUE target, VALUE pattern, rb_num_t flag)
 {
diff --git a/yjit_codegen.c b/yjit_codegen.c
index 577e104fe2..b9226cf738 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -691,6 +691,35 @@ gen_duparray(jitstate_t* jit, ctx_t* ctx)
     return YJIT_KEEP_COMPILING;
 }
 
+VALUE rb_vm_splat_array(VALUE flag, VALUE ary);
+
+// call to_a on the array on the stack
+static codegen_status_t
+gen_splatarray(jitstate_t* jit, ctx_t* ctx)
+{
+    VALUE flag = (VALUE) jit_get_arg(jit, 0);
+
+    // Save the PC and SP because the callee may allocate
+    // Note that this modifies REG_SP, which is why we do it first
+    jit_save_pc(jit, REG0);
+    jit_save_sp(jit, ctx);
+
+    // Get the operands from the stack
+    x86opnd_t ary_opnd = ctx_stack_pop(ctx, 1);
+
+    // Call rb_vm_splat_array(flag, ary)
+    yjit_save_regs(cb);
+    jit_mov_gc_ptr(jit, cb, C_ARG_REGS[0], flag);
+    mov(cb, C_ARG_REGS[1], ary_opnd);
+    call_ptr(cb, REG1, (void *) rb_vm_splat_array);
+    yjit_load_regs(cb);
+
+    x86opnd_t stack_ret = ctx_stack_push(ctx, TYPE_ARRAY);
+    mov(cb, stack_ret, RAX);
+
+    return YJIT_KEEP_COMPILING;
+}
+
 // new hash initialized from top N values
 static codegen_status_t
 gen_newhash(jitstate_t* jit, ctx_t* ctx)
@@ -3470,6 +3499,7 @@ yjit_init_codegen(void)
     yjit_reg_op(BIN(adjuststack), gen_adjuststack);
     yjit_reg_op(BIN(newarray), gen_newarray);
     yjit_reg_op(BIN(duparray), gen_duparray);
+    yjit_reg_op(BIN(splatarray), gen_splatarray);
     yjit_reg_op(BIN(newhash), gen_newhash);
     yjit_reg_op(BIN(concatstrings), gen_concatstrings);
     yjit_reg_op(BIN(putnil), gen_putnil);