ZJIT: Compile sendforward with dynamic dispatch (ruby#14501)

* ZJIT: Compile sendforward with dynamic dispatch

* Reload locals only if it has blockiseq

* Add a test case of ... with other args

Author: k0kubun

test/ruby/test_zjit.rb

@@ -478,7 +478,7 @@ def test(...) = 1
   end
 
   def test_sendforward
-    assert_runs '[1, 2]', %q{
+    assert_compiles '[1, 2]', %q{
       def callee(a, b) = [a, b]
       def test(...) = callee(...)
       test(1, 2)

zjit/src/codegen.rs

@@ -365,6 +365,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::IfTrue { val, target } => no_output!(gen_if_true(jit, asm, opnd!(val), target)),
         Insn::IfFalse { val, target } => no_output!(gen_if_false(jit, asm, opnd!(val), target)),
         &Insn::Send { cd, blockiseq, state, .. } => gen_send(jit, asm, cd, blockiseq, &function.frame_state(state)),
+        &Insn::SendForward { cd, blockiseq, state, .. } => gen_send_forward(jit, asm, cd, blockiseq, &function.frame_state(state)),
         Insn::SendWithoutBlock { cd, state, def_type, .. } => gen_send_without_block(jit, asm, *cd, *def_type, &function.frame_state(*state)),
         // Give up SendWithoutBlockDirect for 6+ args since asm.ccall() doesn't support it.
         Insn::SendWithoutBlockDirect { cd, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self
@@ -1041,6 +1042,29 @@ fn gen_send(
     )
 }
 
+/// Compile a dynamic dispatch with `...`
+fn gen_send_forward(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    cd: *const rb_call_data,
+    blockiseq: IseqPtr,
+    state: &FrameState,
+) -> lir::Opnd {
+    gen_incr_counter(asm, Counter::dynamic_send_count);
+    gen_incr_counter(asm, Counter::dynamic_send_type_send_forward);
+
+    gen_prepare_non_leaf_call(jit, asm, state);
+
+    asm_comment!(asm, "call #{} with dynamic dispatch", ruby_call_method_name(cd));
+    unsafe extern "C" {
+        fn rb_vm_sendforward(ec: EcPtr, cfp: CfpPtr, cd: VALUE, blockiseq: IseqPtr) -> VALUE;
+    }
+    asm.ccall(
+        rb_vm_sendforward as *const u8,
+        vec![EC, CFP, (cd as usize).into(), VALUE(blockiseq as usize).into()],
+    )
+}
+
 /// Compile a dynamic dispatch without block
 fn gen_send_without_block(
     jit: &mut JITState,

zjit/src/hir.rs

@@ -642,6 +642,7 @@ pub enum Insn {
         state: InsnId,
     },
     Send { recv: InsnId, cd: *const rb_call_data, blockiseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
+    SendForward { recv: InsnId, cd: *const rb_call_data, blockiseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
     InvokeSuper { recv: InsnId, cd: *const rb_call_data, blockiseq: IseqPtr, args: Vec<InsnId>, state: InsnId },
     InvokeBlock { cd: *const rb_call_data, args: Vec<InsnId>, state: InsnId },
 
@@ -902,6 +903,13 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             }
+            Insn::SendForward { cd, args, blockiseq, .. } => {
+                write!(f, "SendForward {:p}, :{}", self.ptr_map.map_ptr(blockiseq), ruby_call_method_name(*cd))?;
+                for arg in args {
+                    write!(f, ", {arg}")?;
+                }
+                Ok(())
+            }
             Insn::InvokeSuper { recv, blockiseq, args, .. } => {
                 write!(f, "InvokeSuper {recv}, {:p}", self.ptr_map.map_ptr(blockiseq))?;
                 for arg in args {
@@ -1422,6 +1430,13 @@ impl Function {
                 args: find_vec!(args),
                 state,
             },
+            &SendForward { recv, cd, blockiseq, ref args, state } => SendForward {
+                recv: find!(recv),
+                cd,
+                blockiseq,
+                args: find_vec!(args),
+                state,
+            },
             &InvokeSuper { recv, cd, blockiseq, ref args, state } => InvokeSuper {
                 recv: find!(recv),
                 cd,
@@ -1552,6 +1567,7 @@ impl Function {
             Insn::SendWithoutBlock { .. } => types::BasicObject,
             Insn::SendWithoutBlockDirect { .. } => types::BasicObject,
             Insn::Send { .. } => types::BasicObject,
+            Insn::SendForward { .. } => types::BasicObject,
             Insn::InvokeSuper { .. } => types::BasicObject,
             Insn::InvokeBlock { .. } => types::BasicObject,
             Insn::InvokeBuiltin { return_type, .. } => return_type.unwrap_or(types::BasicObject),
@@ -2430,6 +2446,7 @@ impl Function {
                 worklist.push_back(state);
             }
             &Insn::Send { recv, ref args, state, .. }
+            | &Insn::SendForward { recv, ref args, state, .. }
             | &Insn::SendWithoutBlock { recv, ref args, state, .. }
             | &Insn::CCallVariadic { recv, ref args, state, .. }
             | &Insn::SendWithoutBlockDirect { recv, ref args, state, .. }
@@ -3774,13 +3791,44 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let send = fun.push_insn(block, Insn::Send { recv, cd, blockiseq, args, state: exit_id });
                     state.stack_push(send);
 
-                    // Reload locals that may have been modified by the blockiseq.
-                    // TODO: Avoid reloading locals that are not referenced by the blockiseq
-                    // or not used after this. Max thinks we could eventually DCE them.
-                    for local_idx in 0..state.locals.len() {
-                        let ep_offset = local_idx_to_ep_offset(iseq, local_idx) as u32;
-                        let val = fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0 });
-                        state.setlocal(ep_offset, val);
+                    if !blockiseq.is_null() {
+                        // Reload locals that may have been modified by the blockiseq.
+                        // TODO: Avoid reloading locals that are not referenced by the blockiseq
+                        // or not used after this. Max thinks we could eventually DCE them.
+                        for local_idx in 0..state.locals.len() {
+                            let ep_offset = local_idx_to_ep_offset(iseq, local_idx) as u32;
+                            let val = fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0 });
+                            state.setlocal(ep_offset, val);
+                        }
+                    }
+                }
+                YARVINSN_sendforward => {
+                    let cd: *const rb_call_data = get_arg(pc, 0).as_ptr();
+                    let blockiseq: IseqPtr = get_arg(pc, 1).as_iseq();
+                    let call_info = unsafe { rb_get_call_data_ci(cd) };
+                    let flags = unsafe { rb_vm_ci_flag(call_info) };
+                    let forwarding = (flags & VM_CALL_FORWARDING) != 0;
+                    if let Err(call_type) = unhandled_call_type(flags) {
+                        // Can't handle the 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 + usize::from(forwarding))?;
+                    let recv = state.stack_pop()?;
+                    let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    let send_forward = fun.push_insn(block, Insn::SendForward { recv, cd, blockiseq, args, state: exit_id });
+                    state.stack_push(send_forward);
+
+                    if !blockiseq.is_null() {
+                        // Reload locals that may have been modified by the blockiseq.
+                        for local_idx in 0..state.locals.len() {
+                            let ep_offset = local_idx_to_ep_offset(iseq, local_idx) as u32;
+                            let val = fun.push_insn(block, Insn::GetLocal { ep_offset, level: 0 });
+                            state.setlocal(ep_offset, val);
+                        }
                     }
                 }
                 YARVINSN_invokesuper => {
@@ -5315,7 +5363,6 @@ mod tests {
         fn test@<compiled>:2:
         bb0(v0:BasicObject, v1:BasicObject):
           v6:BasicObject = Send v0, 0x1000, :foo, v1
-          v7:BasicObject = GetLocal l0, EP@3
           CheckInterrupts
           Return v6
         ");
@@ -5457,14 +5504,33 @@ mod tests {
     }
 
     #[test]
-    fn test_cant_compile_forwarding() {
+    fn test_compile_forwarding() {
         eval("
             def test(...) = foo(...)
         ");
         assert_snapshot!(hir_string("test"), @r"
         fn test@<compiled>:2:
         bb0(v0:BasicObject, v1:BasicObject):
-          SideExit UnhandledYARVInsn(sendforward)
+          v6:BasicObject = SendForward 0x1000, :foo, v1
+          CheckInterrupts
+          Return v6
+        ");
+    }
+
+    #[test]
+    fn test_compile_triple_dots_with_positional_args() {
+        eval("
+            def test(a, ...) = foo(a, ...)
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:2:
+        bb0(v0:BasicObject, v1:BasicObject, v2:ArrayExact, v3:BasicObject, v4:BasicObject):
+          v5:NilClass = Const Value(nil)
+          v10:ArrayExact = ToArray v2
+          PatchPoint NoEPEscape(test)
+          GuardBlockParamProxy l0
+          v15:BasicObject[BlockParamProxy] = Const Value(VALUE(0x1000))
+          SideExit UnhandledYARVInsn(splatkw)
         ");
     }
 
@@ -8274,7 +8340,6 @@ mod opt_tests {
           GuardBlockParamProxy l0
           v7:BasicObject[BlockParamProxy] = Const Value(VALUE(0x1000))
           v9:BasicObject = Send v0, 0x1008, :tap, v7
-          v10:BasicObject = GetLocal l0, EP@3
           CheckInterrupts
           Return v9
         ");

zjit/src/stats.rs

@@ -133,6 +133,7 @@ make_counters! {
     dynamic_send_count,
     dynamic_send_type_send_without_block,
     dynamic_send_type_send,
+    dynamic_send_type_send_forward,
     dynamic_send_type_invokeblock,
     dynamic_send_type_invokesuper,