ZJIT: Untag block handler (#15085)

Storing the tagged block handler in profiles is not GC-safe (nice catch,
Kokubun). Store the untagged block handler instead.

Fix bug in https://github.com/ruby/ruby/pull/15051

6 files changed, 28 insertions, 23 deletions

diff --git a/vm_insnhelper.c b/vm_insnhelper.c
index e1ec5e63ec..63dcaba8a3 100644
--- a/vm_insnhelper.c
+++ b/vm_insnhelper.c
@@ -5497,12 +5497,6 @@ vm_invoke_proc_block(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp,
     return vm_invoke_block(ec, reg_cfp, calling, ci, is_lambda, block_handler);
 }
 
-enum rb_block_handler_type
-rb_vm_block_handler_type(VALUE block_handler)
-{
-    return vm_block_handler_type(block_handler);
-}
-
 static inline VALUE
 vm_invoke_block(rb_execution_context_t *ec, rb_control_frame_t *reg_cfp,
                 struct rb_calling_info *calling, const struct rb_callinfo *ci,
@@ -6065,10 +6059,30 @@ vm_define_method(const rb_execution_context_t *ec, VALUE obj, ID id, VALUE iseqv
     }
 }
 
+// Return the untagged block handler:
+// * If it's an ISEQ or an IFUNC, fetch it from its rb_captured_block
+// * If it's a PROC or SYMBOL, return it as is
+static VALUE
+rb_vm_untag_block_handler(VALUE block_handler)
+{
+    switch (vm_block_handler_type(block_handler)) {
+      case block_handler_type_iseq:
+      case block_handler_type_ifunc: {
+        struct rb_captured_block *captured = VM_TAGGED_PTR_REF(block_handler, 0x03);
+        return captured->code.val;
+      }
+      case block_handler_type_proc:
+      case block_handler_type_symbol:
+        return block_handler;
+      default:
+        rb_bug("rb_vm_untag_block_handler: unreachable");
+    }
+}
+
 VALUE
-rb_vm_get_block_handler(rb_control_frame_t *reg_cfp)
+rb_vm_get_untagged_block_handler(rb_control_frame_t *reg_cfp)
 {
-    return VM_CF_BLOCK_HANDLER(reg_cfp);
+    return rb_vm_untag_block_handler(VM_CF_BLOCK_HANDLER(reg_cfp));
 }
 
 static VALUE
diff --git a/zjit.c b/zjit.c
index 72e6fe1424..d1f192801a 100644
--- a/zjit.c
+++ b/zjit.c
@@ -302,8 +302,7 @@ rb_zjit_class_has_default_allocator(VALUE klass)
 }
 
 
-VALUE rb_vm_get_block_handler(rb_control_frame_t *reg_cfp);
-enum rb_block_handler_type rb_vm_block_handler_type(VALUE block_handler);
+VALUE rb_vm_get_untagged_block_handler(rb_control_frame_t *reg_cfp);
 
 // Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them.
 VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
diff --git a/zjit/bindgen/src/main.rs b/zjit/bindgen/src/main.rs
index bbb3b54d6c..95209375dc 100644
--- a/zjit/bindgen/src/main.rs
+++ b/zjit/bindgen/src/main.rs
@@ -399,8 +399,7 @@ fn main() {
         .allowlist_function("rb_yarv_str_eql_internal")
         .allowlist_function("rb_str_neq_internal")
         .allowlist_function("rb_yarv_ary_entry_internal")
-        .allowlist_function("rb_vm_get_block_handler")
-        .allowlist_function("rb_vm_block_handler_type")
+        .allowlist_function("rb_vm_get_untagged_block_handler")
         .allowlist_function("rb_FL_TEST")
         .allowlist_function("rb_FL_TEST_RAW")
         .allowlist_function("rb_RB_TYPE_P")
diff --git a/zjit/src/cruby_bindings.inc.rs b/zjit/src/cruby_bindings.inc.rs
index dc9d0d144c..8623958829 100644
--- a/zjit/src/cruby_bindings.inc.rs
+++ b/zjit/src/cruby_bindings.inc.rs
@@ -571,11 +571,6 @@ pub struct rb_captured_block__bindgen_ty_1 {
     pub val: __BindgenUnionField<VALUE>,
     pub bindgen_union_field: u64,
 }
-pub const block_handler_type_iseq: rb_block_handler_type = 0;
-pub const block_handler_type_ifunc: rb_block_handler_type = 1;
-pub const block_handler_type_symbol: rb_block_handler_type = 2;
-pub const block_handler_type_proc: rb_block_handler_type = 3;
-pub type rb_block_handler_type = u32;
 pub const block_type_iseq: rb_block_type = 0;
 pub const block_type_ifunc: rb_block_type = 1;
 pub const block_type_symbol: rb_block_type = 2;
@@ -1339,8 +1334,7 @@ unsafe extern "C" {
     pub fn rb_zjit_class_initialized_p(klass: VALUE) -> bool;
     pub fn rb_zjit_class_get_alloc_func(klass: VALUE) -> rb_alloc_func_t;
     pub fn rb_zjit_class_has_default_allocator(klass: VALUE) -> bool;
-    pub fn rb_vm_get_block_handler(reg_cfp: *mut rb_control_frame_t) -> VALUE;
-    pub fn rb_vm_block_handler_type(block_handler: VALUE) -> rb_block_handler_type;
+    pub fn rb_vm_get_untagged_block_handler(reg_cfp: *mut rb_control_frame_t) -> VALUE;
     pub fn rb_iseq_encoded_size(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
     pub fn rb_iseq_pc_at_idx(iseq: *const rb_iseq_t, insn_idx: u32) -> *mut VALUE;
     pub fn rb_iseq_opcode_at_pc(iseq: *const rb_iseq_t, pc: *const VALUE) -> ::std::os::raw::c_int;
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index e0a2e0fdff..07ffe4b00a 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -4421,10 +4421,9 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                             let summary = TypeDistributionSummary::new(&self_type_distribution);
                             if summary.is_monomorphic() {
                                 let obj = summary.bucket(0).class();
-                                let bh_type = unsafe { rb_vm_block_handler_type(obj) };
-                                if bh_type == block_handler_type_iseq {
+                                if unsafe { rb_IMEMO_TYPE_P(obj, imemo_iseq) == 1 } {
                                     fun.push_insn(block, Insn::IncrCounter(Counter::invokeblock_handler_monomorphic_iseq));
-                                } else if bh_type == block_handler_type_ifunc {
+                                } else if unsafe { rb_IMEMO_TYPE_P(obj, imemo_ifunc) == 1 } {
                                     fun.push_insn(block, Insn::IncrCounter(Counter::invokeblock_handler_monomorphic_ifunc));
                                 } else {
                                     fun.push_insn(block, Insn::IncrCounter(Counter::invokeblock_handler_monomorphic_other));
diff --git a/zjit/src/profile.rs b/zjit/src/profile.rs
index c58999668e..3366fe8e58 100644
--- a/zjit/src/profile.rs
+++ b/zjit/src/profile.rs
@@ -45,7 +45,7 @@ impl Profiler {
     }
 
     fn peek_at_block_handler(&self) -> VALUE {
-        unsafe { rb_vm_get_block_handler(self.cfp) }
+        unsafe { rb_vm_get_untagged_block_handler(self.cfp) }
     }
 }