RJIT: Support splat args

Author: k0kubun

lib/ruby_vm/rjit/insn_compiler.rb

@@ -4217,7 +4217,7 @@ def jit_call_iseq(jit, ctx, asm, cme, calling, iseq, frame_type: nil, prev_ep: n
         return CantCompile
       end
 
-      # block_arg = flags & C::VM_CALL_ARGS_BLOCKARG != 0
+      block_arg = flags & C::VM_CALL_ARGS_BLOCKARG != 0
       # jit_caller_setup_arg_block already handled send_blockarg_not_nil_or_proxy
 
       # If we have unfilled optional arguments and keyword arguments then we
@@ -4318,9 +4318,17 @@ def jit_call_iseq(jit, ctx, asm, cme, calling, iseq, frame_type: nil, prev_ep: n
         end
       end
       if flags & C::VM_CALL_ARGS_SPLAT != 0 && !iseq_has_rest
-        splat_array_length = false
-        asm.incr_counter(:send_iseq_splat)
-        return CantCompile
+        array = jit.peek_at_stack(block_arg ? 1 : 0)
+        splat_array_length = if array.nil?
+          0
+        else
+          array.length
+        end
+
+        if opt_num == 0 && required_num != splat_array_length + argc - 1
+          asm.incr_counter(:send_iseq_splat_arity_error)
+          return CantCompile
+        end
       end
 
       # We will not have CantCompile from here.
@@ -4350,8 +4358,25 @@ def jit_call_iseq(jit, ctx, asm, cme, calling, iseq, frame_type: nil, prev_ep: n
 
       # push_splat_args does stack manipulation so we can no longer side exit
       if splat_array_length
-        asm.incr_counter(:send_iseq_splat)
-        return CantCompile
+        remaining_opt = (opt_num + required_num) - (splat_array_length + (argc - 1))
+
+        if opt_num > 0
+          # We are going to jump to the correct offset based on how many optional
+          # params are remaining.
+          offset = opt_num - remaining_opt
+          start_pc_offset = iseq.body.param.opt_table[offset]
+        end
+        # We are going to assume that the splat fills
+        # all the remaining arguments. In the generated code
+        # we test if this is true and if not side exit.
+        argc = argc - 1 + splat_array_length + remaining_opt
+        push_splat_args(splat_array_length, jit, ctx, asm)
+
+        remaining_opt.times do
+          # We need to push nil for the optional arguments
+          stack_ret = ctx.stack_push
+          asm.mov(stack_ret, Qnil)
+        end
       end
 
       # This is a .send call and we need to adjust the stack
@@ -4730,7 +4755,7 @@ def jit_call_cfunc_with_frame(jit, ctx, asm, cme, flags, argc, block_handler, kn
         # So the number of args should just equal the number of args the cfunc takes.
         # In the generated code we test if this is true and if not side exit.
         argc = cfunc.argc
-        jit_caller_setup_arg_splat(jit, ctx, asm, required_args)
+        push_splat_args(required_args, jit, ctx, asm)
       end
 
       # We will not have side exits from here. Adjust the stack, which was skipped in jit_call_opt_send.
@@ -5344,7 +5369,7 @@ def move_rest_args_to_stack(array, num_args, jit, ctx, asm)
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
-    def jit_caller_setup_arg_splat(jit, ctx, asm, required_args)
+    def push_splat_args(required_args, jit, ctx, asm)
       side_exit = side_exit(jit, ctx)
 
       asm.comment('push_splat_args')

rjit_c.h

@@ -64,7 +64,6 @@ RJIT_RUNTIME_COUNTERS(
     send_iseq_has_no_kw,
     send_iseq_zsuper,
     send_iseq_materialized_block,
-    send_iseq_splat_with_kw,
     send_iseq_has_rest,
     send_iseq_block_arg0_splat,
     send_iseq_kw_call,
@@ -74,6 +73,8 @@ RJIT_RUNTIME_COUNTERS(
     send_iseq_missing_optional_kw,
     send_iseq_too_many_kwargs,
     send_iseq_kwargs_mismatch,
+    send_iseq_splat_with_kw,
+    send_iseq_splat_arity_error,
 
     send_cfunc_variadic,
     send_cfunc_too_many_args,

rjit_c.rb

@@ -1350,7 +1350,6 @@ def C.rb_rjit_runtime_counters
       send_iseq_has_no_kw: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_has_no_kw)")],
       send_iseq_zsuper: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_zsuper)")],
       send_iseq_materialized_block: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_materialized_block)")],
-      send_iseq_splat_with_kw: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_splat_with_kw)")],
       send_iseq_has_rest: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_has_rest)")],
       send_iseq_block_arg0_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_block_arg0_splat)")],
       send_iseq_kw_call: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_kw_call)")],
@@ -1360,6 +1359,8 @@ def C.rb_rjit_runtime_counters
       send_iseq_missing_optional_kw: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_missing_optional_kw)")],
       send_iseq_too_many_kwargs: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_too_many_kwargs)")],
       send_iseq_kwargs_mismatch: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_kwargs_mismatch)")],
+      send_iseq_splat_with_kw: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_splat_with_kw)")],
+      send_iseq_splat_arity_error: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_iseq_splat_arity_error)")],
       send_cfunc_variadic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc_variadic)")],
       send_cfunc_too_many_args: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc_too_many_args)")],
       send_cfunc_ruby_array_varg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_cfunc_ruby_array_varg)")],