ZJIT: Don't retain Proc and IFUNC block handlers in profiles

Author: k0kubun

zjit/src/codegen_tests.rs

@@ -6697,6 +6697,33 @@ fn test_profiled_singleton_class_does_not_retain_young_object() {
     "), @"true");
 }
 
+#[test]
+fn test_profiled_proc_block_handler_does_not_retain_proc() {
+    // Profile more than one call so that the profile can hold several Procs
+    rb_zjit_prepare_options();
+    let num_profiles = get_option!(num_profiles);
+    set_call_threshold(CallThreshold::from(num_profiles) + 2);
+
+    assert_snapshot!(inspect("
+        def profiled_proc_take = yield
+        def profiled_proc_forward(&blk) = profiled_proc_take(&blk)
+
+        PROFILED_PROC_OBJECTS = ObjectSpace::WeakMap.new
+        def profiled_proc_make(i)
+          obj = Object.new
+          PROFILED_PROC_OBJECTS[i] = obj
+          pr = proc { obj }
+          profiled_proc_forward(&pr)
+          nil
+        end
+
+        100.times { |i| profiled_proc_make(i) }
+        4.times { GC.start(full_mark: true, immediate_sweep: true) }
+        # Allow one object kept alive by conservative stack scanning
+        PROFILED_PROC_OBJECTS.keys.size <= 1
+    "), @"true");
+}
+
 #[test]
 fn test_profile_under_nested_jit_call() {
     assert_snapshot!(inspect("

zjit/src/hir.rs

@@ -15,7 +15,7 @@ use std::{
 use crate::hir_type::{Type, types};
 use crate::hir_effect::{Effect, abstract_heaps, effects};
 use crate::bitset::BitSet;
-use crate::profile::{ProfiledType, SplatLength, TypeDistributionSummary};
+use crate::profile::{ProfiledType, SplatLength, TypeDistributionSummary, PROFILED_IFUNC_BLOCK_HANDLER, PROFILED_PROC_BLOCK_HANDLER};
 use crate::stats::{Counter, incr_counter};
 use SendFallbackReason::*;
 
@@ -8974,7 +8974,7 @@ fn add_iseq_to_hir(
                                 let obj = summary.bucket(0).class();
                                 if unsafe { rb_IMEMO_TYPE_P(obj, imemo_iseq) == 1 } {
                                     fun.count(block, Counter::invokeblock_handler_monomorphic_iseq);
-                                } else if unsafe { rb_IMEMO_TYPE_P(obj, imemo_ifunc) == 1 } {
+                                } else if obj == PROFILED_IFUNC_BLOCK_HANDLER {
                                     fun.count(block, Counter::invokeblock_handler_monomorphic_ifunc);
                                 } else {
                                     fun.count(block, Counter::invokeblock_handler_monomorphic_other);
@@ -9001,15 +9001,15 @@ fn add_iseq_to_hir(
                             let obj = summary.bucket(0).class();
                             if unsafe { rb_IMEMO_TYPE_P(obj, imemo_iseq) == 1} {
                                 fun.count(block, Counter::getblockparamproxy_handler_iseq);
-                            } else if unsafe { rb_IMEMO_TYPE_P(obj, imemo_ifunc) == 1} {
+                            } else if obj == PROFILED_IFUNC_BLOCK_HANDLER {
                                 fun.count(block, Counter::getblockparamproxy_handler_ifunc);
                             }
                             else if obj.nil_p() {
                                 fun.count(block, Counter::getblockparamproxy_handler_nil);
                             }
                             else if obj.symbol_p() {
                                 fun.count(block, Counter::getblockparamproxy_handler_symbol);
-                            } else if unsafe { rb_obj_is_proc(obj).test() } {
+                            } else if obj == PROFILED_PROC_BLOCK_HANDLER {
                                 fun.count(block, Counter::getblockparamproxy_handler_proc);
                             }
                         } else if summary.is_polymorphic() || summary.is_skewed_polymorphic() {
@@ -9699,12 +9699,10 @@ fn add_iseq_to_hir(
                             let obj = profiled_type.class();
                             if obj.nil_p() {
                                 Some(Self::Nil)
-                            } else if unsafe {
-                                rb_IMEMO_TYPE_P(obj, imemo_iseq) == 1
-                                    || rb_IMEMO_TYPE_P(obj, imemo_ifunc) == 1
-                            } {
+                            } else if unsafe { rb_IMEMO_TYPE_P(obj, imemo_iseq) == 1 }
+                                || obj == PROFILED_IFUNC_BLOCK_HANDLER {
                                 Some(Self::IseqOrIfunc)
-                            } else if unsafe { rb_obj_is_proc(obj).test() } {
+                            } else if obj == PROFILED_PROC_BLOCK_HANDLER {
                                 Some(Self::Proc)
                             } else {
                                 None
@@ -10378,7 +10376,7 @@ fn add_iseq_to_hir(
                     });
 
                     let is_ifunc = (flags & (VM_CALL_ARGS_SPLAT | VM_CALL_KW_SPLAT | VM_CALL_KWARG)) == 0
-                        && block_handler_class.is_some_and(|obj| unsafe { rb_IMEMO_TYPE_P(obj, imemo_ifunc) == 1 });
+                        && block_handler_class.is_some_and(|obj| obj == PROFILED_IFUNC_BLOCK_HANDLER);
 
                     // Collect the profiled ISEQ blocks that can be invoked directly with a JIT-to-JIT call.
                     let mut fallback_reason = InvokeBlockNotSpecialized;

zjit/src/profile.rs

@@ -235,7 +235,7 @@ fn profile_block_handler(profiler: &mut Profiler, profile: &mut IseqProfile) {
     if entry.opnd_types.is_empty() {
         entry.opnd_types.resize(1, TypeDistribution::new());
     }
-    let obj = profiler.peek_at_block_handler();
+    let obj = block_handler_profile_value(profiler.peek_at_block_handler());
     let ty = ProfiledType::object(obj);
     VALUE::from(profiler.iseq).write_barrier(ty.class());
     entry.opnd_types[0].observe(ty);
@@ -252,11 +252,29 @@ fn profile_getblockparamproxy(profiler: &mut Profiler, profile: &mut IseqProfile
     let block_handler = unsafe { *ep.offset(VM_ENV_DATA_INDEX_SPECVAL as isize) };
     let untagged = unsafe { rb_vm_untag_block_handler(block_handler) };
 
-    let ty = ProfiledType::object(untagged);
+    let ty = ProfiledType::object(block_handler_profile_value(untagged));
     VALUE::from(profiler.iseq).write_barrier(ty.class());
     entry.opnd_types[0].observe(ty);
 }
 
+/// Fixnum values to represent a profiled block handler type without retaining its reference.
+/// If we kept them in profiles, imemo for IFUNC or a Proc object with a captured environment
+/// would be kept in memory. However, for IFUNC and Proc, we're only interested in the block
+/// handler type. We use Fixnum values because Fixnum cannot be a block handler.
+pub const PROFILED_IFUNC_BLOCK_HANDLER: VALUE = VALUE::fixnum_from_usize(1);
+pub const PROFILED_PROC_BLOCK_HANDLER: VALUE = VALUE::fixnum_from_usize(2);
+
+/// Map an untagged block handler to the value recorded in the profile.
+fn block_handler_profile_value(block_handler: VALUE) -> VALUE {
+    if unsafe { rb_IMEMO_TYPE_P(block_handler, imemo_ifunc) == 1 } {
+        PROFILED_IFUNC_BLOCK_HANDLER
+    } else if unsafe { rb_obj_is_proc(block_handler).test() } {
+        PROFILED_PROC_BLOCK_HANDLER
+    } else {
+        block_handler
+    }
+}
+
 fn profile_invokesuper(profiler: &mut Profiler, profile: &mut IseqProfile) {
     let cme = unsafe { rb_vm_frame_method_entry(profiler.cfp) };
     let cme_value = VALUE(cme as usize);  // CME is a T_IMEMO, which is a VALUE