ZJIT: Generate code for ToArray, ToNewArray

2 files changed, 33 insertions, 2 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 075a7765c8..52f2cddeb7 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -118,6 +118,27 @@ class TestZJIT < Test::Unit::TestCase
     }, insns: [:pushtoarray]
   end
 
+  def test_splatarray_new_array
+    assert_compiles '[1, 2, 3]', %q{
+      def test a
+        [*a, 3]
+      end
+      test [1, 2]
+    }, insns: [:splatarray]
+  end
+
+  def test_splatarray_existing_array
+    assert_compiles '[1, 2, 3]', %q{
+      def foo v
+        [1, 2, v]
+      end
+      def test a
+        foo(*a)
+      end
+      test [3]
+    }, insns: [:splatarray]
+  end
+
   def test_setglobal_with_trace_var_exception
     assert_compiles '"rescued"', %q{
       def test
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index f257e5c138..d87c94cf8c 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -400,6 +400,8 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::CheckInterrupts { state } => no_output!(gen_check_interrupts(jit, asm, &function.frame_state(state))),
         &Insn::HashDup { val, state } => { gen_hash_dup(asm, opnd!(val), &function.frame_state(state)) },
         &Insn::ArrayPush { array, val, state } => { no_output!(gen_array_push(asm, opnd!(array), opnd!(val), &function.frame_state(state))) },
+        &Insn::ToNewArray { val, state } => { gen_to_new_array(jit, asm, opnd!(val), &function.frame_state(state)) },
+        &Insn::ToArray { val, state } => { gen_to_array(jit, asm, opnd!(val), &function.frame_state(state)) },
         &Insn::ArrayExtend { state, .. }
         | &Insn::ArrayMax { state, .. }
         | &Insn::DefinedIvar { state, .. }
@@ -407,8 +409,6 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         | &Insn::FixnumMod { state, .. }
         | &Insn::Send { state, .. }
         | &Insn::Throw { state, .. }
-        | &Insn::ToArray { state, .. }
-        | &Insn::ToNewArray { state, .. }
         => return Err(state),
     };
 
@@ -694,6 +694,16 @@ fn gen_array_push(asm: &mut Assembler, array: Opnd, val: Opnd, state: &FrameStat
     asm_ccall!(asm, rb_ary_push, array, val);
 }
 
+fn gen_to_new_array(jit: &mut JITState, asm: &mut Assembler, val: Opnd, state: &FrameState) -> lir::Opnd {
+    gen_prepare_non_leaf_call(jit, asm, state);
+    asm_ccall!(asm, rb_vm_splat_array, Opnd::Value(Qtrue), val)
+}
+
+fn gen_to_array(jit: &mut JITState, asm: &mut Assembler, val: Opnd, state: &FrameState) -> lir::Opnd {
+    gen_prepare_non_leaf_call(jit, asm, state);
+    asm_ccall!(asm, rb_vm_splat_array, Opnd::Value(Qfalse), val)
+}
+
 /// Compile an interpreter entry block to be inserted into an ISEQ
 fn gen_entry_prologue(asm: &mut Assembler, iseq: IseqPtr) {
     asm_comment!(asm, "ZJIT entry point: {}", iseq_get_location(iseq, 0));