Fix splat args (#6385)

* Fix splat args

Cfuncs were not working properly so I disabled them right now.

There were some checks above that were also actually preventing splat args from being called.

Finally I did some basic code cleanup after realizing I didn't need to mutate argc so much

* Add can't compile for direct cfunc splat call

* Fix typo

* Update yjit/src/codegen.rs

Co-authored-by: Maxime Chevalier-Boisvert <maximechevalierb@gmail.com>

2 files changed, 35 insertions, 29 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index f1586cee4b..71831a96ed 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4023,6 +4023,11 @@ fn gen_send_cfunc(
         return CantCompile;
     }
 
+    if flags & VM_CALL_ARGS_SPLAT != 0  {
+        gen_counter_incr!(asm, send_args_splat_cfunc);
+        return CantCompile;
+    }
+
     let kw_arg = unsafe { vm_ci_kwarg(ci) };
     let kw_arg_num = if kw_arg.is_null() {
         0
@@ -4060,18 +4065,6 @@ fn gen_send_cfunc(
     asm.cmp(CFP, stack_limit);
     asm.jbe(counted_exit!(ocb, side_exit, send_se_cf_overflow).into());
 
-    if flags & VM_CALL_ARGS_SPLAT != 0 {
-        if cfunc_argc > 0 {
-            // push_splat_args does a ctx.stack_push so we can no longer side exit
-            argc = push_splat_args(argc, ctx, asm, ocb, side_exit, cfunc_argc as u32);
-        } else {
-            // This is a variadic c function and we'd need to pop
-            // each of the elements off the array, but the array may be dynamically sized
-            gen_counter_incr!(asm, send_args_splat_variadic);
-            return CantCompile
-        }
-    }
-
     // Number of args which will be passed through to the callee
     // This is adjusted by the kwargs being combined into a hash.
     let passed_argc = if kw_arg.is_null() {
@@ -4257,19 +4250,16 @@ fn gen_return_branch(
 }
 
 
+
 /// Pushes arguments from an array to the stack that are passed with a splat (i.e. *args)
 /// It optimistically compiles to a static size that is the exact number of arguments
 /// needed for the function.
-fn push_splat_args(argc: i32, ctx: &mut Context, asm: &mut Assembler, ocb: &mut OutlinedCb, side_exit: CodePtr, num_params: u32) -> i32 {
-
-    let mut argc = argc;
+fn push_splat_args(required_args: i32, ctx: &mut Context, asm: &mut Assembler, ocb: &mut OutlinedCb, side_exit: CodePtr) {
 
     asm.comment("push_splat_args");
 
     let array_opnd = ctx.stack_opnd(0);
 
-    argc = argc - 1;
-
     let array_reg = asm.load(array_opnd);
     guard_object_is_heap(
         asm,
@@ -4299,11 +4289,9 @@ fn push_splat_args(argc: i32, ctx: &mut Context, asm: &mut Assembler, ocb: &mut
     );
     let array_len_opnd = asm.csel_nz(emb_len_opnd, array_len_opnd);
 
-    let required_args = num_params - argc as u32;
-
     // Only handle the case where the number of values in the array is equal to the number requested
     asm.cmp(array_len_opnd, required_args.into());
-    asm.jne(counted_exit!(ocb, side_exit, send_splatarray_rhs_too_small).into());
+    asm.jne(counted_exit!(ocb, side_exit, send_splatarray_length_not_equal).into());
 
     let array_opnd = ctx.stack_pop(1);
 
@@ -4329,10 +4317,8 @@ fn push_splat_args(argc: i32, ctx: &mut Context, asm: &mut Assembler, ocb: &mut
         for i in (0..required_args as i32) {
             let top = ctx.stack_push(Type::Unknown);
             asm.mov(top, Opnd::mem(64, ary_opnd, i * (SIZEOF_VALUE as i32)));
-            argc += 1;
         }
     }
-    argc
 }
 
 fn gen_send_iseq(
@@ -4448,12 +4434,25 @@ fn gen_send_iseq(
     let opt_num = unsafe { get_iseq_body_param_opt_num(iseq) };
     let opts_missing: i32 = opt_num - opts_filled;
 
+
     if opt_num > 0 && flags & VM_CALL_ARGS_SPLAT != 0 {
         gen_counter_incr!(asm, send_iseq_complex_callee);
         return CantCompile;
     }
 
-    if opts_filled < 0 || opts_filled > opt_num {
+    if doing_kw_call && flags & VM_CALL_ARGS_SPLAT != 0 {
+        gen_counter_incr!(asm, send_iseq_complex_callee);
+        return CantCompile;
+    }
+
+    if opts_filled < 0 && flags & VM_CALL_ARGS_SPLAT == 0  {
+        // Too few arguments and no splat to make up for it
+        gen_counter_incr!(asm, send_iseq_arity_error);
+        return CantCompile;
+    }
+
+    if opts_filled > opt_num {
+        // Too many arguments
         gen_counter_incr!(asm, send_iseq_arity_error);
         return CantCompile;
     }
@@ -4596,9 +4595,14 @@ fn gen_send_iseq(
     asm.cmp(CFP, stack_limit);
     asm.jbe(counted_exit!(ocb, side_exit, send_se_cf_overflow).into());
 
-    // push_splat_args does a ctx.stack_push so we can no longer side exit
+    // push_splat_args does stack manipulation so we can no longer side exit
     if flags & VM_CALL_ARGS_SPLAT != 0 {
-        argc = push_splat_args(argc, ctx, asm, ocb, side_exit, num_params);
+        let required_args = num_params as i32 - (argc - 1);
+        // 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 = num_params as i32;
+        push_splat_args(required_args, ctx, asm, ocb, side_exit)
     }
 
     if doing_kw_call {
@@ -5077,8 +5081,9 @@ fn gen_send_general(
     // To handle the aliased method case (VM_METHOD_TYPE_ALIAS)
     loop {
         let def_type = unsafe { get_cme_def_type(cme) };
-        if flags & VM_CALL_ARGS_SPLAT != 0 && (def_type != VM_METHOD_TYPE_ISEQ && def_type != VM_METHOD_TYPE_CFUNC)  {
-            // We can't handle splat calls to non-iseq methods
+
+        if flags & VM_CALL_ARGS_SPLAT != 0 && def_type != VM_METHOD_TYPE_ISEQ  {
+            gen_counter_incr!(asm, send_args_splat_non_iseq);
             return CantCompile;
         }
 
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 56aa0b4081..5fc83c3896 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -193,9 +193,10 @@ make_counters! {
     send_getter_arity,
     send_se_cf_overflow,
     send_se_protected_check_failed,
-    send_splatarray_rhs_too_small,
+    send_splatarray_length_not_equal,
     send_splat_not_array,
-    send_args_splat_variadic,
+    send_args_splat_non_iseq,
+    send_args_splat_cfunc,
     send_iseq_ruby2_keywords,
 
     traced_cfunc_return,