RJIT: Implement ifunc invokeblock

5 files changed, 51 insertions, 4 deletions

diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
index 058195a62b..013f961829 100644
--- a/lib/ruby_vm/rjit/insn_compiler.rb
+++ b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -1536,8 +1536,47 @@ module RubyVM::RJIT
           send_shift: 0, frame_type: C::VM_FRAME_MAGIC_BLOCK,
         )
       elsif comptime_handler & 0x3 == 0x3 # VM_BH_IFUNC_P
-        asm.incr_counter(:invokeblock_ifunc)
-        CantCompile
+        # We aren't handling CALLER_SETUP_ARG and CALLER_REMOVE_EMPTY_KW_SPLAT yet.
+        if calling.flags & C::VM_CALL_ARGS_SPLAT != 0
+          asm.incr_counter(:invokeblock_ifunc_args_splat)
+          return CantCompile
+        end
+        if calling.flags & C::VM_CALL_KW_SPLAT != 0
+          asm.incr_counter(:invokeblock_ifunc_kw_splat)
+          return CantCompile
+        end
+
+        asm.comment('get local EP')
+        jit_get_lep(jit, asm, reg: :rax)
+        asm.mov(:rcx, [:rax, C.VALUE.size * C::VM_ENV_DATA_INDEX_SPECVAL]) # block_handler_opnd
+
+        asm.comment('guard block_handler type');
+        side_exit = side_exit(jit, ctx)
+        asm.mov(:rax, :rcx) # block_handler_opnd
+        asm.and(:rax, 0x3) # tag_opnd: block_handler is a tagged pointer
+        asm.cmp(:rax, 0x3) # VM_BH_IFUNC_P
+        tag_changed_exit = counted_exit(side_exit, :invokeblock_tag_changed)
+        jit_chain_guard(:jne, jit, ctx, asm, tag_changed_exit)
+
+        # The cfunc may not be leaf
+        jit_prepare_routine_call(jit, ctx, asm) # clobbers :rax
+
+        asm.comment('call ifunc')
+        asm.and(:rcx, ~0x3) # captured_opnd
+        asm.lea(:rax, ctx.sp_opnd(-calling.argc * C.VALUE.size)) # argv
+        asm.mov(C_ARGS[0], EC)
+        asm.mov(C_ARGS[1], :rcx) # captured_opnd
+        asm.mov(C_ARGS[2], calling.argc)
+        asm.mov(C_ARGS[3], :rax) # argv
+        asm.call(C.rb_vm_yield_with_cfunc)
+
+        ctx.stack_pop(calling.argc)
+        stack_ret = ctx.stack_push
+        asm.mov(stack_ret, C_RET)
+
+        # Share the successor with other chains
+        jump_to_next_insn(jit, ctx, asm)
+        EndBlock
       elsif symbol?(comptime_handler)
         asm.incr_counter(:invokeblock_symbol)
         CantCompile
diff --git a/rjit_c.c b/rjit_c.c
index a72ae4eec8..f4b34691ec 100644
--- a/rjit_c.c
+++ b/rjit_c.c
@@ -505,6 +505,7 @@ extern VALUE rb_reg_new_ary(VALUE ary, int opt);
 extern void rb_vm_setclassvariable(const rb_iseq_t *iseq, const rb_control_frame_t *cfp, ID id, VALUE val, ICVARC ic);
 extern VALUE rb_str_bytesize(VALUE str);
 extern const rb_callable_method_entry_t *rb_callable_method_entry_or_negative(VALUE klass, ID mid);
+extern VALUE rb_vm_yield_with_cfunc(rb_execution_context_t *ec, const struct rb_captured_block *captured, int argc, const VALUE *argv);
 
 #include "rjit_c.rbinc"
 
diff --git a/rjit_c.h b/rjit_c.h
index 2a60586621..9a7f2ba542 100644
--- a/rjit_c.h
+++ b/rjit_c.h
@@ -84,13 +84,14 @@ RJIT_RUNTIME_COUNTERS(
     invokesuper_same_me,
 
     invokeblock_none,
-    invokeblock_ifunc,
     invokeblock_symbol,
     invokeblock_proc,
     invokeblock_tag_changed,
     invokeblock_iseq_block_changed,
     invokeblock_iseq_arity,
     invokeblock_iseq_arg0_splat,
+    invokeblock_ifunc_args_splat,
+    invokeblock_ifunc_kw_splat,
 
     getivar_megamorphic,
     getivar_not_heap,
diff --git a/rjit_c.rb b/rjit_c.rb
index f6f8a5c898..c50911c55a 100644
--- a/rjit_c.rb
+++ b/rjit_c.rb
@@ -701,6 +701,10 @@ module RubyVM::RJIT # :nodoc: all
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_vm_throw) }
   end
 
+  def C.rb_vm_yield_with_cfunc
+    Primitive.cexpr! %q{ SIZET2NUM((size_t)rb_vm_yield_with_cfunc) }
+  end
+
   def C.rjit_full_cfunc_return
     Primitive.cexpr! %q{ SIZET2NUM((size_t)rjit_full_cfunc_return) }
   end
@@ -1316,13 +1320,14 @@ module RubyVM::RJIT # :nodoc: all
       invokesuper_me_changed: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokesuper_me_changed)")],
       invokesuper_same_me: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokesuper_same_me)")],
       invokeblock_none: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_none)")],
-      invokeblock_ifunc: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_ifunc)")],
       invokeblock_symbol: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_symbol)")],
       invokeblock_proc: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_proc)")],
       invokeblock_tag_changed: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_tag_changed)")],
       invokeblock_iseq_block_changed: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_iseq_block_changed)")],
       invokeblock_iseq_arity: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_iseq_arity)")],
       invokeblock_iseq_arg0_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_iseq_arg0_splat)")],
+      invokeblock_ifunc_args_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_ifunc_args_splat)")],
+      invokeblock_ifunc_kw_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), invokeblock_ifunc_kw_splat)")],
       getivar_megamorphic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), getivar_megamorphic)")],
       getivar_not_heap: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), getivar_not_heap)")],
       getivar_special_const: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), getivar_special_const)")],
diff --git a/tool/rjit/bindgen.rb b/tool/rjit/bindgen.rb
index 9916ebe0d0..cd8502a2a7 100755
--- a/tool/rjit/bindgen.rb
+++ b/tool/rjit/bindgen.rb
@@ -553,6 +553,7 @@ generator = BindingGenerator.new(
     rjit_str_simple_append
     rb_str_buf_append
     rb_str_dup
+    rb_vm_yield_with_cfunc
   ],
   types: %w[
     CALL_DATA