YJIT: Implement VM_CALL_ARGS_BLOCKARG with Proc for ISeq calls

Rack uses this. Guard that the `obj` in `the_call(&obj)` is always going to be a Proc if the compile time sample is a Proc. Co-authored-by: Takashi Kokubun <[email protected]> Co-authored-by: Maxime Chevalier-Boisvert <[email protected]> co-authored-by: Aaron Patterson <[email protected]>
XrXr · Jul 25, 2023 · b9ac410 · b9ac410
1 parent e110401
commit b9ac410
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 5 deletions.
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
@@ -441,6 +441,7 @@ fn main() {
         .allowlist_function("rb_method_basic_definition_p")
         .allowlist_function("rb_yjit_array_len")
         .allowlist_function("rb_obj_class")
+        .allowlist_function("rb_obj_is_proc")
 
         // We define VALUE manually, don't import it
         .blocklist_type("VALUE")

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
@@ -4741,6 +4741,8 @@ enum SpecVal {
     BlockParamProxy,
     PrevEP(*const VALUE),
     PrevEPOpnd(Opnd),
+    /// Block argument that is a material heap object. Proc or Symbol.
+    MaterialBlockArg(Opnd),
 }
 
 struct ControlFrame {
@@ -4814,6 +4816,9 @@ fn gen_push_frame(
 
             block_handler
         }
+        SpecVal::MaterialBlockArg(arg) => {
+            arg
+        }
         SpecVal::PrevEP(prev_ep) => {
             let tagged_prev_ep = (prev_ep as usize) | 1;
             VALUE(tagged_prev_ep).into()
@@ -5484,7 +5489,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 material_block_arg = exit_if_unsupported_block_arg_type(jit, asm, block_arg_type)?;
 
     // Block parameter handling. This mirrors setup_parameters_complex().
     if unsafe { get_iseq_flags_has_block(iseq) } {
@@ -5679,6 +5684,7 @@ fn gen_send_iseq(
 
     // We will not have None from here. You can use stack_pop / stack_pop.
 
+    let mut material_spec_val = SpecVal::None;
     match block_arg_type {
         Some(Type::Nil) => {
             // We have a nil block arg, so let's pop it off the args
@@ -5692,7 +5698,14 @@ fn gen_send_iseq(
             // Nothing to do
         }
         _ => {
-            assert!(false);
+            if material_block_arg {
+                let ret = asm.ccall(rb_obj_is_proc as _, vec![asm.stack_opnd(0)]);
+                asm.cmp(ret, Qfalse.into());
+                asm.je(Target::side_exit(Counter::send_block_arg_type_changed));
+
+                let popped = asm.stack_pop(1);
+                material_spec_val = SpecVal::MaterialBlockArg(asm.load(popped));
+            }
         }
     }
 
@@ -6052,7 +6065,7 @@ fn gen_send_iseq(
     } else {
         None
     };
-
+    
     // Points to the receiver operand on the stack unless a captured environment is used
     let recv = match captured_opnd {
         Some(captured_opnd) => asm.load(Opnd::mem(64, captured_opnd, 0)), // captured->self
@@ -6085,6 +6098,8 @@ fn gen_send_iseq(
         SpecVal::BlockParamProxy
     } else if let Some(block_val) = block {
         SpecVal::BlockISeq(block_val)
+    } else if matches!(material_spec_val, SpecVal::MaterialBlockArg(_)) {
+        material_spec_val
     } else {
         SpecVal::None
     };
@@ -6297,11 +6312,18 @@ 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, block_arg_type: Option<Type>) -> Option<bool> {
     match block_arg_type {
         Some(Type::Nil | Type::BlockParamProxy) => {
             // We'll handle this later
         }
+        Some(_) if {
+            let block_arg = jit.peek_at_stack(&asm.ctx, 0);
+            unsafe { rb_obj_is_proc(block_arg) }.test()
+        } => {
+            // The block argument is a object we can pass as-is
+            return Some(true);
+        }
         None => {
             // Nothing to do
         }
@@ -6310,7 +6332,7 @@ fn exit_if_unsupported_block_arg_type(asm: &mut Assembler, block_arg_type: Optio
             return None
         }
     }
-    Some(())
+    Some(false)
 }
 
 #[must_use]

diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
@@ -1099,6 +1099,7 @@ extern "C" {
     pub fn rb_hash_aref(hash: VALUE, key: VALUE) -> VALUE;
     pub fn rb_hash_aset(hash: VALUE, key: VALUE, val: VALUE) -> VALUE;
     pub fn rb_hash_bulk_insert(argc: ::std::os::raw::c_long, argv: *const VALUE, hash: VALUE);
+    pub fn rb_obj_is_proc(recv: VALUE) -> VALUE;
     pub fn rb_sym2id(obj: VALUE) -> ID;
     pub fn rb_id2sym(id: ID) -> VALUE;
     pub fn rb_intern(name: *const ::std::os::raw::c_char) -> ID;

diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
@@ -210,6 +210,7 @@ make_counters! {
     send_args_splat_super,
     send_iseq_zsuper,
     send_block_arg,
+    send_block_arg_type_changed,
     send_ivar_set_method,
     send_zsuper_method,
     send_undef_method,