ZJIT: Implement ToArray and ArrayExtend

- ToArray is used in both `concattoarray` and `splatarray` instructions
  to convert a value to an array OR put the value in an array
- ArrayExtend is used in `concattoarray` to extend `ary1` with `ary2`

With this change, ZJIT can actually execute `concattoarray` and
`splatarray` instructions.

Author: st0012

test/ruby/test_zjit.rb

@@ -489,6 +489,42 @@ def test = [1,2,3]
     }
   end
 
+  def test_array_splat_custom_to_a
+    assert_compiles '[1, 2, 3]', %q{
+      class Foo
+        def to_a
+          [1, 2, 3]
+        end
+      end
+      def bar(*a) = a
+      def test = bar(*Foo.new)
+      test
+    }, insns: [:splatarray]
+  end
+
+  def test_array_splat_nil
+    assert_compiles '1', %q{
+      def foo = 1
+      def test = foo(*nil)
+      test
+    }, insns: [:splatarray]
+  end
+
+  def test_array_splat
+    assert_compiles '"foo"', %q{
+      def foo(b) = b
+      def test(a) = foo(*a)
+      test("foo")
+    }, insns: [:splatarray]
+  end
+
+  def test_array_concat_implicit
+    assert_compiles '[1, 2, 3]', %q{
+      def test(a) = [1, *a]
+      test([2, 3])
+    }, insns: [:concattoarray]
+  end
+
   def test_new_range_inclusive
     assert_compiles '1..5', %q{
       def test(a, b) = a..b

zjit/bindgen/src/main.rs

@@ -118,6 +118,7 @@ fn main() {
         .allowlist_function("rb_obj_is_kind_of")
         .allowlist_function("rb_obj_frozen_p")
         .allowlist_function("rb_class_inherited_p")
+        .allowlist_function("rb_Array")
 
         // From ruby/internal/encoding/encoding.h
         .allowlist_type("ruby_encoding_consts")
@@ -149,6 +150,7 @@ fn main() {
         .allowlist_function("rb_ary_dup")
         .allowlist_function("rb_ary_push")
         .allowlist_function("rb_ary_unshift_m")
+        .allowlist_function("rb_ary_concat")
 
         // From internal/array.h
         .allowlist_function("rb_ec_ary_new_from_values")

zjit/src/codegen.rs

@@ -367,7 +367,8 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::AnyToString { val, str, state } => gen_anytostring(asm, opnd!(val), opnd!(str), &function.frame_state(*state))?,
         Insn::Defined { op_type, obj, pushval, v } => gen_defined(jit, asm, *op_type, *obj, *pushval, opnd!(v))?,
         &Insn::IncrCounter(counter) => return Some(gen_incr_counter(asm, counter)),
-        Insn::ArrayExtend { .. }
+        Insn::ToArray { val, state } => gen_to_array(jit, asm, opnd!(val), &function.frame_state(*state))?,
+        Insn::ArrayExtend { left, right, state } => return gen_array_extend(jit, asm, opnd!(left), opnd!(right), &function.frame_state(*state)),
         | Insn::ArrayMax { .. }
         | Insn::ArrayPush { .. }
         | Insn::DefinedIvar { .. }
@@ -379,7 +380,6 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         | Insn::Send { .. }
         | Insn::StringIntern { .. }
         | Insn::Throw { .. }
-        | Insn::ToArray { .. }
         | Insn::ToNewArray { .. }
         | Insn::Const { .. }
         => {
@@ -768,12 +768,9 @@ fn gen_send_without_block(
     self_val: Opnd,
     args: Vec<Opnd>,
 ) -> Option<lir::Opnd> {
-    // Spill locals onto the stack.
-    // TODO: Don't spill locals eagerly; lazily reify frames
-    asm_comment!(asm, "spill locals");
-    for (idx, &insn_id) in state.locals().enumerate() {
-        asm.mov(Opnd::mem(64, SP, (-local_idx_to_ep_offset(jit.iseq, idx) - 1) * SIZEOF_VALUE_I32), jit.get_opnd(insn_id)?);
-    }
+
+    spill_locals(jit, asm, state);
+
     // Spill the receiver and the arguments onto the stack.
     // They need to be on the interpreter stack to let the interpreter access them.
     // TODO: Avoid spilling operands that have been spilled before.
@@ -819,13 +816,8 @@ fn gen_send_without_block_direct(
 
     // Spill the virtual stack and the locals of the caller onto the stack
     // TODO: Lazily materialize caller frames on side exits or when needed
-    asm_comment!(asm, "spill locals and stack");
-    for (idx, &insn_id) in state.locals().enumerate() {
-        asm.mov(Opnd::mem(64, SP, (-local_idx_to_ep_offset(jit.iseq, idx) - 1) * SIZEOF_VALUE_I32), jit.get_opnd(insn_id)?);
-    }
-    for (idx, &insn_id) in state.stack().enumerate() {
-        asm.mov(Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32), jit.get_opnd(insn_id)?);
-    }
+    spill_locals(jit, asm, state);
+    spill_stack(jit, asm, state);
 
     // Set up the new frame
     // TODO: Lazily materialize caller frames on side exits or when needed
@@ -1042,6 +1034,38 @@ fn gen_anytostring(asm: &mut Assembler, val: lir::Opnd, str: lir::Opnd, state: &
     Some(asm_ccall!(asm, rb_obj_as_string_result, str, val))
 }
 
+fn gen_to_array(jit: &mut JITState, asm: &mut Assembler, val: lir::Opnd, state: &FrameState) -> Option<lir::Opnd> {
+    // Save PC
+    gen_save_pc(asm, state);
+
+    // Save SP, dump locals and stack as rb_Array can call user-defined `to_a` method
+    gen_save_sp(asm, state.stack().len());
+    spill_locals(jit, asm, state);
+    spill_stack(jit, asm, state);
+
+    // ToArray is used in both concattoarray and splatarray instructions
+    // It needs to satisfy the following conditions:
+    // - Check if the value is already an array. If it is, return it
+    // - Check if the value can be converted into an array. If it is, return it
+    // - Return an array with just the value in it
+    //
+    // rb_Array satisfies the above conditions, so we can use it here
+    Some(asm_ccall!(asm, rb_Array, val))
+}
+
+fn gen_array_extend(jit: &mut JITState, asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd, state: &FrameState) -> Option<()> {
+    // Save PC
+    gen_save_pc(asm, state);
+
+    // Save SP, dump locals and stack as rb_ary_concat calls `rb_to_array_type`, which can call user-defined `to_ary` method
+    gen_save_sp(asm, state.stack().len());
+    spill_locals(jit, asm, state);
+    spill_stack(jit, asm, state);
+
+    asm_ccall!(asm, rb_ary_concat, left, right);
+    Some(())
+}
+
 /// Evaluate if a value is truthy
 /// Produces a CBool type (0 or 1)
 /// In Ruby, only nil and false are falsy
@@ -1263,6 +1287,28 @@ fn max_num_params(function: &Function) -> usize {
     }).max().unwrap_or(0)
 }
 
+/// Spill locals onto the stack.
+/// TODO: Don't spill locals eagerly; lazily reify frames
+fn spill_locals(jit: &mut JITState, asm: &mut Assembler, state: &FrameState) -> Option<()> {
+    asm_comment!(asm, "spill locals");
+    for (idx, &insn_id) in state.locals().enumerate() {
+        let local_mem = Opnd::mem(64, SP, (-local_idx_to_ep_offset(jit.iseq, idx) - 1) * SIZEOF_VALUE_I32);
+        asm.mov(local_mem, jit.get_opnd(insn_id)?);
+    }
+    Some(())
+}
+
+/// Spill stack onto the stack.
+/// TODO: Don't spill stack eagerly; lazily reify frames
+fn spill_stack(jit: &mut JITState, asm: &mut Assembler, state: &FrameState) -> Option<()> {
+    asm_comment!(asm, "spill stack");
+    for (idx, &insn_id) in state.stack().enumerate() {
+        let stack_mem = Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32);
+        asm.mov(stack_mem, jit.get_opnd(insn_id)?);
+    }
+    Some(())
+}
+
 #[cfg(target_arch = "x86_64")]
 macro_rules! c_callable {
     ($(#[$outer:meta])*