From 548a1b2ed778ed5a1fe36178f5a964393ed85fb7 Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Fri, 16 Feb 2024 19:49:43 -0500
Subject: [PATCH] YJIT: Support `**nil` for cfuncs

Similar to the iseq call support. Fairly straight forward.
---
 test/ruby/test_yjit.rb |  4 ++--
 yjit/src/codegen.rs    | 30 +++++++++++++++++++++---------
 yjit/src/stats.rs      |  2 +-
 3 files changed, 24 insertions(+), 12 deletions(-)

diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 3ccc1be1dd3e3d..35036aad4e2a01 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -1574,11 +1574,11 @@ def test_send_polymorphic_method_name
   end
 
   def test_kw_splat_nil
-    assert_compiles(<<~'RUBY', result: %i[ok ok], no_send_fallbacks: true)
+    assert_compiles(<<~'RUBY', result: %i[ok ok ok], no_send_fallbacks: true)
       def id(x) = x
       def kw_fw(arg, **) = id(arg, **)
       def fw(...) = id(...)
-      def use = [fw(:ok), kw_fw(:ok)]
+      def use = [fw(:ok), kw_fw(:ok), :ok.itself(**nil)]
 
       use
     RUBY
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 6ccbfaae666db8..693a5d183a87d4 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -6037,11 +6037,8 @@ fn gen_send_cfunc(
         return None;
     }
 
-    // Don't JIT calls with keyword splat
-    if flags & VM_CALL_KW_SPLAT != 0 {
-        gen_counter_incr(asm, Counter::send_cfunc_kw_splat);
-        return None;
-    }
+    exit_if_kwsplat_non_nil(asm, flags, Counter::send_cfunc_kw_splat_non_nil)?;
+    let kw_splat = flags & VM_CALL_KW_SPLAT != 0;
 
     let kw_arg = unsafe { vm_ci_kwarg(ci) };
     let kw_arg_num = if kw_arg.is_null() {
@@ -6065,7 +6062,11 @@ fn gen_send_cfunc(
     gen_counter_incr(asm, Counter::num_send_cfunc);
 
     // Delegate to codegen for C methods if we have it.
-    if kw_arg.is_null() && flags & VM_CALL_OPT_SEND == 0 && flags & VM_CALL_ARGS_SPLAT == 0 && (cfunc_argc == -1 || argc == cfunc_argc) {
+    if kw_arg.is_null() &&
+            !kw_splat &&
+            flags & VM_CALL_OPT_SEND == 0 &&
+            flags & VM_CALL_ARGS_SPLAT == 0 &&
+            (cfunc_argc == -1 || argc == cfunc_argc) {
         let expected_stack_after = asm.ctx.get_stack_size() as i32 - argc;
         if let Some(known_cfunc_codegen) = lookup_cfunc_codegen(unsafe { (*cme).def }) {
             if perf_call!("gen_send_cfunc: ", known_cfunc_codegen(jit, asm, ocb, ci, cme, block, argc, recv_known_class)) {
@@ -6098,6 +6099,10 @@ fn gen_send_cfunc(
         argc - kw_arg_num + 1
     };
 
+    if kw_splat {
+        passed_argc -= 1;
+    }
+
     // If the argument count doesn't match
     if cfunc_argc >= 0 && cfunc_argc != passed_argc && flags & VM_CALL_ARGS_SPLAT == 0 {
         gen_counter_incr(asm, Counter::send_cfunc_argc_mismatch);
@@ -6147,6 +6152,13 @@ fn gen_send_cfunc(
         }
     }
 
+    if kw_splat {
+        // Only `**nil` is supported right now. Checked in exit_if_kwsplat_non_nil()
+        assert_eq!(Type::Nil, asm.ctx.get_opnd_type(StackOpnd(0)));
+        asm.stack_pop(1);
+        argc -= 1;
+    }
+
     // push_splat_args does stack manipulation so we can no longer side exit
     if flags & VM_CALL_ARGS_SPLAT != 0 {
         assert!(cfunc_argc >= 0);
@@ -6559,7 +6571,7 @@ fn gen_send_iseq(
     exit_if_stack_too_large(iseq)?;
     exit_if_tail_call(asm, ci)?;
     exit_if_has_post(asm, iseq)?;
-    exit_if_kwsplat_non_nil(asm, flags)?;
+    exit_if_kwsplat_non_nil(asm, flags, Counter::send_iseq_kw_splat_non_nil)?;
     exit_if_has_rest_and_captured(asm, iseq_has_rest, captured_opnd)?;
     exit_if_has_kwrest_and_captured(asm, has_kwrest, captured_opnd)?;
     exit_if_has_rest_and_supplying_kws(asm, iseq_has_rest, supplying_kws)?;
@@ -7560,10 +7572,10 @@ fn exit_if_has_post(asm: &mut Assembler, iseq: *const rb_iseq_t) -> Option<()> {
 }
 
 #[must_use]
-fn exit_if_kwsplat_non_nil(asm: &mut Assembler, flags: u32) -> Option<()> {
+fn exit_if_kwsplat_non_nil(asm: &mut Assembler, flags: u32, counter: Counter) -> Option<()> {
     let kw_splat = flags & VM_CALL_KW_SPLAT != 0;
     let kw_splat_stack = StackOpnd((flags & VM_CALL_ARGS_BLOCKARG != 0).into());
-    exit_if(asm, kw_splat && asm.ctx.get_opnd_type(kw_splat_stack) != Type::Nil, Counter::send_iseq_kw_splat_non_nil)
+    exit_if(asm, kw_splat && asm.ctx.get_opnd_type(kw_splat_stack) != Type::Nil, counter)
 }
 
 #[must_use]
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 1472444b8aaf18..8137ccd17c0c96 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -353,7 +353,7 @@ make_counters! {
     send_missing_method,
     send_refined_method,
     send_private_not_fcall,
-    send_cfunc_kw_splat,
+    send_cfunc_kw_splat_non_nil,
     send_cfunc_ruby_array_varg,
     send_cfunc_argc_mismatch,
     send_cfunc_block_arg,