From f3a117605c86b44ef48f313dc921c7fb86ff3481 Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Fri, 13 Dec 2024 10:41:04 -0500
Subject: [PATCH] YJIT: Speculate block arg for `c_func_method(&nil)` calls
 (#12326)

A good amount of call sites always pass nil as block argument, but the
nil doesn't show up in the context. Put a runtime guard for those
cases to handle it. Particular relevant for the `ruby-lsp` benchmark in
`yjit-bench`. Up to a 2% speedup across headline benchmarks.

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>
Co-authored-by: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>
Co-authored-by: Aaron Patterson <tenderlove@ruby-lang.org>
Co-authored-by: Kevin Menard <kevin@nirvdrum.com>
Co-authored-by: Randy Stauner <randy.stauner@shopify.com>
---
 yjit/src/codegen.rs | 30 +++++++++++-------------------
 yjit/src/stats.rs   |  1 +
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index e4283b786f..91d3aff532 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -6949,7 +6949,7 @@ fn gen_send_cfunc(
         return None;
     }
 
-    let block_arg_type = if block_arg {
+    let mut block_arg_type = if block_arg {
         Some(asm.ctx.get_opnd_type(StackOpnd(0)))
     } else {
         None
@@ -6957,7 +6957,15 @@ fn gen_send_cfunc(
 
     match block_arg_type {
         Some(Type::Nil | Type::BlockParamProxy) => {
-            // We'll handle this later
+            // We don't need the actual stack value for these
+            asm.stack_pop(1);
+        }
+        Some(Type::Unknown | Type::UnknownImm) if jit.peek_at_stack(&asm.ctx, 0).nil_p() => {
+            // The sample blockarg is nil, so speculate that's the case.
+            asm.cmp(asm.stack_opnd(0), Qnil.into());
+            asm.jne(Target::side_exit(Counter::guard_send_cfunc_block_not_nil));
+            block_arg_type = Some(Type::Nil);
+            asm.stack_pop(1);
         }
         None => {
             // Nothing to do
@@ -6967,23 +6975,7 @@ fn gen_send_cfunc(
             return None;
         }
     }
-
-    match block_arg_type {
-        Some(Type::Nil) => {
-            // We have a nil block arg, so let's pop it off the args
-            asm.stack_pop(1);
-        }
-        Some(Type::BlockParamProxy) => {
-            // We don't need the actual stack value
-            asm.stack_pop(1);
-        }
-        None => {
-            // Nothing to do
-        }
-        _ => {
-            assert!(false);
-        }
-    }
+    let block_arg_type = block_arg_type; // drop `mut`
 
     // Pop the empty kw_splat hash
     if kw_splat {
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 56ecfc5bc1..3a3965577b 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -470,6 +470,7 @@ make_counters! {
     guard_send_str_aref_not_fixnum,
 
     guard_send_cfunc_bad_splat_vargs,
+    guard_send_cfunc_block_not_nil,
 
     guard_invokesuper_me_changed,