YJIT: Implement VM_CALL_ARGS_BLOCKARG with Proc for ISeq calls

Rack uses this. Speculate that the `obj` in `the_call(&obj)`
will be a proc when the compile-time sample is a proc.

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>

3 files changed, 67 insertions, 18 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 77b8747e3c..7e44a05f00 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -5068,6 +5068,9 @@ enum SpecVal {
     BlockParamProxy,
     PrevEP(*const VALUE),
     PrevEPOpnd(Opnd),
+    // To avoid holding values across C calls for complex calls,
+    // we might need to set the SpecVal earlier in the call sequence
+    AlreadySet,
 }
 
 // Each variant represents a branch in vm_caller_setup_arg_block.
@@ -5160,9 +5163,14 @@ fn gen_push_frame(
         }
         SpecVal::PrevEPOpnd(ep_opnd) => {
             asm.or(ep_opnd, 1.into())
-        },
+        }
+        SpecVal::AlreadySet => 0.into(), // unused
     };
-    asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), specval);
+    if let SpecVal::AlreadySet = frame.specval {
+        asm.comment("specval should have been set");
+    } else {
+        asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), specval);
+    }
 
     // Write env flags at sp[-1]
     // sp[-1] = frame_type;
@@ -5794,13 +5802,7 @@ fn gen_send_iseq(
     }
     let mut opts_missing: i32 = opt_num - opts_filled;
 
-
     let block_arg = flags & VM_CALL_ARGS_BLOCKARG != 0;
-    let block_arg_type = if block_arg {
-        Some(asm.ctx.get_opnd_type(StackOpnd(0)))
-    } else {
-        None
-    };
 
     exit_if_stack_too_large(iseq)?;
     exit_if_tail_call(asm, ci)?;
@@ -5816,7 +5818,7 @@ fn gen_send_iseq(
     exit_if_wrong_number_arguments(asm, opts_filled, flags, opt_num, iseq_has_rest)?;
     exit_if_doing_kw_and_opts_missing(asm, doing_kw_call, opts_missing)?;
     exit_if_has_rest_and_optional_and_block(asm, iseq_has_rest, opt_num, iseq, block_arg)?;
-    exit_if_unsupported_block_arg_type(asm, block_arg_type)?;
+    let block_arg_type = exit_if_unsupported_block_arg_type(jit, asm, block_arg)?;
 
     // Block parameter handling. This mirrors setup_parameters_complex().
     if iseq_has_block_param {
@@ -6003,12 +6005,35 @@ fn gen_send_iseq(
             // We don't need the actual stack value
             asm.stack_pop(1);
         }
+        Some(Type::TProc) => {
+            // Place the proc as the block handler. We do this early because
+            // the block arg being at the top of the stack gets in the way of
+            // rest param handling later. Also, since there are C calls that
+            // come later, we can't hold this value in a register and place it
+            // near the end when we push a new control frame.
+            asm.comment("guard block arg is a proc");
+            // Simple predicate, no need for jit_prepare_routine_call().
+            let is_proc = asm.ccall(rb_obj_is_proc as _, vec![asm.stack_opnd(0)]);
+            asm.cmp(is_proc, Qfalse.into());
+            jit_chain_guard(
+                JCC_JE,
+                jit,
+                asm,
+                ocb,
+                SEND_MAX_CHAIN_DEPTH,
+                Counter::guard_send_block_arg_type,
+            );
+
+            let proc = asm.stack_pop(1);
+            let callee_ep = -argc + num_locals + VM_ENV_DATA_SIZE as i32 - 1;
+            let callee_specval = callee_ep + VM_ENV_DATA_INDEX_SPECVAL;
+            let callee_specval = asm.ctx.sp_opnd(callee_specval as isize * SIZEOF_VALUE as isize);
+            asm.store(callee_specval, proc);
+        }
         None => {
             // Nothing to do
         }
-        _ => {
-            assert!(false);
-        }
+        _ => unreachable!(),
     }
 
     let builtin_attrs = unsafe { rb_yjit_iseq_builtin_attrs(iseq) };
@@ -6434,6 +6459,8 @@ fn gen_send_iseq(
         SpecVal::PrevEPOpnd(ep_opnd)
     } else if block_arg_type == Some(Type::BlockParamProxy) {
         SpecVal::BlockParamProxy
+    } else if let Some(Type::TProc) = block_arg_type {
+        SpecVal::AlreadySet
     } else if let Some(block_handler) = block {
         SpecVal::BlockHandler(block_handler)
     } else {
@@ -6642,20 +6669,35 @@ fn exit_if_has_rest_and_optional_and_block(asm: &mut Assembler, iseq_has_rest: b
 }
 
 #[must_use]
-fn exit_if_unsupported_block_arg_type(asm: &mut Assembler, block_arg_type: Option<Type>) -> Option<()> {
+fn exit_if_unsupported_block_arg_type(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    supplying_block_arg: bool
+) -> Option<Option<Type>> {
+    let block_arg_type = if supplying_block_arg {
+        asm.ctx.get_opnd_type(StackOpnd(0))
+    } else {
+        // Passing no block argument
+        return Some(None);
+    };
+
     match block_arg_type {
-        Some(Type::Nil | Type::BlockParamProxy) => {
+        Type::Nil | Type::BlockParamProxy => {
             // We'll handle this later
+            Some(Some(block_arg_type))
         }
-        None => {
-            // Nothing to do
+        _ if {
+            let sample_block_arg = jit.peek_at_stack(&asm.ctx, 0);
+            unsafe { rb_obj_is_proc(sample_block_arg) }.test()
+        } => {
+            // Speculate that we'll have a proc as the block arg
+            Some(Some(Type::TProc))
         }
         _ => {
             gen_counter_incr(asm, Counter::send_block_arg);
-            return None
+            None
         }
     }
-    Some(())
 }
 
 #[must_use]
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index 89061ac4e5..40646ebd34 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -59,6 +59,8 @@ pub enum Type {
     TArray, // An object with the T_ARRAY flag set, possibly an rb_cArray
     CArray, // An un-subclassed string of type rb_cArray (can have instance vars in some cases)
 
+    TProc, // A proc object. Could be an instance of a subclass of ::rb_cProc
+
     BlockParamProxy, // A special sentinel value indicating the block parameter should be read from
                      // the current surrounding cfp
 }
@@ -110,6 +112,8 @@ impl Type {
                 RUBY_T_ARRAY => Type::TArray,
                 RUBY_T_HASH => Type::Hash,
                 RUBY_T_STRING => Type::TString,
+                #[cfg(not(test))]
+                RUBY_T_DATA if unsafe { rb_obj_is_proc(val).test() } => Type::TProc,
                 _ => Type::UnknownHeap,
             }
         }
@@ -155,6 +159,7 @@ impl Type {
             Type::TString => true,
             Type::CString => true,
             Type::BlockParamProxy => true,
+            Type::TProc => true,
             _ => false,
         }
     }
@@ -189,6 +194,7 @@ impl Type {
             Type::Hash => Some(RUBY_T_HASH),
             Type::ImmSymbol | Type::HeapSymbol => Some(RUBY_T_SYMBOL),
             Type::TString | Type::CString => Some(RUBY_T_STRING),
+            Type::TProc => Some(RUBY_T_DATA),
             Type::Unknown | Type::UnknownImm | Type::UnknownHeap => None,
             Type::BlockParamProxy => None,
         }
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index ef4d25d051..27e3c6d08b 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -296,6 +296,7 @@ make_counters! {
     invokeblock_symbol,
 
     // Method calls that exit to the interpreter
+    guard_send_block_arg_type,
     guard_send_klass_megamorphic,
     guard_send_se_cf_overflow,
     guard_send_se_protected_check_failed,