ZJIT: Support variadic C calls

This reduces the `dynamic_send_count` in `liquid-render` by ~21%

Author: st0012

test/ruby/test_zjit.rb

@@ -2675,6 +2675,44 @@ def test
     }, insns: [:invokeblock]
   end
 
+  def test_ccall_variadic_with_multiple_args
+    assert_compiles "[1, 2, 3]", %q{
+      def test
+        a = []
+        a.push(1, 2, 3)
+        a
+      end
+
+      test
+      test
+    }, insns: [:opt_send_without_block]
+  end
+
+  def test_ccall_variadic_with_no_args
+    assert_compiles "[1]", %q{
+      def test
+        a = [1]
+        a.push
+      end
+
+      test
+      test
+    }, insns: [:opt_send_without_block]
+  end
+
+  def test_ccall_variadic_with_no_args_causing_argument_error
+    assert_compiles ":error", %q{
+      def test
+        format
+      rescue ArgumentError
+        :error
+      end
+
+      test
+      test
+    }, insns: [:opt_send_without_block]
+  end
+
   private
 
   # Assert that every method call in `test_script` can be compiled by ZJIT

zjit/src/codegen.rs

@@ -395,6 +395,9 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::GuardBitEquals { val, expected, state } => gen_guard_bit_equals(jit, asm, opnd!(val), *expected, &function.frame_state(*state)),
         Insn::PatchPoint { invariant, state } => no_output!(gen_patch_point(jit, asm, invariant, &function.frame_state(*state))),
         Insn::CCall { cfun, args, name: _, return_type: _, elidable: _ } => gen_ccall(asm, *cfun, opnds!(args)),
+        Insn::CCallVariadic { cfun, recv, args, name: _, cme, state } => {
+            gen_ccall_variadic(jit, asm, *cfun, opnd!(recv), opnds!(args), *cme, &function.frame_state(*state))
+        }
         Insn::GetIvar { self_val, id, state: _ } => gen_getivar(asm, opnd!(self_val), *id),
         Insn::SetGlobal { id, val, state } => no_output!(gen_setglobal(jit, asm, *id, opnd!(val), &function.frame_state(*state))),
         Insn::GetGlobal { id, state } => gen_getglobal(jit, asm, *id, &function.frame_state(*state)),
@@ -640,6 +643,79 @@ fn gen_ccall(asm: &mut Assembler, cfun: *const u8, args: Vec<Opnd>) -> lir::Opnd
     asm.ccall(cfun, args)
 }
 
+/// Push a C frame for cfunc calls
+fn gen_push_cfunc_frame(asm: &mut Assembler, recv: Opnd, cme: *const rb_callable_method_entry_t, sp_offset: i32) {
+    asm_comment!(asm, "push C frame");
+
+    let new_sp = asm.lea(Opnd::mem(64, SP, sp_offset * SIZEOF_VALUE_I32));
+    // sp[-3]: cme
+    asm.store(Opnd::mem(64, SP, (sp_offset - 3) * SIZEOF_VALUE_I32), VALUE::from(cme).into());
+    // sp[-2]: block_handler (specval)
+    asm.store(Opnd::mem(64, SP, (sp_offset - 2) * SIZEOF_VALUE_I32), VM_BLOCK_HANDLER_NONE.into());
+    // sp[-1]: frame type (must be a valid FIXNUM)
+    let frame_type = VM_FRAME_MAGIC_CFUNC | VM_FRAME_FLAG_CFRAME | VM_ENV_FLAG_LOCAL;
+    asm.store(Opnd::mem(64, SP, (sp_offset - 1) * SIZEOF_VALUE_I32), frame_type.into());
+
+    fn cfp_opnd(offset: i32) -> Opnd {
+        Opnd::mem(64, CFP, offset - (RUBY_SIZEOF_CONTROL_FRAME as i32))
+    }
+
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_PC), 0.into()); // C frames don't have a PC
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SP), new_sp);   // SP for new frame
+
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_ISEQ), 0.into()); // C frames don't have an iseq
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), recv);
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_BLOCK_CODE), 0.into());
+
+    let ep = asm.sub(new_sp, SIZEOF_VALUE.into());
+    asm.mov(cfp_opnd(RUBY_OFFSET_CFP_EP), ep);
+}
+
+/// Pop a C frame after cfunc call
+fn gen_pop_cfunc_frame(asm: &mut Assembler) {
+    asm_comment!(asm, "pop C frame");
+    let new_cfp = asm.add(CFP, RUBY_SIZEOF_CONTROL_FRAME.into());
+    asm.mov(CFP, new_cfp);
+    asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), CFP);
+}
+
+/// Generate code for a variadic C function call
+/// func(int argc, VALUE *argv, VALUE recv)
+fn gen_ccall_variadic(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    cfun: *const u8,
+    recv: Opnd,
+    args: Vec<Opnd>,
+    cme: *const rb_callable_method_entry_t,
+    state: &FrameState,
+) -> lir::Opnd {
+    gen_prepare_non_leaf_call(jit, asm, state);
+
+    asm_comment!(asm, "pushing C frame for cme=0x{:x}", cme as usize);
+    // Following YJIT: allocate 3 slots for frame metadata (cme, block_handler, frame_type)
+    // The new frame's SP will be at current_stack_size + 3
+    let sp_offset = state.stack().len() as i32 + 3;
+    gen_push_cfunc_frame(asm, recv, cme, sp_offset);
+
+    // Move to the new frame
+    let new_cfp = asm.lea(Opnd::mem(64, CFP, -(RUBY_SIZEOF_CONTROL_FRAME as i32)));
+    asm.mov(CFP, new_cfp);
+    asm.store(Opnd::mem(64, EC, RUBY_OFFSET_EC_CFP), new_cfp);
+
+    let argc = args.len();
+
+    asm_comment!(asm, "call variadic cfunc: argc={}", argc);
+
+    let argv_ptr = gen_push_opnds(jit, asm, &args);
+    let result = asm.ccall(cfun, vec![argc.into(), argv_ptr, recv]);
+    gen_pop_opnds(asm, &args);
+
+    gen_pop_cfunc_frame(asm);
+
+    result
+}
+
 /// Emit an uncached instance variable lookup
 fn gen_getivar(asm: &mut Assembler, recv: Opnd, id: ID) -> Opnd {
     asm_ccall!(asm, rb_ivar_get, recv, id.0.into())

zjit/src/hir.rs

@@ -581,6 +581,17 @@ pub enum Insn {
     /// `name` is for printing purposes only
     CCall { cfun: *const u8, args: Vec<InsnId>, name: ID, return_type: Type, elidable: bool },
 
+    /// Call a variadic C function with signature: func(int argc, VALUE *argv, VALUE recv)
+    /// This handles frame setup, argv creation, and frame teardown all in one
+    CCallVariadic {
+        cfun: *const u8,
+        recv: InsnId,
+        args: Vec<InsnId>,
+        cme: *const rb_callable_method_entry_t,
+        name: ID,
+        state: InsnId,
+    },
+
     /// Un-optimized fallback implementation (dynamic dispatch) for send-ish instructions
     /// Ignoring keyword arguments etc for now
     SendWithoutBlock { recv: InsnId, cd: *const rb_call_data, args: Vec<InsnId>, state: InsnId },
@@ -713,6 +724,7 @@ impl Insn {
             Insn::LoadIvarEmbedded { .. } => false,
             Insn::LoadIvarExtended { .. } => false,
             Insn::CCall { elidable, .. } => !elidable,
+            Insn::CCallVariadic { .. } => true,
             // TODO: NewRange is effects free if we can prove the two ends to be Fixnum,
             // but we don't have type information here in `impl Insn`. See rb_range_new().
             Insn::NewRange { .. } => true,
@@ -888,6 +900,13 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
                 }
                 Ok(())
             },
+            Insn::CCallVariadic { cfun,  recv, args, name, .. } => {
+                write!(f, "CCallVariadic {}@{:p}, {recv}", name.contents_lossy(), self.ptr_map.map_ptr(cfun))?;
+                for arg in args {
+                    write!(f, ", {arg}")?;
+                }
+                Ok(())
+            },
             Insn::Snapshot { state } => write!(f, "Snapshot {}", state.print(self.ptr_map)),
             Insn::Defined { op_type, v, .. } => {
                 // op_type (enum defined_type) printing logic from iseq.c.
@@ -1368,6 +1387,9 @@ impl Function {
             &HashDup { val, state } => HashDup { val: find!(val), state },
             &ObjectAlloc { val, state } => ObjectAlloc { val: find!(val), state },
             &CCall { cfun, ref args, name, return_type, elidable } => CCall { cfun, args: find_vec!(args), name, return_type, elidable },
+            &CCallVariadic { cfun, recv, ref args, cme, name, state } => CCallVariadic {
+                cfun, recv: find!(recv), args: find_vec!(args), cme, name, state
+            },
             &Defined { op_type, obj, pushval, v, state } => Defined { op_type, obj, pushval, v: find!(v), state: find!(state) },
             &DefinedIvar { self_val, pushval, id, state } => DefinedIvar { self_val: find!(self_val), pushval, id, state },
             &NewArray { ref elements, state } => NewArray { elements: find_vec!(elements), state: find!(state) },
@@ -1455,6 +1477,7 @@ impl Function {
             Insn::NewRangeFixnum { .. } => types::RangeExact,
             Insn::ObjectAlloc { .. } => types::HeapObject,
             Insn::CCall { return_type, .. } => *return_type,
+            Insn::CCallVariadic { .. } => types::BasicObject,
             Insn::GuardType { val, guard_type, .. } => self.type_of(*val).intersection(*guard_type),
             Insn::GuardTypeNot { .. } => types::BasicObject,
             Insn::GuardBitEquals { val, expected, .. } => self.type_of(*val).intersection(Type::from_value(*expected)),
@@ -1982,9 +2005,42 @@ impl Function {
                         return Ok(());
                     }
                 }
+                // Variadic method
                 -1 => {
-                    // (argc, argv, self) parameter form
-                    // Falling through for now
+                    // The method gets a pointer to the first argument
+                    // rb_f_puts(int argc, VALUE *argv, VALUE recv)
+                    let ci_flags = unsafe { vm_ci_flag(call_info) };
+                    if ci_flags & VM_CALL_ARGS_SIMPLE != 0 {
+                        // We need profiled type information for variadic calls
+                        // to ensure we're calling the right method
+                        if profiled_type.is_none() {
+                            return Err(());
+                        }
+
+                        // Add method redefinition guard
+                        fun.push_insn(block, Insn::PatchPoint {
+                            invariant: Invariant::MethodRedefined {
+                                klass: recv_class,
+                                method: method_id,
+                                cme: method
+                            },
+                            state
+                        });
+
+                        let cfun = unsafe { get_mct_func(cfunc) }.cast();
+                        let ccall = fun.push_insn(block, Insn::CCallVariadic {
+                            cfun,
+                            recv,
+                            args,
+                            cme: method,
+                            name: method_id,
+                            state,
+                        });
+
+                        fun.make_equal_to(send_insn_id, ccall);
+                        return Ok(());
+                    }
+                    // Fall through for complex cases (splat, kwargs, etc.)
                 }
                 -2 => {
                     // (self, args_ruby_array) parameter form
@@ -2297,6 +2353,11 @@ impl Function {
                 worklist.push_back(state)
             }
             Insn::CCall { args, .. } => worklist.extend(args),
+            Insn::CCallVariadic { recv, args, state, .. } => {
+                worklist.push_back(*recv);
+                worklist.extend(args);
+                worklist.push_back(*state);
+            }
             &Insn::GetIvar { self_val, state, .. } | &Insn::DefinedIvar { self_val, state, .. } => {
                 worklist.push_back(self_val);
                 worklist.push_back(state);
@@ -6797,6 +6858,26 @@ mod opt_tests {
         ");
     }
 
+    #[test]
+    fn test_optimize_variadic_ccall() {
+        eval("
+            def test
+              puts 'Hello'
+            end
+            test; test
+        ");
+        assert_snapshot!(hir_string("test"), @r"
+        fn test@<compiled>:3:
+        bb0(v0:BasicObject):
+          v4:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000))
+          v6:StringExact = StringCopy v4
+          PatchPoint MethodRedefined(Object@0x1008, puts@0x1010, cme:0x1018)
+          v15:BasicObject = CCallVariadic puts@0x1040, v0, v6
+          CheckInterrupts
+          Return v15
+        ");
+    }
+
     #[test]
     fn test_dont_optimize_fixnum_add_if_redefined() {
         eval("