ZJIT: Generate code for ArrayExtend

4 files changed, 18 insertions, 2 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 736c6306d1..e2a0120eb0 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -139,6 +139,15 @@ class TestZJIT < Test::Unit::TestCase
     }, insns: [:splatarray]
   end
 
+  def test_concattoarray
+    assert_compiles '[1, 2, 3]', %q{
+      def test(*a)
+        [1, 2, *a]
+      end
+      test 3
+    }, insns: [:concattoarray]
+  end
+
   def test_definedivar
     assert_compiles '[nil, "instance-variable", nil]', %q{
       def test
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
index e4c073d48a..0eb800e1f9 100644
--- a/zjit/bindgen/src/main.rs
+++ b/zjit/bindgen/src/main.rs
@@ -145,6 +145,7 @@ fn main() {
         .allowlist_function("rb_ary_store")
         .allowlist_function("rb_ary_resurrect")
         .allowlist_function("rb_ary_cat")
+        .allowlist_function("rb_ary_concat")
         .allowlist_function("rb_ary_clear")
         .allowlist_function("rb_ary_dup")
         .allowlist_function("rb_ary_push")
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index cdc62043fd..189741b07c 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -403,8 +403,8 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &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::DefinedIvar { self_val, id, pushval, .. } => { gen_defined_ivar(asm, opnd!(self_val), id, pushval) },
-        &Insn::ArrayExtend { state, .. }
-        | &Insn::ArrayMax { state, .. }
+        &Insn::ArrayExtend { left, right, state } => { no_output!(gen_array_extend(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state))) },
+        &Insn::ArrayMax { state, .. }
         | &Insn::FixnumDiv { state, .. }
         | &Insn::FixnumMod { state, .. }
         | &Insn::Send { state, .. }
@@ -708,6 +708,11 @@ fn gen_defined_ivar(asm: &mut Assembler, self_val: Opnd, id: ID, pushval: VALUE)
     asm_ccall!(asm, rb_zjit_defined_ivar, self_val, id.0.into(), Opnd::Value(pushval))
 }
 
+fn gen_array_extend(jit: &mut JITState, asm: &mut Assembler, left: Opnd, right: Opnd, state: &FrameState) {
+    gen_prepare_non_leaf_call(jit, asm, state);
+    asm_ccall!(asm, rb_ary_concat, left, right);
+}
+
 /// 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));
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index 33f7e0b3e6..4975308f68 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -783,6 +783,7 @@ unsafe extern "C" {
     pub fn rb_ary_cat(ary: VALUE, train: *const VALUE, len: ::std::os::raw::c_long) -> VALUE;
     pub fn rb_ary_push(ary: VALUE, elem: VALUE) -> VALUE;
     pub fn rb_ary_clear(ary: VALUE) -> VALUE;
+    pub fn rb_ary_concat(lhs: VALUE, rhs: VALUE) -> VALUE;
     pub fn rb_hash_new() -> VALUE;
     pub fn rb_hash_aref(hash: VALUE, key: VALUE) -> VALUE;
     pub fn rb_hash_aset(hash: VALUE, key: VALUE, val: VALUE) -> VALUE;