YJIT: Handle splat with opt more fully (#7209)

* YJIT: Handle splat with opt more fully

* Update yjit/src/codegen.rs

---------

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

5 files changed, 44 insertions, 15 deletions

diff --git a/yjit.c b/yjit.c
index a1710f9c05..bc25880de2 100644
--- a/yjit.c
+++ b/yjit.c
@@ -103,6 +103,12 @@ rb_yjit_mark_unused(void *mem_block, uint32_t mem_size)
     return mprotect(mem_block, mem_size, PROT_NONE) == 0;
 }
 
+long
+rb_yjit_array_len(VALUE a)
+{
+    return rb_array_len(a);
+}
+
 // `start` is inclusive and `end` is exclusive.
 void
 rb_yjit_icache_invalidate(void *start, void *end)
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 229b02bdcd..35b09e1abc 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -410,6 +410,7 @@ fn main() {
         .allowlist_function("rb_METHOD_ENTRY_VISI")
         .allowlist_function("rb_RCLASS_ORIGIN")
         .allowlist_function("rb_method_basic_definition_p")
+        .allowlist_function("rb_yjit_array_len")
 
         // We define VALUE manually, don't import it
         .blocklist_type("VALUE")
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 7fd94f746f..83ce7e2ea6 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -5222,24 +5222,14 @@ fn gen_send_iseq(
         return CantCompile;
     }
 
+    // We will handle splat case later
     if opt_num > 0 && flags & VM_CALL_ARGS_SPLAT == 0 {
         num_params -= opts_missing as u32;
         unsafe {
             let opt_table = get_iseq_body_param_opt_table(iseq);
             start_pc_offset = (*opt_table.offset(opts_filled as isize)).as_u32();
         }
-    } else if opt_num > 0 && flags & VM_CALL_ARGS_SPLAT != 0 {
-        // We are going to assume that args_splat fills all of the optional arguments.
-        // We should actually get the array length instead at compile time.
-        // We don't set num_params here because we use it for the splat.
-        // If our assumption is wrong, we will exit early in the splat code.
-        // We could do that a bit smarter and not exit but instead change
-        // the offset we jump to. But we aren't currently doing that.
-        unsafe {
-            let opt_table = get_iseq_body_param_opt_table(iseq);
-            start_pc_offset = (*opt_table.offset(opt_num as isize)).as_u32();
-        };
-    };
+    }
 
     if doing_kw_call {
         // Here we're calling a method with keyword arguments and specifying
@@ -5393,12 +5383,42 @@ fn gen_send_iseq(
 
     // push_splat_args does stack manipulation so we can no longer side exit
     if flags & VM_CALL_ARGS_SPLAT != 0 {
-        let required_args = num_params - (argc as u32 - 1);
+
+        let array = jit_peek_at_stack(jit, ctx, if block_arg { 1 } else { 0 }) ;
+        let array_length = if array == Qnil {
+            0
+        } else {
+            unsafe { rb_yjit_array_len(array) as u32}
+        };
+
+        if opt_num == 0 && required_num != array_length as i32 {
+            gen_counter_incr!(asm, send_iseq_splat_arity_error);
+            return CantCompile;
+        }
+
+        let remaining_opt = (opt_num as u32 + required_num as u32).saturating_sub(array_length + (argc as u32 - 1));
+
+        if opt_num > 0 {
+
+            // We are going to jump to the correct offset based on how many optional
+            // params are remaining.
+            unsafe {
+                let opt_table = get_iseq_body_param_opt_table(iseq);
+                let offset = (opt_num - remaining_opt as i32) as isize;
+                start_pc_offset = (*opt_table.offset(offset)).as_u32();
+            };
+        }
         // 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)
+        argc = argc - 1 + array_length as i32 + remaining_opt as i32;
+        push_splat_args(array_length, ctx, asm, ocb, side_exit);
+
+        for _ in 0..remaining_opt as u32 {
+            // We need to push nil for the optional arguments
+            let stack_ret = ctx.stack_push(Type::Unknown);
+            asm.mov(stack_ret, Qnil.into());
+        }
     }
 
     // This is a .send call and we need to adjust the stack
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index b8a8c91f38..ea57ea67fd 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1194,6 +1194,7 @@ extern "C" {
     pub fn rb_yjit_mark_writable(mem_block: *mut ::std::os::raw::c_void, mem_size: u32) -> bool;
     pub fn rb_yjit_mark_executable(mem_block: *mut ::std::os::raw::c_void, mem_size: u32);
     pub fn rb_yjit_mark_unused(mem_block: *mut ::std::os::raw::c_void, mem_size: u32) -> bool;
+    pub fn rb_yjit_array_len(a: VALUE) -> ::std::os::raw::c_long;
     pub fn rb_yjit_icache_invalidate(
         start: *mut ::std::os::raw::c_void,
         end: *mut ::std::os::raw::c_void,
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index ababf60cc9..bd75d7fc05 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -222,6 +222,7 @@ make_counters! {
     send_args_splat_cfunc_var_args,
     send_args_splat_cfunc_zuper,
     send_args_splat_cfunc_ruby2_keywords,
+    send_iseq_splat_arity_error,
     send_iseq_ruby2_keywords,
     send_send_not_imm,
     send_send_wrong_args,