YJIT: Support invokeblock (#6640)

* YJIT: Support invokeblock

* Update yjit/src/backend/arm64/mod.rs

* Update yjit/src/codegen.rs

Co-authored-by: Maxime Chevalier-Boisvert <[email protected]>

10 files changed, 198 insertions, 30 deletions

diff --git a/yjit.c b/yjit.c
index 1694e1edd7..b943277d61 100644
--- a/yjit.c
+++ b/yjit.c
@@ -627,6 +627,12 @@ rb_get_iseq_body_stack_max(const rb_iseq_t *iseq)
 }
 
 bool
+rb_get_iseq_flags_has_lead(const rb_iseq_t *iseq)
+{
+    return iseq->body->param.flags.has_lead;
+}
+
+bool
 rb_get_iseq_flags_has_opt(const rb_iseq_t *iseq)
 {
     return iseq->body->param.flags.has_opt;
@@ -669,7 +675,13 @@ rb_get_iseq_flags_has_block(const rb_iseq_t *iseq)
 }
 
 bool
-rb_get_iseq_flags_has_accepts_no_kwarg(const rb_iseq_t *iseq)
+rb_get_iseq_flags_ambiguous_param0(const rb_iseq_t *iseq)
+{
+    return iseq->body->param.flags.ambiguous_param0;
+}
+
+bool
+rb_get_iseq_flags_accepts_no_kwarg(const rb_iseq_t *iseq)
 {
     return iseq->body->param.flags.accepts_no_kwarg;
 }
diff --git a/yjit.rb b/yjit.rb
index ac49a30e90..bb344948eb 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -187,6 +187,7 @@ module RubyVM::YJIT
       $stderr.puts("***YJIT: Printing YJIT statistics on exit***")
 
       print_counters(stats, prefix: 'send_', prompt: 'method call exit reasons: ')
+      print_counters(stats, prefix: 'invokeblock_', prompt: 'invokeblock exit reasons: ')
       print_counters(stats, prefix: 'invokesuper_', prompt: 'invokesuper exit reasons: ')
       print_counters(stats, prefix: 'leave_', prompt: 'leave exit reasons: ')
       print_counters(stats, prefix: 'gbpp_', prompt: 'getblockparamproxy exit reasons: ')
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 21aaec84cb..167ab2a74f 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -250,6 +250,7 @@ fn main() {
         .blocklist_type("rb_control_frame_struct")
         .opaque_type("rb_control_frame_struct")
         .allowlist_function("rb_vm_bh_to_procval")
+        .allowlist_function("rb_vm_ep_local_ep")
         .allowlist_type("vm_special_object_type")
         .allowlist_var("VM_ENV_DATA_INDEX_SPECVAL")
         .allowlist_var("VM_ENV_DATA_INDEX_FLAGS")
@@ -353,13 +354,15 @@ fn main() {
         .allowlist_function("rb_get_iseq_body_parent_iseq")
         .allowlist_function("rb_get_iseq_body_iseq_encoded")
         .allowlist_function("rb_get_iseq_body_stack_max")
+        .allowlist_function("rb_get_iseq_flags_has_lead")
         .allowlist_function("rb_get_iseq_flags_has_opt")
         .allowlist_function("rb_get_iseq_flags_has_kw")
         .allowlist_function("rb_get_iseq_flags_has_rest")
         .allowlist_function("rb_get_iseq_flags_has_post")
         .allowlist_function("rb_get_iseq_flags_has_kwrest")
         .allowlist_function("rb_get_iseq_flags_has_block")
-        .allowlist_function("rb_get_iseq_flags_has_accepts_no_kwarg")
+        .allowlist_function("rb_get_iseq_flags_ambiguous_param0")
+        .allowlist_function("rb_get_iseq_flags_accepts_no_kwarg")
         .allowlist_function("rb_get_iseq_flags_ruby2_keywords")
         .allowlist_function("rb_get_iseq_body_local_table_size")
         .allowlist_function("rb_get_iseq_body_param_keyword")
diff --git a/yjit/src/backend/arm64/mod.rs b/yjit/src/backend/arm64/mod.rs
index 0c784c0bea..ce1dd2e43c 100644
--- a/yjit/src/backend/arm64/mod.rs
+++ b/yjit/src/backend/arm64/mod.rs
@@ -165,8 +165,8 @@ impl Assembler
                 Opnd::Reg(_) | Opnd::InsnOut { .. } => opnd,
                 Opnd::Mem(_) => split_load_operand(asm, opnd),
                 Opnd::Imm(imm) => {
-                    if imm <= 0 {
-                        asm.load(opnd)
+                    if imm == 0 {
+                        Opnd::Reg(XZR_REG)
                     } else if (dest_num_bits == 64 &&
                                 BitmaskImmediate::try_from(imm as u64).is_ok()) ||
                             (dest_num_bits == 32 &&
@@ -1352,8 +1352,8 @@ mod tests {
         asm.test(Opnd::Reg(X0_REG), Opnd::Imm(-7));
         asm.compile_with_num_regs(&mut cb, 1);
 
-        // Assert that a load and a test instruction were written.
-        assert_eq!(8, cb.get_write_pos());
+        // Assert that a test instruction is written.
+        assert_eq!(4, cb.get_write_pos());
     }
 
     #[test]
diff --git a/yjit/src/backend/x86_64/mod.rs b/yjit/src/backend/x86_64/mod.rs
index ac5ac0fff4..dc5f21221d 100644
--- a/yjit/src/backend/x86_64/mod.rs
+++ b/yjit/src/backend/x86_64/mod.rs
@@ -93,6 +93,7 @@ impl Assembler
         vec![
             RAX_REG,
             RCX_REG,
+            RDX_REG,
         ]
     }
 
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index de0a2dec0f..9ec6c26f89 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -4011,6 +4011,7 @@ enum SpecVal {
     BlockISeq(IseqPtr),
     BlockParamProxy,
     PrevEP(*const VALUE),
+    PrevEPOpnd(Opnd),
 }
 
 struct ControlFrame {
@@ -4050,17 +4051,6 @@ fn gen_push_frame(
 
     let sp = frame.sp;
 
-    let num_locals = frame.local_size;
-    if num_locals > 0 {
-        asm.comment("initialize locals");
-
-        // Initialize local variables to Qnil
-        for i in 0..num_locals {
-            let offs = (SIZEOF_VALUE as i32) * (i - num_locals - 3);
-            asm.store(Opnd::mem(64, sp, offs), Qnil.into());
-        }
-    }
-
     asm.comment("push cme, specval, frame type");
 
     // Write method entry at sp[-3]
@@ -4099,9 +4089,25 @@ fn gen_push_frame(
             let tagged_prev_ep = (prev_ep as usize) | 1;
             VALUE(tagged_prev_ep).into()
         }
+        SpecVal::PrevEPOpnd(ep_opnd) => {
+            asm.or(ep_opnd, 1.into())
+        },
     };
     asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -2), specval);
 
+    // Arm requires another register to load the immediate value of Qnil before storing it.
+    // So doing this after releasing the register for specval to avoid register spill.
+    let num_locals = frame.local_size;
+    if num_locals > 0 {
+        asm.comment("initialize locals");
+
+        // Initialize local variables to Qnil
+        for i in 0..num_locals {
+            let offs = (SIZEOF_VALUE as i32) * (i - num_locals - 3);
+            asm.store(Opnd::mem(64, sp, offs), Qnil.into());
+        }
+    }
+
     // Write env flags at sp[-1]
     // sp[-1] = frame_type;
     asm.store(Opnd::mem(64, sp, SIZEOF_VALUE_I32 * -1), frame.frame_type.into());
@@ -4522,7 +4528,7 @@ fn gen_send_bmethod(
     }
 
     let frame_type = VM_FRAME_MAGIC_BLOCK | VM_FRAME_FLAG_BMETHOD | VM_FRAME_FLAG_LAMBDA;
-    gen_send_iseq(jit, ctx, asm, ocb, iseq, ci, frame_type, Some(capture.ep), cme, block, flags, argc)
+    gen_send_iseq(jit, ctx, asm, ocb, iseq, ci, frame_type, Some(capture.ep), cme, block, flags, argc, None)
 }
 
 fn gen_send_iseq(
@@ -4538,10 +4544,10 @@ fn gen_send_iseq(
     block: Option<IseqPtr>,
     flags: u32,
     argc: i32,
+    captured_opnd: Option<Opnd>,
 ) -> CodegenStatus {
     let mut argc = argc;
 
-
     // Create a side-exit to fall back to the interpreter
     let side_exit = get_side_exit(jit, ocb, ctx);
 
@@ -4592,7 +4598,7 @@ fn gen_send_iseq(
 
     // If we have a method accepting no kwargs (**nil), exit if we have passed
     // it any kwargs.
-    if supplying_kws && unsafe { get_iseq_flags_has_accepts_no_kwarg(iseq) } {
+    if supplying_kws && unsafe { get_iseq_flags_accepts_no_kwarg(iseq) } {
         gen_counter_incr!(asm, send_iseq_complex_callee);
         return CantCompile;
     }
@@ -4997,12 +5003,17 @@ fn gen_send_iseq(
         asm.mov(ctx.stack_opnd(-1), unspec_opnd.into());
     }
 
-    // Points to the receiver operand on the stack
-    let recv = ctx.stack_opnd(argc);
+    // Points to the receiver operand on the stack unless a captured environment is used
+    let recv = match captured_opnd {
+        Some(captured_opnd) => asm.load(Opnd::mem(64, captured_opnd, 0)), // captured->self
+        _ => ctx.stack_opnd(argc),
+    };
+    let captured_self = captured_opnd.is_some();
+    let sp_offset = (argc as isize) + if captured_self { 0 } else { 1 };
 
     // Store the updated SP on the current frame (pop arguments and receiver)
     asm.comment("store caller sp");
-    let caller_sp = asm.lea(ctx.sp_opnd((SIZEOF_VALUE as isize) * -((argc as isize) + 1)));
+    let caller_sp = asm.lea(ctx.sp_opnd((SIZEOF_VALUE as isize) * -sp_offset));
     asm.store(Opnd::mem(64, CFP, RUBY_OFFSET_CFP_SP), caller_sp);
 
     // Store the next PC in the current frame
@@ -5017,6 +5028,9 @@ fn gen_send_iseq(
         // We've already side-exited if the callee expects a block, so we
         // ignore any supplied block here
         SpecVal::PrevEP(prev_ep)
+    } else if let Some(captured_opnd) = captured_opnd {
+        let ep_opnd = asm.load(Opnd::mem(64, captured_opnd, SIZEOF_VALUE_I32)); // captured->ep
+        SpecVal::PrevEPOpnd(ep_opnd)
     } else if block_arg_type == Some(Type::BlockParamProxy) {
         SpecVal::BlockParamProxy
     } else if let Some(block_val) = block {
@@ -5058,7 +5072,11 @@ fn gen_send_iseq(
         callee_ctx.set_local_type(arg_idx.try_into().unwrap(), arg_type);
     }
 
-    let recv_type = ctx.get_opnd_type(StackOpnd(argc.try_into().unwrap()));
+    let recv_type = if captured_self {
+        ctx.get_opnd_type(CapturedSelfOpnd)
+    } else {
+        ctx.get_opnd_type(StackOpnd(argc.try_into().unwrap()))
+    };
     callee_ctx.upgrade_opnd_type(SelfOpnd, recv_type);
 
     // The callee might change locals through Kernel#binding and other means.
@@ -5068,7 +5086,7 @@ fn gen_send_iseq(
     // After the return, sp_offset will be 1. The codegen for leave writes
     // the return value in case of JIT-to-JIT return.
     let mut return_ctx = *ctx;
-    return_ctx.stack_pop((argc + 1).try_into().unwrap());
+    return_ctx.stack_pop(sp_offset.try_into().unwrap());
     return_ctx.stack_push(Type::Unknown);
     return_ctx.set_sp_offset(1);
     return_ctx.reset_chain_depth();
@@ -5317,7 +5335,7 @@ fn gen_send_general(
             VM_METHOD_TYPE_ISEQ => {
                 let iseq = unsafe { get_def_iseq_ptr((*cme).def) };
                 let frame_type = VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL;
-                return gen_send_iseq(jit, ctx, asm, ocb, iseq, ci, frame_type, None, cme, block, flags, argc);
+                return gen_send_iseq(jit, ctx, asm, ocb, iseq, ci, frame_type, None, cme, block, flags, argc, None);
             }
             VM_METHOD_TYPE_CFUNC => {
                 return gen_send_cfunc(
@@ -5642,6 +5660,95 @@ fn gen_send(
     return gen_send_general(jit, ctx, asm, ocb, cd, block);
 }
 
+fn gen_invokeblock(
+    jit: &mut JITState,
+    ctx: &mut Context,
+    asm: &mut Assembler,
+    ocb: &mut OutlinedCb,
+) -> CodegenStatus {
+    if !jit_at_current_insn(jit) {
+        defer_compilation(jit, ctx, asm, ocb);
+        return EndBlock;
+    }
+
+    // Get call info
+    let cd = jit_get_arg(jit, 0).as_ptr();
+    let ci = unsafe { get_call_data_ci(cd) };
+    let argc: i32 = unsafe { vm_ci_argc(ci) }.try_into().unwrap();
+    let flags = unsafe { vm_ci_flag(ci) };
+
+    // Get block_handler
+    let cfp = unsafe { get_ec_cfp(jit.ec.unwrap()) };
+    let lep = unsafe { rb_vm_ep_local_ep(get_cfp_ep(cfp)) };
+    let comptime_handler = unsafe { *lep.offset(VM_ENV_DATA_INDEX_SPECVAL.try_into().unwrap()) };
+
+    // Handle each block_handler type
+    if comptime_handler.0 == VM_BLOCK_HANDLER_NONE as usize { // no block given
+        gen_counter_incr!(asm, invokeblock_none);
+        CantCompile
+    } else if comptime_handler.0 & 0x3 == 0x1 { // VM_BH_ISEQ_BLOCK_P
+        asm.comment("get local EP");
+        let ep_opnd = gen_get_lep(jit, asm);
+        let block_handler_opnd = asm.load(
+            Opnd::mem(64, ep_opnd, (SIZEOF_VALUE as i32) * (VM_ENV_DATA_INDEX_SPECVAL as i32))
+        );
+
+        asm.comment("guard block_handler type");
+        let side_exit = get_side_exit(jit, ocb, ctx);
+        let tag_opnd = asm.and(block_handler_opnd, 0x3.into()); // block_handler is a tagged pointer
+        asm.cmp(tag_opnd, 0x1.into()); // VM_BH_ISEQ_BLOCK_P
+        asm.jne(counted_exit!(ocb, side_exit, invokeblock_iseq_tag_changed).into());
+
+        // Not supporting vm_callee_setup_block_arg_arg0_splat for now
+        let comptime_captured = unsafe { ((comptime_handler.0 & !0x3) as *const rb_captured_block).as_ref().unwrap() };
+        let comptime_iseq = unsafe { *comptime_captured.code.iseq.as_ref() };
+        if argc == 1 && unsafe { get_iseq_flags_has_lead(comptime_iseq) && !get_iseq_flags_ambiguous_param0(comptime_iseq) } {
+            gen_counter_incr!(asm, invokeblock_iseq_arg0_splat);
+            return CantCompile;
+        }
+
+        asm.comment("guard known ISEQ");
+        let captured_opnd = asm.and(block_handler_opnd, Opnd::Imm(!0x3));
+        let iseq_opnd = asm.load(Opnd::mem(64, captured_opnd, SIZEOF_VALUE_I32 * 2));
+        asm.cmp(iseq_opnd, (comptime_iseq as usize).into());
+        let block_changed_exit = counted_exit!(ocb, side_exit, invokeblock_iseq_block_changed);
+        jit_chain_guard(
+            JCC_JNE,
+            jit,
+            ctx,
+            asm,
+            ocb,
+            SEND_MAX_CHAIN_DEPTH,
+            block_changed_exit,
+        );
+
+        gen_send_iseq(
+            jit,
+            ctx,
+            asm,
+            ocb,
+            comptime_iseq,
+            ci,
+            VM_FRAME_MAGIC_BLOCK,
+            None,
+            0 as _,
+            None,
+            flags,
+            argc,
+            Some(captured_opnd),
+        )
+    } else if comptime_handler.0 & 0x3 == 0x3 { // VM_BH_IFUNC_P
+        gen_counter_incr!(asm, invokeblock_ifunc);
+        CantCompile
+    } else if comptime_handler.symbol_p() {
+        gen_counter_incr!(asm, invokeblock_symbol);
+        CantCompile
+    } else { // Proc
+        gen_counter_incr!(asm, invokeblock_proc);
+        CantCompile
+    }
+}
+
 fn gen_invokesuper(
     jit: &mut JITState,
     ctx: &mut Context,
@@ -5781,7 +5888,7 @@ fn gen_invokesuper(
         VM_METHOD_TYPE_ISEQ => {
             let iseq = unsafe { get_def_iseq_ptr((*cme).def) };
             let frame_type = VM_FRAME_MAGIC_METHOD | VM_ENV_FLAG_LOCAL;
-            gen_send_iseq(jit, ctx, asm, ocb, iseq, ci, frame_type, None, cme, block, ci_flags, argc)
+            gen_send_iseq(jit, ctx, asm, ocb, iseq, ci, frame_type, None, cme, block, ci_flags, argc, None)
         }
         VM_METHOD_TYPE_CFUNC => {
             gen_send_cfunc(jit, ctx, asm, ocb, ci, cme, block, ptr::null(), ci_flags, argc)
@@ -6548,6 +6655,7 @@ fn get_gen_fn(opcode: VALUE) -> Option<InsnGenFn> {
         YARVINSN_getblockparam => Some(gen_getblockparam),
         YARVINSN_opt_send_without_block => Some(gen_opt_send_without_block),
         YARVINSN_send => Some(gen_send),
+        YARVINSN_invokeblock => Some(gen_invokeblock),
         YARVINSN_invokesuper => Some(gen_invokesuper),
         YARVINSN_leave => Some(gen_leave),
 
diff --git a/yjit/src/core.rs b/yjit/src/core.rs
index c0e48e87b2..afa7604aaf 100644
--- a/yjit/src/core.rs
+++ b/yjit/src/core.rs
@@ -269,6 +269,9 @@ pub enum InsnOpnd {
     // The value is self
     SelfOpnd,
 
+    // Captured block's self
+    CapturedSelfOpnd,
+
     // Temporary stack operand with stack index
     StackOpnd(u16),
 }
@@ -297,6 +300,9 @@ pub struct Context {
     // Type we track for self
     self_type: Type,
 
+    // Type we track for captured block's self
+    captured_self_type: Type,
+
     // Mapping of temp stack entries to types we track
     temp_mapping: [TempMapping; MAX_TEMP_TYPES],
 }
@@ -1160,6 +1166,7 @@ impl Context {
     pub fn get_opnd_type(&self, opnd: InsnOpnd) -> Type {
         match opnd {
             SelfOpnd => self.self_type,
+            CapturedSelfOpnd => self.captured_self_type,
             StackOpnd(idx) => {
                 let idx = idx as u16;
                 assert!(idx < self.stack_size);
@@ -1201,6 +1208,7 @@ impl Context {
 
         match opnd {
             SelfOpnd => self.self_type.upgrade(opnd_type),
+            CapturedSelfOpnd => self.self_type.upgrade(opnd_type),
             StackOpnd(idx) => {
                 let idx = idx as u16;
                 assert!(idx < self.stack_size);
@@ -1236,6 +1244,7 @@ impl Context {
 
         match opnd {
             SelfOpnd => (MapToSelf, opnd_type),
+            CapturedSelfOpnd => unreachable!("not used for captured self"),
             StackOpnd(idx) => {
                 let idx = idx as u16;
                 assert!(idx < self.stack_size);
@@ -1257,6 +1266,7 @@ impl Context {
     pub fn set_opnd_mapping(&mut self, opnd: InsnOpnd, (mapping, opnd_type): (TempMapping, Type)) {
         match opnd {
             SelfOpnd => unreachable!("self always maps to self"),
+            CapturedSelfOpnd => unreachable!("not used for captured self"),
             StackOpnd(idx) => {
                 assert!(idx < self.stack_size);
                 let stack_idx = (self.stack_size - 1 - idx) as usize;
diff --git a/yjit/src/cruby.rs b/yjit/src/cruby.rs
index d8dbdb4019..168443e6f0 100644
--- a/yjit/src/cruby.rs
+++ b/yjit/src/cruby.rs
@@ -162,6 +162,7 @@ pub use rb_iseq_encoded_size as get_iseq_encoded_size;
 pub use rb_get_iseq_body_local_iseq as get_iseq_body_local_iseq;
 pub use rb_get_iseq_body_iseq_encoded as get_iseq_body_iseq_encoded;
 pub use rb_get_iseq_body_stack_max as get_iseq_body_stack_max;
+pub use rb_get_iseq_flags_has_lead as get_iseq_flags_has_lead;
 pub use rb_get_iseq_flags_has_opt as get_iseq_flags_has_opt;
 pub use rb_get_iseq_flags_has_kw as get_iseq_flags_has_kw;
 pub use rb_get_iseq_flags_has_rest as get_iseq_flags_has_rest;
@@ -169,7 +170,8 @@ pub use rb_get_iseq_flags_ruby2_keywords as get_iseq_flags_ruby2_keywords;
 pub use rb_get_iseq_flags_has_post as get_iseq_flags_has_post;
 pub use rb_get_iseq_flags_has_kwrest as get_iseq_flags_has_kwrest;
 pub use rb_get_iseq_flags_has_block as get_iseq_flags_has_block;
-pub use rb_get_iseq_flags_has_accepts_no_kwarg as get_iseq_flags_has_accepts_no_kwarg;
+pub use rb_get_iseq_flags_ambiguous_param0 as get_iseq_flags_ambiguous_param0;
+pub use rb_get_iseq_flags_accepts_no_kwarg as get_iseq_flags_accepts_no_kwarg;
 pub use rb_get_iseq_body_local_table_size as get_iseq_body_local_table_size;
 pub use rb_get_iseq_body_param_keyword as get_iseq_body_param_keyword;
 pub use rb_get_iseq_body_param_size as get_iseq_body_param_size;
@@ -346,13 +348,27 @@ impl VALUE {
         (cval & mask) == flag
     }
 
-    /// Return true for a static (non-heap) Ruby symbol
+    /// Return true if the value is a Ruby symbol (RB_SYMBOL_P)
+    pub fn symbol_p(self) -> bool {
+        self.static_sym_p() || self.dynamic_sym_p()
+    }
+
+    /// Return true for a static (non-heap) Ruby symbol (RB_STATIC_SYM_P)
     pub fn static_sym_p(self) -> bool {
         let VALUE(cval) = self;
         let flag = RUBY_SYMBOL_FLAG as usize;
         (cval & 0xff) == flag
     }
 
+    /// Return true for a dynamic Ruby symbol (RB_DYNAMIC_SYM_P)
+    fn dynamic_sym_p(self) -> bool {
+        return if self.special_const_p() {
+            false
+        } else {
+            self.builtin_type() == RUBY_T_SYMBOL
+        }
+    }
+
     /// Returns true or false depending on whether the value is nil
     pub fn nil_p(self) -> bool {
         self == Qnil
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index da7b332e07..3373c6d76e 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -966,6 +966,9 @@ pub const VM_ENV_FLAG_WB_REQUIRED: vm_frame_env_flags = 8;
 pub const VM_ENV_FLAG_ISOLATED: vm_frame_env_flags = 16;
 pub type vm_frame_env_flags = u32;
 extern "C" {
+    pub fn rb_vm_ep_local_ep(ep: *const VALUE) -> *const VALUE;
+}
+extern "C" {
     pub fn rb_iseq_path(iseq: *const rb_iseq_t) -> VALUE;
 }
 extern "C" {
@@ -1431,6 +1434,9 @@ extern "C" {
     pub fn rb_get_iseq_body_stack_max(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
 }
 extern "C" {
+    pub fn rb_get_iseq_flags_has_lead(iseq: *const rb_iseq_t) -> bool;
+}
+extern "C" {
     pub fn rb_get_iseq_flags_has_opt(iseq: *const rb_iseq_t) -> bool;
 }
 extern "C" {
@@ -1452,7 +1458,10 @@ extern "C" {
     pub fn rb_get_iseq_flags_has_block(iseq: *const rb_iseq_t) -> bool;
 }
 extern "C" {
-    pub fn rb_get_iseq_flags_has_accepts_no_kwarg(iseq: *const rb_iseq_t) -> bool;
+    pub fn rb_get_iseq_flags_ambiguous_param0(iseq: *const rb_iseq_t) -> bool;
+}
+extern "C" {
+    pub fn rb_get_iseq_flags_accepts_no_kwarg(iseq: *const rb_iseq_t) -> bool;
 }
 extern "C" {
     pub fn rb_get_iseq_body_param_keyword(
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index e07b475a9f..b6b2d92b6d 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -217,6 +217,14 @@ make_counters! {
     invokesuper_me_changed,
     invokesuper_block,
 
+    invokeblock_none,
+    invokeblock_iseq_arg0_splat,
+    invokeblock_iseq_block_changed,
+    invokeblock_iseq_tag_changed,
+    invokeblock_ifunc,
+    invokeblock_proc,
+    invokeblock_symbol,
+
     leave_se_interrupt,
     leave_interp_return,
     leave_start_pc_non_zero,