ZJIT: Track block ISEQ pointers baked into yield guards for compaction

The `yield` dispatch guards `captured->code.iseq` against the profiled
block ISEQ, but it baked the ISEQ pointer into JIT code as a raw
`Const::CPtr`. The GC neither marked it nor updated it, so once
compaction moved the block ISEQ, the guard kept comparing against the
old address. Any block ISEQ allocated at that address later passed the
guard, and the JIT-to-JIT call that had already been patched for the
original block ran the original block's code with the new block's
`self` and EP. This showed up as `Array#each` running an unrelated
Bundler block in place of a test's block.

Bake the ISEQ in as a `Const::Value` so it is recorded as a GC offset,
marked, and updated on compaction. Add an `Iseq` HIR type so that
`Type::from_value` can type an `imemo_iseq`, like `CallableMethodEntry`.

Author: k0kubun

zjit/src/codegen_tests.rs

@@ -1127,6 +1127,25 @@ fn test_yield_inlined_caller_block_dispatches_without_guards() {
     });
 }
 
+#[test]
+fn test_yield_block_iseq_guard_survives_compaction() {
+    // The yield dispatch guards captured->code.iseq against the profiled block ISEQ baked into
+    // JIT code. ISEQs are movable GC objects, so compaction must update the baked pointer.
+    // Otherwise the guard would side-exit after compaction, or worse, match an unrelated ISEQ
+    // allocated at the old address later and dispatch to the wrong block.
+    // Disable inlining: when foo is inlined into test, yield takes the guard-free path instead.
+    with_inlining_threshold(0, || {
+        eval("
+            def foo = yield
+            def test = foo { 42 }
+            # Call it enough times to compile both test and foo (through test's JIT-to-JIT stub)
+            4.times { test }
+            GC.verify_compaction_references(expand_heap: true, toward: :empty) if GC.respond_to?(:compact)
+        ");
+        assert_snapshot!(assert_compiles("test"), @"42");
+    });
+}
+
 #[test]
 fn test_yield_with_lambda_arg() {
     // A lambda passed via &l is a proc handler (not imemo_iseq): yield falls back but runs.

zjit/src/cruby.rs

@@ -757,6 +757,12 @@ impl VALUE {
         unsafe { rb_IMEMO_TYPE_P(self, imemo_ment) == 1 }
     }
 
+    /// Return true if `self` is an ISEQ (`imemo_iseq`)
+    pub fn iseq_p(self) -> bool {
+        if self == VALUE(0) { return false; }
+        unsafe { rb_IMEMO_TYPE_P(self, imemo_iseq) == 1 }
+    }
+
     /// Assert that `self` is a method entry in debug builds
     pub fn as_cme(self) -> *const rb_callable_method_entry_t {
         let ptr: *const rb_callable_method_entry_t = self.as_ptr();

zjit/src/hir.rs

@@ -3226,7 +3226,7 @@ impl Function {
     /// Load `captured->code.iseq` from a `struct rb_captured_block *`.
     fn load_captured_code_iseq(&mut self, block: BlockId, captured: InsnId) -> InsnId {
         let offset: i32 = std::mem::offset_of!(rb_captured_block, code).try_into().unwrap();
-        self.load_field(block, captured, FieldName::code_iseq, offset, types::CPtr)
+        self.load_field(block, captured, FieldName::code_iseq, offset, types::Iseq)
     }
 
     /// Untag an ISEQ block handler into its `struct rb_captured_block *`:
@@ -3272,10 +3272,11 @@ impl Function {
             self.push_insn(block, Insn::GuardBitEquals { val: tag, expected: Const::CInt64(0x1), reason: Box::new(SideExitReason::InvokeBlockHandlerNotIseq), state, recompile: Some(Recompile) });
             let captured = self.untag_block_handler(block, block_handler);
 
-            // Guard captured->code.iseq is the profiled block iseq. Compare the raw imemo pointer:
-            // type inference (from_value) can't type an iseq imemo, so guard it as a CPtr identity.
+            // Guard captured->code.iseq is the profiled block iseq. The ISEQ is a GC object that
+            // can be moved by compaction, so bake it in as a `Const::Value` rather than a raw
+            // `Const::CPtr` to let the GC mark it and update the pointer embedded in JIT code.
             let captured_iseq = self.load_captured_code_iseq(block, captured);
-            self.push_insn(block, Insn::GuardBitEquals { val: captured_iseq, expected: Const::CPtr(block_iseq as *const u8), reason: Box::new(SideExitReason::InvokeBlockIseqChanged), state, recompile: Some(Recompile) });
+            self.push_insn(block, Insn::GuardBitEquals { val: captured_iseq, expected: Const::Value(VALUE::from(block_iseq)), reason: Box::new(SideExitReason::InvokeBlockIseqChanged), state, recompile: Some(Recompile) });
 
             let result = self.push_insn(block, Insn::InvokeBlockIseqDirect { iseq: block_iseq, captured, args, state });
             return (block, result);
@@ -3300,7 +3301,8 @@ impl Function {
 
         let mut compare_block = dispatch_block;
         for &block_iseq in iseqs {
-            let expected = self.push_insn(compare_block, Insn::Const { val: Const::CPtr(block_iseq as *const u8) });
+            // Use `Const::Value` so that the GC updates the ISEQ pointer on compaction. See above.
+            let expected = self.push_insn(compare_block, Insn::Const { val: Const::Value(VALUE::from(block_iseq)) });
             let iseq_matches = self.push_insn(compare_block, Insn::IsBitEqual { left: captured_iseq, right: expected });
             let direct_block = self.new_block(insn_idx);
             let miss_block = self.new_block(insn_idx);

zjit/src/hir/opt_tests.rs

@@ -4436,9 +4436,9 @@ mod hir_opt_tests {
           v30:CInt64[1] = GuardBitEquals v29, CInt64(1) recompile
           v31:CInt64[-4] = Const CInt64(-4)
           v32:CInt64 = IntAnd v27, v31
-          v33:CPtr = LoadField v32, :code_iseq@0x1001
-          v34:CPtr[CPtr(0x1002)] = GuardBitEquals v33, CPtr(0x1002) recompile
-          v35:BasicObject = InvokeBlockIseqDirect (0x1002), v32, v10, v12, v14, v16, v18, v20, v22, v24
+          v33:Iseq = LoadField v32, :code_iseq@0x1001
+          v34:Iseq[VALUE(0x1008)] = GuardBitEquals v33, Value(VALUE(0x1008)) recompile
+          v35:BasicObject = InvokeBlockIseqDirect (0x1008), v32, v10, v12, v14, v16, v18, v20, v22, v24
           CheckInterrupts
           Return v35
         ");
@@ -4558,19 +4558,19 @@ mod hir_opt_tests {
         bb5():
           v20:CInt64[-4] = Const CInt64(-4)
           v21:CInt64 = IntAnd v13, v20
-          v22:CPtr = LoadField v21, :code_iseq@0x1001
-          v23:CPtr[CPtr(0x1002)] = Const CPtr(0x1002)
+          v22:Iseq = LoadField v21, :code_iseq@0x1001
+          v23:Iseq[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           v24:CBool = IsBitEqual v22, v23
           CondBranch v24, bb7(), bb8()
         bb7():
-          v26:BasicObject = InvokeBlockIseqDirect (0x1002), v21, v10
+          v26:BasicObject = InvokeBlockIseqDirect (0x1008), v21, v10
           Jump bb4(v26)
         bb8():
-          v28:CPtr[CPtr(0x1003)] = Const CPtr(0x1003)
+          v28:Iseq[VALUE(0x1010)] = Const Value(VALUE(0x1010))
           v29:CBool = IsBitEqual v22, v28
           CondBranch v29, bb9(), bb10()
         bb9():
-          v31:BasicObject = InvokeBlockIseqDirect (0x1003), v21, v10
+          v31:BasicObject = InvokeBlockIseqDirect (0x1010), v21, v10
           Jump bb4(v31)
         bb10():
           Jump bb6()
@@ -4623,19 +4623,19 @@ mod hir_opt_tests {
         bb5():
           v20:CInt64[-4] = Const CInt64(-4)
           v21:CInt64 = IntAnd v13, v20
-          v22:CPtr = LoadField v21, :code_iseq@0x1001
-          v23:CPtr[CPtr(0x1002)] = Const CPtr(0x1002)
+          v22:Iseq = LoadField v21, :code_iseq@0x1001
+          v23:Iseq[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           v24:CBool = IsBitEqual v22, v23
           CondBranch v24, bb7(), bb8()
         bb7():
-          v26:BasicObject = InvokeBlockIseqDirect (0x1002), v21, v10
+          v26:BasicObject = InvokeBlockIseqDirect (0x1008), v21, v10
           Jump bb4(v26)
         bb8():
-          v28:CPtr[CPtr(0x1003)] = Const CPtr(0x1003)
+          v28:Iseq[VALUE(0x1010)] = Const Value(VALUE(0x1010))
           v29:CBool = IsBitEqual v22, v28
           CondBranch v29, bb9(), bb10()
         bb9():
-          v31:BasicObject = InvokeBlockIseqDirect (0x1003), v21, v10
+          v31:BasicObject = InvokeBlockIseqDirect (0x1010), v21, v10
           Jump bb4(v31)
         bb10():
           Jump bb6()
@@ -4687,19 +4687,19 @@ mod hir_opt_tests {
         bb5():
           v20:CInt64[-4] = Const CInt64(-4)
           v21:CInt64 = IntAnd v13, v20
-          v22:CPtr = LoadField v21, :code_iseq@0x1001
-          v23:CPtr[CPtr(0x1002)] = Const CPtr(0x1002)
+          v22:Iseq = LoadField v21, :code_iseq@0x1001
+          v23:Iseq[VALUE(0x1008)] = Const Value(VALUE(0x1008))
           v24:CBool = IsBitEqual v22, v23
           CondBranch v24, bb7(), bb8()
         bb7():
-          v26:BasicObject = InvokeBlockIseqDirect (0x1002), v21, v10
+          v26:BasicObject = InvokeBlockIseqDirect (0x1008), v21, v10
           Jump bb4(v26)
         bb8():
-          v28:CPtr[CPtr(0x1003)] = Const CPtr(0x1003)
+          v28:Iseq[VALUE(0x1010)] = Const Value(VALUE(0x1010))
           v29:CBool = IsBitEqual v22, v28
           CondBranch v29, bb9(), bb10()
         bb9():
-          v31:BasicObject = InvokeBlockIseqDirect (0x1003), v21, v10
+          v31:BasicObject = InvokeBlockIseqDirect (0x1010), v21, v10
           Jump bb4(v31)
         bb10():
           Jump bb6()
@@ -4752,12 +4752,12 @@ mod hir_opt_tests {
         bb5():
           v20:CInt64[-4] = Const CInt64(-4)
           v21:CInt64 = IntAnd v13, v20
-          v22:CPtr = LoadField v21, :code_iseq@0x1011
-          v23:CPtr[CPtr(0x1012)] = Const CPtr(0x1012)
+          v22:Iseq = LoadField v21, :code_iseq@0x1011
+          v23:Iseq[VALUE(0x1018)] = Const Value(VALUE(0x1018))
           v24:CBool = IsBitEqual v22, v23
           CondBranch v24, bb7(), bb8()
         bb7():
-          v26:BasicObject = InvokeBlockIseqDirect (0x1012), v21, v10
+          v26:BasicObject = InvokeBlockIseqDirect (0x1018), v21, v10
           Jump bb4(v26)
         bb8():
           Jump bb6()
@@ -19914,9 +19914,9 @@ mod hir_opt_tests {
           v83:CInt64[1] = GuardBitEquals v82, CInt64(1) recompile
           v84:CInt64[-4] = Const CInt64(-4)
           v85:CInt64 = IntAnd v80, v84
-          v86:CPtr = LoadField v85, :code_iseq@0x1001
-          v87:CPtr[CPtr(0x1002)] = GuardBitEquals v86, CPtr(0x1002) recompile
-          v88:BasicObject = InvokeBlockIseqDirect (0x1002), v85, v77
+          v86:Iseq = LoadField v85, :code_iseq@0x1001
+          v87:Iseq[VALUE(0x1008)] = GuardBitEquals v86, Value(VALUE(0x1008)) recompile
+          v88:BasicObject = InvokeBlockIseqDirect (0x1008), v85, v77
           v92:Fixnum[1] = Const Value(1)
           v93:Fixnum = FixnumAdd v49, v92
           PatchPoint NoEPEscape(each)

zjit/src/hir_type/gen_hir_type.rb

@@ -51,6 +51,7 @@ def to_graphviz type, f
 undef_ = value.subtype "Undef"
 # TODO(max): Figure out if CME should be a subtype of IMemo
 value.subtype "CallableMethodEntry"  # rb_callable_method_entry_t*
+value.subtype "Iseq"  # rb_iseq_t*
 basic_object = value.subtype "BasicObject"
 basic_object_exact = basic_object.subtype "BasicObjectExact"
 basic_object_subclass = basic_object.subtype "BasicObjectSubclass"

zjit/src/hir_type/hir_type.inc.rs

@@ -333,6 +333,10 @@ impl Type {
             // valid on imemo.
             Type::new(bits::CallableMethodEntry, Specialization::Object(val))
         }
+        else if val.iseq_p() {
+            // Same as CME: an ISEQ is an imemo, so class_of is not valid on it.
+            Type::new(bits::Iseq, Specialization::Object(val))
+        }
         else {
             Self::from_heap_object(val)
         }
@@ -1119,6 +1123,18 @@ mod tests {
         });
     }
 
+    #[test]
+    fn iseq() {
+        crate::cruby::with_rubyvm(|| {
+            let iseq = crate::cruby::test_utils::get_proc_iseq("proc { 1 }");
+            assert!(!iseq.is_null());
+            let iseq_value: VALUE = iseq.into();
+            let ty = Type::from_value(iseq_value);
+            assert_subtype(ty, types::Iseq);
+            assert!(ty.ruby_object_known());
+        });
+    }
+
     #[test]
     fn string_subclass_is_string_subtype() {
         crate::cruby::with_rubyvm(|| {