ZJIT: Keep surviving JIT frame across longjmp

A tag established inside a live ZJIT frame can catch a longjmp without unwinding that native frame. Record that fact on rb_vm_tag and materialize only frames above the target, preserving its JITFrame link and keeping JITFrame saves at a single store.

Author: k0kubun

eval_intern.h

@@ -106,11 +106,14 @@ extern int select_large_fdset(int, fd_set *, fd_set *, fd_set *, struct timeval
   EC_SAVE_TAG_CFP(_tag, _ec); \
   rb_vm_tag_jmpbuf_init(&_tag.buf); \
 
-// Remember the CFP as of EC_PUSH_TAG so that ZJIT can materialize frames
-// only up to longjmp's target CFP. When a C method does longjmp inside it,
-// the target CFP may not be equal to the VM_FRAME_FLAG_FINISH frame.
+// Remember the CFP and whether it was already running in ZJIT as of
+// EC_PUSH_TAG. When a C method does longjmp inside it, the target CFP may not
+// be equal to the VM_FRAME_FLAG_FINISH frame. If its ZJIT frame predates this
+// tag, its native stack survives the jump and should not be materialized.
 #if USE_ZJIT
-# define EC_SAVE_TAG_CFP(_tag, _ec) _tag.cfp = _ec->cfp
+# define EC_SAVE_TAG_CFP(_tag, _ec) \
+    _tag.cfp = _ec->cfp; \
+    _tag.zjit_frame_active = CFP_ZJIT_FRAME_P(_ec->cfp)
 #else
 # define EC_SAVE_TAG_CFP(_tag, _ec)
 #endif
@@ -166,7 +169,7 @@ rb_ec_tag_jump(const rb_execution_context_t *ec, enum ruby_tag_type st)
 {
     RUBY_ASSERT(st > TAG_NONE && st <= TAG_FATAL, ": Invalid tag jump: %d", (int)st);
 #if USE_ZJIT
-    rb_zjit_materialize_frames(ec, ec->cfp);
+    rb_zjit_materialize_frames_for_longjmp(ec, ec->cfp);
 #endif
     ec->tag->state = st;
     ruby_longjmp(RB_VM_TAG_JMPBUF_GET(ec->tag->buf), 1);

vm.c

@@ -2884,14 +2884,15 @@ vm_exec_loop(rb_execution_context_t *ec, enum ruby_tag_type state,
 
 #if USE_ZJIT
 // Materialize JITFrame-enabled CFP into interpreter-compatible CFP
-void
-rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cfp)
+static void
+zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cfp, bool materialize_target)
 {
     if (!rb_zjit_enabled_p) return;
     const rb_control_frame_t *end_cfp = ec->tag->cfp;
     VM_ASSERT(cfp <= end_cfp);
 
     while (true) {
+        if (cfp == end_cfp && !materialize_target) break;
         if (CFP_ZJIT_FRAME_P(cfp)) {
             const zjit_jit_frame_t *jit_frame = CFP_ZJIT_FRAME(cfp);
             cfp->pc = jit_frame->pc;
@@ -2931,6 +2932,20 @@ rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t
         cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp);
     }
 }
+
+void
+rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cfp)
+{
+    zjit_materialize_frames(ec, cfp, true);
+}
+
+void
+rb_zjit_materialize_frames_for_longjmp(const rb_execution_context_t *ec, rb_control_frame_t *cfp)
+{
+    // A ZJIT frame active before the tag's setjmp is below it on the native
+    // stack and survives longjmp. Materialize only the frames unwound above it.
+    zjit_materialize_frames(ec, cfp, !ec->tag->zjit_frame_active);
+}
 #endif
 
 static inline VALUE

vm_core.h

@@ -1029,6 +1029,8 @@ struct rb_vm_tag {
 #if USE_ZJIT
     // ec->cfp as of EC_PUSH_TAG, which is saved for materializing JITFrame.
     rb_control_frame_t *cfp;
+    // Whether cfp had a ZJIT frame before this tag's setjmp was established.
+    bool zjit_frame_active;
 #endif
 };
 

zjit.h

@@ -92,6 +92,7 @@ void rb_zjit_invalidate_no_singleton_class(VALUE klass);
 void rb_zjit_invalidate_root_box(void);
 void rb_zjit_jit_frame_update_references(zjit_jit_frame_t *jit_frame);
 void rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cfp);
+void rb_zjit_materialize_frames_for_longjmp(const rb_execution_context_t *ec, rb_control_frame_t *cfp);
 size_t rb_zjit_hash_new_size(void);
 bool rb_zjit_class_allocate_instance_fastpath(VALUE klass, size_t *size_out, shape_id_t *shape_id_out);
 bool rb_zjit_str_resurrect_fastpath(VALUE str, bool chilled, size_t *size_out, VALUE *flags_out, long *len_out, size_t *byte_size_out);
@@ -136,6 +137,7 @@ static inline void rb_zjit_invalidate_no_singleton_class(VALUE klass) {}
 static inline void rb_zjit_invalidate_root_box(void) {}
 static inline void rb_zjit_jit_frame_update_references(zjit_jit_frame_t *jit_frame) {}
 static inline void rb_zjit_materialize_frames(const rb_execution_context_t *ec, rb_control_frame_t *cfp) {}
+static inline void rb_zjit_materialize_frames_for_longjmp(const rb_execution_context_t *ec, rb_control_frame_t *cfp) {}
 static inline const zjit_jit_frame_t *CFP_ZJIT_FRAME(const rb_control_frame_t *cfp) { return NULL; }
 #endif // #if USE_ZJIT
 

zjit/src/codegen_tests.rs

@@ -6,7 +6,7 @@ use crate::backend::lir::Assembler;
 use crate::codegen::max_iseq_versions;
 use crate::cruby::*;
 use crate::hir::{Insn, iseq_to_hir};
-use crate::options::{get_option, rb_zjit_prepare_options, set_call_threshold, set_inline_threshold};
+use crate::options::{get_option, rb_zjit_prepare_options, set_call_threshold, set_inline_threshold, set_max_versions};
 use crate::payload::IseqVersion;
 use crate::hir::tests::hir_build_tests::assert_contains_opcode;
 use crate::payload::*;
@@ -7198,6 +7198,61 @@ fn test_regression_gc_stress_with_lazy_block_code() {
     "#), @":ok");
 }
 
+// Hash recursion uses catch/throw internally. The target frame remains in JIT
+// code after the caught throw, so longjmp must not materialize and detach it
+// before a callee side exit uses its updated PC and stack map.
+#[test]
+fn test_keep_jit_frame_for_caught_jump() {
+    rb_zjit_prepare_options();
+    let old_call_threshold = unsafe { crate::options::rb_zjit_call_threshold };
+    let old_inline_threshold = get_option!(inline_threshold);
+    let old_max_versions = get_option!(max_versions);
+    set_call_threshold(1);
+    set_inline_threshold(0);
+    set_max_versions(2);
+    let result = inspect(r#"
+        module KeepJITFrameAssertions
+          def assert_receiver(*)
+            raise unless is_a?(KeepJITFrameBase)
+          end
+        end
+
+        class KeepJITFrameBase
+          include KeepJITFrameAssertions
+
+          def hash_class = Hash
+
+          def test
+            hash = hash_class[]
+            recursive = [hash]
+            hash[:x] = recursive
+            object = Object.new
+            lookup = { hash => object }
+
+            [recursive, [hash]].each do |key|
+              key = { x: key }
+              assert_receiver(object, lookup[key], -> { key.inspect })
+            end
+          end
+        end
+
+        class KeepJITFrameHash < Hash
+        end
+
+        class KeepJITFrameSubclass < KeepJITFrameBase
+          def hash_class = KeepJITFrameHash
+        end
+
+        KeepJITFrameBase.new.test
+        KeepJITFrameSubclass.new.test
+        :ok
+    "#);
+    set_max_versions(old_max_versions);
+    set_inline_threshold(old_inline_threshold);
+    set_call_threshold(old_call_threshold);
+    assert_snapshot!(result, @":ok");
+}
+
 #[test]
 fn test_float_arithmetic() {
     set_call_threshold(1);