Partially implement send of cfunc

6 files changed, 175 insertions, 36 deletions

diff --git a/lib/ruby_vm/mjit/assembler.rb b/lib/ruby_vm/mjit/assembler.rb
index 40e8827046..f8543c4a66 100644
--- a/lib/ruby_vm/mjit/assembler.rb
+++ b/lib/ruby_vm/mjit/assembler.rb
@@ -3,6 +3,9 @@ module RubyVM::MJIT
   # 32-bit memory access
   class DwordPtr < Data.define(:reg, :disp); end
 
+  # C call argument registers
+  C_ARG_OPNDS = [:rdi, :rsi, :rdx, :rcx, :r8, :r9]
+
   # https://www.intel.com/content/dam/develop/public/us/en/documents/325383-sdm-vol-2abcd.pdf
   # Mostly an x86_64 assembler, but this also has some stuff that is useful for any architecture.
   class Assembler
@@ -118,11 +121,24 @@ module RubyVM::MJIT
       end
     end
 
-    # @param addr [Integer]
-    def call(addr)
+    def call(dst)
+      case dst
       # CALL rel32
-      # E8 cd
-      insn(opcode: 0xe8, imm: rel32(addr))
+      in Integer => dst_addr
+        # E8 cd
+        # D: Operand 1: Offset
+        insn(opcode: 0xe8, imm: rel32(dst_addr))
+      # CALL r/m64 (Mod 11: reg)
+      in Symbol => dst_reg
+        # FF /2
+        # M: Operand 1: ModRM:r/m (r)
+        insn(
+          opcode: 0xff,
+          mod_rm: ModRM[mod: Mod11, reg: 2, rm: dst_reg],
+        )
+      else
+        raise NotImplementedError, "call: not-implemented operands: #{dst.inspect}"
+      end
     end
 
     def cmovl(dst, src)
@@ -458,6 +474,16 @@ module RubyVM::MJIT
             mod_rm: ModRM[mod: Mod01, reg: src_reg, rm: dst_reg],
             disp: dst_disp,
           )
+        # MOV r/m64, r64 (Mod 10: [reg]+disp32)
+        in Symbol => src_reg if r64?(dst_reg) && imm32?(dst_disp) && r64?(src_reg)
+          # REX.W + 89 /r
+          # MR: Operand 1: ModRM:r/m (w), Operand 2: ModRM:reg (r)
+          insn(
+            prefix: REX_W,
+            opcode: 0x89,
+            mod_rm: ModRM[mod: Mod10, reg: src_reg, rm: dst_reg],
+            disp: imm32(dst_disp),
+          )
         else
           raise NotImplementedError, "mov: not-implemented operands: #{dst.inspect}, #{src.inspect}"
         end
diff --git a/lib/ruby_vm/mjit/context.rb b/lib/ruby_vm/mjit/context.rb
index fb9db212d5..e834b42999 100644
--- a/lib/ruby_vm/mjit/context.rb
+++ b/lib/ruby_vm/mjit/context.rb
@@ -23,7 +23,7 @@ module RubyVM::MJIT
       [SP, C.VALUE.size * (self.sp_offset - 1 - depth_from_top)]
     end
 
-    def sp_opnd(offset_bytes)
+    def sp_opnd(offset_bytes = 0)
       [SP, (C.VALUE.size * self.sp_offset) + offset_bytes]
     end
   end
diff --git a/lib/ruby_vm/mjit/insn_compiler.rb b/lib/ruby_vm/mjit/insn_compiler.rb
index fbf8d564b3..139d0851d0 100644
--- a/lib/ruby_vm/mjit/insn_compiler.rb
+++ b/lib/ruby_vm/mjit/insn_compiler.rb
@@ -370,7 +370,7 @@ module RubyVM::MJIT
         iseq: jit.iseq,
         shape: Default,
         target0: BranchTarget.new(ctx:, pc: jit.pc + C.VALUE.size * (jit.insn.len + jump_offset)), # branch target
-        target1: BranchTarget.new(ctx:, pc: jit.pc + C.VALUE.size * jit.insn.len),                    # fallthrough
+        target1: BranchTarget.new(ctx:, pc: jit.pc + C.VALUE.size * jit.insn.len),                 # fallthrough
       )
       branch_stub.target0.address = Assembler.new.then do |ocb_asm|
         @exit_compiler.compile_branch_stub(ctx, ocb_asm, branch_stub, true)
@@ -785,7 +785,7 @@ module RubyVM::MJIT
           @ocb.write(ocb_asm)
         end
         branch_stub.compile = proc do |branch_asm|
-          branch_asm.comment('jit_chain_guard')
+          # Not using `asm.comment` here since it's usually put before cmp/test before this.
           branch_asm.stub(branch_stub) do
             case branch_stub.shape
             in Default
@@ -865,9 +865,9 @@ module RubyVM::MJIT
 
     # @param jit [RubyVM::MJIT::JITState]
     # @param asm [RubyVM::MJIT::Assembler]
-    def jit_save_pc(jit, asm)
+    def jit_save_pc(jit, asm, comment: 'save PC to CFP')
       next_pc = jit.pc + jit.insn.len * C.VALUE.size # Use the next one for backtrace and side exits
-      asm.comment('save PC to CFP')
+      asm.comment(comment)
       asm.mov(:rax, next_pc)
       asm.mov([CFP, C.rb_control_frame_t.offsetof(:pc)], :rax)
     end
@@ -878,7 +878,7 @@ module RubyVM::MJIT
     def jit_save_sp(jit, ctx, asm)
       if ctx.sp_offset != 0
         asm.comment('save SP to CFP')
-        asm.lea(SP, ctx.sp_opnd(0))
+        asm.lea(SP, ctx.sp_opnd)
         asm.mov([CFP, C.rb_control_frame_t.offsetof(:sp)], SP)
         ctx.sp_offset = 0
       end
@@ -1012,7 +1012,7 @@ module RubyVM::MJIT
         asm.incr_counter(:send_kw_splat)
         return CantCompile
       end
-      recv_index = argc + ((flags & C.VM_CALL_ARGS_BLOCKARG == 0) ? 0 : 1)
+      recv_index = argc # TODO: +1 for VM_CALL_ARGS_BLOCKARG
 
       # Get a compile-time receiver and its class
       comptime_recv = jit.peek_at_stack(recv_index)
@@ -1066,8 +1066,7 @@ module RubyVM::MJIT
         jit_call_iseq_setup(jit, ctx, asm, ci, cme, flags, argc)
       # when C.VM_METHOD_TYPE_NOTIMPLEMENTED
       when C.VM_METHOD_TYPE_CFUNC
-        asm.incr_counter(:send_cfunc)
-        return CantCompile
+        jit_call_cfunc(jit, ctx, asm, ci, cme, flags, argc)
       when C.VM_METHOD_TYPE_ATTRSET
         asm.incr_counter(:send_attrset)
         return CantCompile
@@ -1123,14 +1122,11 @@ module RubyVM::MJIT
     def jit_call_iseq_setup_normal(jit, ctx, asm, ci, cme, flags, argc, iseq)
       # Save caller SP and PC before pushing a callee frame for backtrace and side exits
       asm.comment('save SP to caller CFP')
-      sp_index = ctx.sp_offset - 1 - argc - ((flags & C.VM_CALL_ARGS_BLOCKARG == 0) ? 0 : 1) # Pop receiver and arguments for side exits
-      asm.lea(:rax, [SP, C.VALUE.size * sp_index])
+      # Not setting this to SP register. This cfp->sp will be copied to SP on leave insn.
+      sp_index = -(1 + argc) # Pop receiver and arguments for side exits # TODO: subtract one more for VM_CALL_ARGS_BLOCKARG
+      asm.lea(:rax, ctx.sp_opnd(C.VALUE.size * sp_index))
       asm.mov([CFP, C.rb_control_frame_t.offsetof(:sp)], :rax)
-
-      asm.comment('save PC to caller CFP')
-      next_pc = jit.pc + jit.insn.len * C.VALUE.size # Use the next one for backtrace and side exits
-      asm.mov(:rax, next_pc)
-      asm.mov([CFP, C.rb_control_frame_t.offsetof(:pc)], :rax)
+      jit_save_pc(jit, asm, comment: 'save PC to caller CFP')
 
       frame_type = C.VM_FRAME_MAGIC_METHOD | C.VM_ENV_FLAG_LOCAL
       jit_push_frame(
@@ -1147,6 +1143,90 @@ module RubyVM::MJIT
       EndBlock
     end
 
+    # vm_call_cfunc
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def jit_call_cfunc(jit, ctx, asm, ci, cme, flags, argc)
+      if jit_caller_setup_arg(jit, ctx, asm, flags) == CantCompile
+        return CantCompile
+      end
+      if jit_caller_remove_empty_kw_splat(jit, ctx, asm, flags) == CantCompile
+        return CantCompile
+      end
+
+      # Disabled until we implement TracePoint invalidation
+      disabled = true
+      if disabled
+        return CantCompile
+      end
+
+      jit_call_cfunc_with_frame(jit, ctx, asm, ci, cme, flags, argc)
+    end
+
+    # jit_call_cfunc_with_frame
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def jit_call_cfunc_with_frame(jit, ctx, asm, ci, cme, flags, argc)
+      cfunc = cme.def.body.cfunc
+
+      # TODO: support them
+      if cfunc.argc < 0
+        asm.incr_counter(:send_cfunc_variadic)
+        return CantCompile
+      end
+      if argc + 1 > 6
+        asm.incr_counter(:send_cfunc_too_many_args)
+        return CantCompile
+      end
+
+      frame_type = C.VM_FRAME_MAGIC_CFUNC | C.VM_FRAME_FLAG_CFRAME | C.VM_ENV_FLAG_LOCAL
+      if flags & C.VM_CALL_KW_SPLAT != 0
+        frame_type |= C.VM_FRAME_FLAG_CFRAME_KW
+      end
+
+      # rb_check_arity
+      if argc != cfunc.argc
+        asm.incr_counter(:send_arity)
+        return CantCompile
+      end
+
+      # Save caller SP and PC before pushing a callee frame for backtrace and side exits
+      asm.comment('save SP to caller CFP')
+      sp_index = -(1 + argc) # Pop receiver and arguments for side exits # TODO: subtract one more for VM_CALL_ARGS_BLOCKARG
+      asm.lea(SP, ctx.sp_opnd(C.VALUE.size * sp_index))
+      asm.mov([CFP, C.rb_control_frame_t.offsetof(:sp)], SP)
+      ctx.sp_offset = -sp_index
+      jit_save_pc(jit, asm, comment: 'save PC to caller CFP')
+
+      jit_check_ints(jit, ctx, asm)
+
+      # Push a callee frame. SP register and ctx are not modified inside this.
+      jit_push_frame(jit, ctx, asm, ci, cme, flags, argc, frame_type)
+
+      asm.comment('call C function')
+      # Push receiver and args
+      (1 + argc).times do |i|
+        asm.mov(C_ARG_OPNDS[i], ctx.stack_opnd(argc - i)) # TODO: +1 for VM_CALL_ARGS_BLOCKARG
+      end
+      asm.mov(:rax, cfunc.func)
+      asm.call(:rax) # TODO: use rel32 if close enough
+      ctx.stack_pop(1 + argc)
+
+      asm.comment('push the return value')
+      stack_ret = ctx.stack_push
+      asm.mov(stack_ret, :rax)
+
+      asm.comment('pop the stack frame')
+      asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], CFP)
+
+      # Let guard chains share the same successor (ctx.sp_offset == 1)
+      assert_equal(1, ctx.sp_offset)
+      jump_to_next_insn(jit, ctx, asm)
+      EndBlock
+    end
+
     # vm_call_ivar
     # @param jit [RubyVM::MJIT::JITState]
     # @param ctx [RubyVM::MJIT::Context]
@@ -1205,30 +1285,28 @@ module RubyVM::MJIT
       asm.mov([SP, C.VALUE.size * (ep_offset - 1)], C.VM_BLOCK_HANDLER_NONE)
       asm.mov([SP, C.VALUE.size * (ep_offset - 0)], frame_type)
 
-      # This moves SP register. Don't side-exit after this.
-      asm.comment('move SP register to callee stack')
-      sp_offset = ctx.sp_offset + local_size + 3
-      asm.add(SP, C.VALUE.size * sp_offset)
-
       asm.comment('set up new frame')
       cfp_offset = -C.rb_control_frame_t.size # callee CFP
       # Not setting PC since JIT code will do that as needed
-      asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:sp)], SP)
       asm.mov(:rax, iseq.to_i)
       asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:iseq)], :rax)
-      self_index = -(1 + argc + ((flags & C.VM_CALL_ARGS_BLOCKARG == 0) ? 0 : 1) + local_size + 3)
+      self_index = ctx.sp_offset - (1 + argc) # TODO: +1 for VM_CALL_ARGS_BLOCKARG
       asm.mov(:rax, [SP, C.VALUE.size * self_index])
       asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:self)], :rax)
-      asm.lea(:rax, [SP, C.VALUE.size * -1])
+      asm.lea(:rax, [SP, C.VALUE.size * ep_offset])
       asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:ep)], :rax)
       asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:block_code)], 0)
-      asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:__bp__)], SP) # TODO: get rid of this!!
+      # Update SP register only for ISEQ calls. SP-relative operations should be done above this.
+      sp_reg = iseq ? SP : :rax
+      asm.lea(sp_reg, [SP, C.VALUE.size * (ctx.sp_offset + local_size + 3)])
+      asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:sp)], sp_reg)
+      asm.mov([CFP, cfp_offset + C.rb_control_frame_t.offsetof(:__bp__)], sp_reg) # TODO: get rid of this!!
 
       # cfp->jit_return is used only for ISEQs
       if iseq
         # Stub cfp->jit_return
         return_ctx = ctx.dup
-        return_ctx.stack_size -= argc + ((flags & C.VM_CALL_ARGS_BLOCKARG == 0) ? 0 : 1) # Pop args
+        return_ctx.stack_size -= argc # Pop args # TODO: subtract 1 more for VM_CALL_ARGS_BLOCKARG
         return_ctx.sp_offset = 1 # SP is in the position after popping a receiver and arguments
         branch_stub = BranchStub.new(
           iseq: jit.iseq,
@@ -1253,8 +1331,10 @@ module RubyVM::MJIT
       end
 
       asm.comment('switch to callee CFP')
-      asm.sub(CFP, C.rb_control_frame_t.size)
-      asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], CFP)
+      # Update CFP register only for ISEQ calls
+      cfp_reg = iseq ? CFP : :rax
+      asm.lea(cfp_reg, [CFP, cfp_offset])
+      asm.mov([EC, C.rb_execution_context_t.offsetof(:cfp)], cfp_reg)
     end
 
     # vm_callee_setup_arg: Set up args and return opt_pc (or CantCompile)
@@ -1279,11 +1359,13 @@ module RubyVM::MJIT
           return 0
         else
           # We don't support the remaining `else if`s yet.
+          asm.incr_counter(:send_iseq_not_simple)
           return CantCompile
         end
       end
 
       # We don't support setup_parameters_complex
+      asm.incr_counter(:send_iseq_kw_splat)
       return CantCompile
     end
 
diff --git a/mjit_c.h b/mjit_c.h
index 8392cf2ad5..8aa91250d8 100644
--- a/mjit_c.h
+++ b/mjit_c.h
@@ -128,6 +128,12 @@ MJIT_RUNTIME_COUNTERS(
     send_stackoverflow,
     send_arity,
 
+    send_iseq_not_simple,
+    send_iseq_kw_splat,
+
+    send_cfunc_variadic,
+    send_cfunc_too_many_args,
+
     send_ivar,
     send_ivar_splat,
     send_ivar_opt_send,
diff --git a/mjit_c.rb b/mjit_c.rb
index 283697ea6c..479c6d0abf 100644
--- a/mjit_c.rb
+++ b/mjit_c.rb
@@ -418,6 +418,18 @@ module RubyVM::MJIT # :nodoc: all
     Primitive.cexpr! %q{ UINT2NUM(VM_ENV_FLAG_LOCAL) }
   end
 
+  def C.VM_FRAME_FLAG_CFRAME
+    Primitive.cexpr! %q{ UINT2NUM(VM_FRAME_FLAG_CFRAME) }
+  end
+
+  def C.VM_FRAME_FLAG_CFRAME_KW
+    Primitive.cexpr! %q{ UINT2NUM(VM_FRAME_FLAG_CFRAME_KW) }
+  end
+
+  def C.VM_FRAME_MAGIC_CFUNC
+    Primitive.cexpr! %q{ UINT2NUM(VM_FRAME_MAGIC_CFUNC) }
+  end
+
   def C.VM_FRAME_MAGIC_METHOD
     Primitive.cexpr! %q{ UINT2NUM(VM_FRAME_MAGIC_METHOD) }
   end
@@ -879,6 +891,15 @@ module RubyVM::MJIT # :nodoc: all
     )
   end
 
+  def C.rb_method_cfunc_t
+    @rb_method_cfunc_t ||= CType::Struct.new(
+      "rb_method_cfunc_struct", Primitive.cexpr!("SIZEOF(struct rb_method_cfunc_struct)"),
+      func: [CType::Immediate.parse("void *"), Primitive.cexpr!("OFFSETOF((*((struct rb_method_cfunc_struct *)NULL)), func)")],
+      invoker: [CType::Immediate.parse("void *"), Primitive.cexpr!("OFFSETOF((*((struct rb_method_cfunc_struct *)NULL)), invoker)")],
+      argc: [CType::Immediate.parse("int"), Primitive.cexpr!("OFFSETOF((*((struct rb_method_cfunc_struct *)NULL)), argc)")],
+    )
+  end
+
   def C.rb_method_definition_struct
     @rb_method_definition_struct ||= CType::Struct.new(
       "rb_method_definition_struct", Primitive.cexpr!("SIZEOF(struct rb_method_definition_struct)"),
@@ -948,6 +969,10 @@ module RubyVM::MJIT # :nodoc: all
       send_refined: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_refined)")],
       send_stackoverflow: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_stackoverflow)")],
       send_arity: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_arity)")],
+      send_iseq_not_simple: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_iseq_not_simple)")],
+      send_iseq_kw_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_iseq_kw_splat)")],
+      send_cfunc_variadic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_cfunc_variadic)")],
+      send_cfunc_too_many_args: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_cfunc_too_many_args)")],
       send_ivar: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_ivar)")],
       send_ivar_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_ivar_splat)")],
       send_ivar_opt_send: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_ivar_opt_send)")],
@@ -1105,10 +1130,6 @@ module RubyVM::MJIT # :nodoc: all
     CType::Stub.new(:rb_event_flag_t)
   end
 
-  def C.rb_method_cfunc_t
-    CType::Stub.new(:rb_method_cfunc_t)
-  end
-
   def C.rb_method_alias_t
     CType::Stub.new(:rb_method_alias_t)
   end
diff --git a/tool/mjit/bindgen.rb b/tool/mjit/bindgen.rb
index 0d8cdee83f..df900803e4 100755
--- a/tool/mjit/bindgen.rb
+++ b/tool/mjit/bindgen.rb
@@ -379,6 +379,9 @@ generator = BindingGenerator.new(
       VM_CALL_OPT_SEND
       VM_ENV_FLAG_LOCAL
       VM_FRAME_MAGIC_METHOD
+      VM_FRAME_MAGIC_CFUNC
+      VM_FRAME_FLAG_CFRAME
+      VM_FRAME_FLAG_CFRAME_KW
       VM_METHOD_TYPE_CFUNC
       VM_METHOD_TYPE_ISEQ
       VM_METHOD_TYPE_IVAR
@@ -445,6 +448,7 @@ generator = BindingGenerator.new(
     rb_shape
     rb_shape_t
     rb_method_attr_t
+    rb_method_cfunc_t
   ],
   dynamic_types: %w[
     VALUE