ZJIT: Support invokeblock insn (ruby#14496)

ZJIT: Support invokeblock insn

Basically like other send-ish insns, except that it doesn't pop the receiver.

Author: st0012

test/ruby/test_zjit.rb

@@ -2584,6 +2584,46 @@ def test(x)
     }, insns: [:opt_case_dispatch]
   end
 
+  def test_invokeblock
+    assert_compiles '42', %q{
+      def test
+        yield
+      end
+      test { 42 }
+    }, insns: [:invokeblock]
+  end
+
+  def test_invokeblock_with_args
+    assert_compiles '3', %q{
+      def test(x, y)
+        yield x, y
+      end
+      test(1, 2) { |a, b| a + b }
+    }, insns: [:invokeblock]
+  end
+
+  def test_invokeblock_no_block_given
+    assert_compiles ':error', %q{
+      def test
+        yield rescue :error
+      end
+      test
+    }, insns: [:invokeblock]
+  end
+
+  def test_invokeblock_multiple_yields
+    assert_compiles "[1, 2, 3]", %q{
+      results = []
+      def test
+        yield 1
+        yield 2
+        yield 3
+      end
+      test { |x| results << x }
+      results
+    }, insns: [:invokeblock]
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT

zjit/src/codegen.rs

@@ -369,6 +369,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
             gen_send_without_block(jit, asm, *cd, &function.frame_state(*state)),
         Insn::SendWithoutBlockDirect { cme, iseq, self_val, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(self_val), opnds!(args), &function.frame_state(*state)),
         &Insn::InvokeSuper { cd, blockiseq, state, .. } => gen_invokesuper(jit, asm, cd, blockiseq, &function.frame_state(state)),
+        Insn::InvokeBlock { cd, state, .. } => gen_invoke_block(jit, asm, *cd, &function.frame_state(*state)),
         // Ensure we have enough room fit ec, self, and arguments
         // TODO remove this check when we have stack args (we can use Time.new to test it)
         Insn::InvokeBuiltin { bf, state, .. } if bf.argc + 2 > (C_ARG_OPNDS.len() as i32) => return Err(*state),
@@ -1093,6 +1094,31 @@ fn gen_send_without_block_direct(
     ret
 }
 
+/// Compile for invokeblock
+fn gen_invoke_block(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    cd: *const rb_call_data,
+    state: &FrameState,
+) -> lir::Opnd {
+    gen_incr_counter(asm, Counter::dynamic_send_count);
+
+    // Save PC and SP, spill locals and stack
+    gen_prepare_call_with_gc(asm, state);
+    gen_save_sp(asm, state.stack().len());
+    gen_spill_locals(jit, asm, state);
+    gen_spill_stack(jit, asm, state);
+
+    asm_comment!(asm, "call invokeblock");
+    unsafe extern "C" {
+        fn rb_vm_invokeblock(ec: EcPtr, cfp: CfpPtr, cd: VALUE) -> VALUE;
+    }
+    asm.ccall(
+        rb_vm_invokeblock as *const u8,
+        vec![EC, CFP, (cd as usize).into()],
+    )
+}
+
 /// Compile a dynamic dispatch for `super`
 fn gen_invokesuper(
     jit: &mut JITState,

zjit/src/hir.rs

@@ -585,6 +585,7 @@ pub enum Insn {
     SendWithoutBlock { self_val: InsnId, cd: *const rb_call_data, args: Vec<InsnId>, state: InsnId },
     Send { self_val: InsnId, cd: *const rb_call_data, blockiseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
     InvokeSuper { self_val: InsnId, cd: *const rb_call_data, blockiseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
+    InvokeBlock { cd: *const rb_call_data, args: Vec<InsnId>, state: InsnId },
 
     /// Optimized ISEQ call
     SendWithoutBlockDirect {
@@ -845,6 +846,13 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             }
+            Insn::InvokeBlock { args, .. } => {
+                write!(f, "InvokeBlock")?;
+                for arg in args {
+                    write!(f, ", {arg}")?;
+                }
+                Ok(())
+            }
             Insn::InvokeBuiltin { bf, args, .. } => {
                 write!(f, "InvokeBuiltin {}", unsafe { CStr::from_ptr(bf.name) }.to_str().unwrap())?;
                 for arg in args {
@@ -1349,6 +1357,11 @@ impl Function {
                 args: find_vec!(args),
                 state,
             },
+            &InvokeBlock { cd, ref args, state } => InvokeBlock {
+                cd,
+                args: find_vec!(args),
+                state,
+            },
             &InvokeBuiltin { bf, ref args, state, return_type } => InvokeBuiltin { bf, args: find_vec!(args), state, return_type },
             &ArrayDup { val, state } => ArrayDup { val: find!(val), state },
             &HashDup { val, state } => HashDup { val: find!(val), state },
@@ -1463,6 +1476,7 @@ impl Function {
             Insn::SendWithoutBlockDirect { .. } => types::BasicObject,
             Insn::Send { .. } => types::BasicObject,
             Insn::InvokeSuper { .. } => types::BasicObject,
+            Insn::InvokeBlock { .. } => types::BasicObject,
             Insn::InvokeBuiltin { return_type, .. } => return_type.unwrap_or(types::BasicObject),
             Insn::Defined { pushval, .. } => Type::from_value(*pushval).union(types::NilClass),
             Insn::DefinedIvar { .. } => types::BasicObject,
@@ -2276,7 +2290,8 @@ impl Function {
                 worklist.extend(args);
                 worklist.push_back(state);
             }
-            &Insn::InvokeBuiltin { ref args, state, .. } => {
+            &Insn::InvokeBuiltin { ref args, state, .. }
+            | &Insn::InvokeBlock { ref args, state, .. } => {
                 worklist.extend(args);
                 worklist.push_back(state)
             }
@@ -3651,6 +3666,21 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         }
                     }
                 }
+                YARVINSN_invokeblock => {
+                    let cd: *const rb_call_data = get_arg(pc, 0).as_ptr();
+                    let call_info = unsafe { rb_get_call_data_ci(cd) };
+                    if let Err(call_type) = unknown_call_type(unsafe { rb_vm_ci_flag(call_info) }) {
+                        // Unknown call type; side-exit into the interpreter
+                        let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                        fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::UnhandledCallType(call_type) });
+                        break;  // End the block
+                    }
+                    let argc = unsafe { vm_ci_argc((*cd).ci) };
+                    let args = state.stack_pop_n(argc as usize)?;
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    let result = fun.push_insn(block, Insn::InvokeBlock { cd, args, state: exit_id });
+                    state.stack_push(result);
+                }
                 YARVINSN_getglobal => {
                     let id = ID(get_arg(pc, 0).as_u64());
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
@@ -5835,11 +5865,14 @@ mod tests {
           v19:CBool = Test v16
           IfFalse v19, bb1(v0, v1, v2, v3, v4, v10)
           PatchPoint NoEPEscape(open)
-          SideExit UnhandledYARVInsn(invokeblock)
-        bb1(v27:BasicObject, v28:BasicObject, v29:BasicObject, v30:BasicObject, v31:BasicObject, v32:BasicObject):
+          v26:BasicObject = InvokeBlock, v10
+          v30:BasicObject = InvokeBuiltin dir_s_close, v0, v10
+          CheckInterrupts
+          Return v26
+        bb1(v36:BasicObject, v37:BasicObject, v38:BasicObject, v39:BasicObject, v40:BasicObject, v41:BasicObject):
           PatchPoint NoEPEscape(open)
           CheckInterrupts
-          Return v32
+          Return v41
         ");
     }
 
@@ -6028,6 +6061,38 @@ mod tests {
           Throw TAG_BREAK, v6
         ");
     }
+
+    #[test]
+    fn test_invokeblock() {
+        eval(r#"
+            def test
+              yield
+            end
+        "#);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject):
+          v5:BasicObject = InvokeBlock
+          CheckInterrupts
+          Return v5
+        ");
+    }
+
+    #[test]
+    fn test_invokeblock_with_args() {
+        eval(r#"
+            def test(x, y)
+              yield x, y
+            end
+        "#);
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject, v1:BasicObject, v2:BasicObject):
+          v7:BasicObject = InvokeBlock, v1, v2
+          CheckInterrupts
+          Return v7
+        ");
+    }
 }
 
 #[cfg(test)]