Write JITFrame pointer for C frames, skip iseq/block_code writes

Instead of writing 3 fields to C frames (jit_return=0, iseq=0,
block_code=0), write a single JITFrame pointer with {pc=NULL,
iseq=NULL, materialize_block_code=true}. This eliminates 2 memory
writes per C method call.

Add rb_zjit_cfp_has_iseq() and rb_zjit_cfp_has_pc() helpers that
check JITFrame first when present, since cfp->iseq and cfp->pc may
be stale for frames with JITFrame. Replace all cfp->iseq ||
CFP_JIT_RETURN(cfp) and cfp->pc || CFP_JIT_RETURN(cfp) patterns
with the new helpers.

Author: k0kubun

cont.c

@@ -1348,7 +1348,7 @@ rb_jit_cont_each_iseq(rb_iseq_callback callback, void *data)
 
         const rb_control_frame_t *cfp = cont->ec->cfp;
         while (!RUBY_VM_CONTROL_FRAME_STACK_OVERFLOW_P(cont->ec, cfp)) {
-            if ((cfp->pc || CFP_JIT_RETURN(cfp)) && (cfp->iseq || CFP_JIT_RETURN(cfp))) {
+            if (rb_zjit_cfp_has_pc(cfp) && rb_zjit_cfp_has_iseq(cfp)) {
                 const rb_iseq_t *iseq = rb_zjit_cfp_iseq(cfp);
                 if (iseq && imemo_type((VALUE)iseq) == imemo_iseq) {
                     callback(iseq, data);

gc.c

@@ -1000,7 +1000,7 @@ gc_validate_pc(VALUE obj)
 
     rb_execution_context_t *ec = GET_EC();
     const rb_control_frame_t *cfp = ec->cfp;
-    if (cfp && VM_FRAME_RUBYFRAME_P(cfp) && (cfp->pc || CFP_JIT_RETURN(cfp))) {
+    if (cfp && VM_FRAME_RUBYFRAME_P(cfp) && rb_zjit_cfp_has_pc(cfp)) {
         const VALUE *iseq_encoded = ISEQ_BODY(rb_zjit_cfp_iseq(cfp))->iseq_encoded;
         const VALUE *iseq_encoded_end = iseq_encoded + ISEQ_BODY(rb_zjit_cfp_iseq(cfp))->iseq_size;
         RUBY_ASSERT(rb_zjit_cfp_pc(cfp) >= iseq_encoded, "PC not set when allocating, breaking tracing");

vm.c

@@ -941,7 +941,7 @@ rb_control_frame_t *
 rb_vm_get_binding_creatable_next_cfp(const rb_execution_context_t *ec, const rb_control_frame_t *cfp)
 {
     while (!RUBY_VM_CONTROL_FRAME_STACK_OVERFLOW_P(ec, cfp)) {
-        if (cfp->iseq || CFP_JIT_RETURN(cfp)) {
+        if (rb_zjit_cfp_has_iseq(cfp)) {
             return (rb_control_frame_t *)cfp;
         }
         cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp);
@@ -1997,7 +1997,7 @@ rb_vm_invoke_proc_with_self(rb_execution_context_t *ec, rb_proc_t *proc, VALUE s
 VALUE *
 rb_vm_svar_lep(const rb_execution_context_t *ec, const rb_control_frame_t *cfp)
 {
-    while ((cfp->pc == 0 && !CFP_JIT_RETURN(cfp)) || (cfp->iseq == 0 && !CFP_JIT_RETURN(cfp))) {
+    while (!rb_zjit_cfp_has_pc(cfp) || !rb_zjit_cfp_has_iseq(cfp)) {
         if (VM_FRAME_TYPE(cfp) == VM_FRAME_MAGIC_IFUNC) {
             struct vm_ifunc *ifunc = (struct vm_ifunc *)rb_zjit_cfp_iseq(cfp);
             return ifunc->svar_lep;

vm_backtrace.c

@@ -102,7 +102,7 @@ calc_node_id(const rb_iseq_t *iseq, const VALUE *pc)
 int
 rb_vm_get_sourceline(const rb_control_frame_t *cfp)
 {
-    if (VM_FRAME_RUBYFRAME_P(cfp) && (cfp->iseq || CFP_JIT_RETURN(cfp))) {
+    if (VM_FRAME_RUBYFRAME_P(cfp) && rb_zjit_cfp_has_iseq(cfp)) {
         const rb_iseq_t *iseq = rb_zjit_cfp_iseq(cfp);
         int line = calc_lineno(iseq, rb_zjit_cfp_pc(cfp));
         if (line != 0) {
@@ -688,8 +688,8 @@ rb_ec_partial_backtrace_object(const rb_execution_context_t *ec, long start_fram
     }
 
     for (; cfp != end_cfp && (bt->backtrace_size < num_frames); cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp)) {
-        if (cfp->iseq || CFP_JIT_RETURN(cfp)) {
-            if (cfp->pc || CFP_JIT_RETURN(cfp)) {
+        if (rb_zjit_cfp_has_iseq(cfp)) {
+            if (rb_zjit_cfp_has_pc(cfp)) {
                 if (start_frame > 0) {
                     start_frame--;
                 }
@@ -753,7 +753,7 @@ rb_ec_partial_backtrace_object(const rb_execution_context_t *ec, long start_fram
     // is the one of the caller Ruby frame, so if the last entry is a C frame we find the caller Ruby frame here.
     if (backpatch_counter > 0) {
         for (; cfp != end_cfp; cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp)) {
-            if ((cfp->iseq || CFP_JIT_RETURN(cfp)) && (cfp->pc || CFP_JIT_RETURN(cfp)) && !(skip_internal && is_internal_location(rb_zjit_cfp_iseq(cfp)))) {
+            if (rb_zjit_cfp_has_iseq(cfp) && rb_zjit_cfp_has_pc(cfp) && !(skip_internal && is_internal_location(rb_zjit_cfp_iseq(cfp)))) {
                 VM_ASSERT(!skip_next_frame); // ISEQ_TYPE_RESCUE/ISEQ_TYPE_ENSURE should have a caller Ruby ISEQ, not a cfunc
                 bt_backpatch_loc(backpatch_counter, loc, rb_zjit_cfp_iseq(cfp), rb_zjit_cfp_pc(cfp));
                 RB_OBJ_WRITTEN(btobj, Qundef, rb_zjit_cfp_iseq(cfp));
@@ -1020,8 +1020,8 @@ backtrace_each(const rb_execution_context_t *ec,
     /* SDR(); */
     for (i=0, cfp = start_cfp; i<size; i++, cfp = RUBY_VM_NEXT_CONTROL_FRAME(cfp)) {
         /* fprintf(stderr, "cfp: %d\n", (rb_control_frame_t *)(ec->vm_stack + ec->vm_stack_size) - cfp); */
-        if (cfp->iseq || CFP_JIT_RETURN(cfp)) {
-            if (cfp->pc || CFP_JIT_RETURN(cfp)) {
+        if (rb_zjit_cfp_has_iseq(cfp)) {
+            if (rb_zjit_cfp_has_pc(cfp)) {
                 iter_iseq(arg, cfp);
             }
         }
@@ -1747,7 +1747,7 @@ thread_profile_frames(rb_execution_context_t *ec, int start, int limit, VALUE *b
     end_cfp = RUBY_VM_NEXT_CONTROL_FRAME(end_cfp);
 
     for (i=0; i<limit && cfp != end_cfp; cfp = RUBY_VM_PREVIOUS_CONTROL_FRAME(cfp)) {
-        if (VM_FRAME_RUBYFRAME_P_UNCHECKED(cfp) && (cfp->pc != 0 || CFP_JIT_RETURN(cfp))) {
+        if (VM_FRAME_RUBYFRAME_P_UNCHECKED(cfp) && rb_zjit_cfp_has_pc(cfp)) {
             if (start > 0) {
                 start--;
                 continue;

vm_dump.c

@@ -119,7 +119,7 @@ control_frame_dump(const rb_execution_context_t *ec, const rb_control_frame_t *c
         selfstr = "";
     }
 
-    if (cfp->iseq || CFP_JIT_RETURN(cfp)) {
+    if (rb_zjit_cfp_has_iseq(cfp)) {
         iseq = rb_zjit_cfp_iseq(cfp);
 #define RUBY_VM_IFUNC_P(ptr) IMEMO_TYPE_P(ptr, imemo_ifunc)
         if (RUBY_VM_IFUNC_P(iseq)) {
@@ -132,7 +132,7 @@ control_frame_dump(const rb_execution_context_t *ec, const rb_control_frame_t *c
             line = -1;
         }
         else {
-            if (cfp->pc || CFP_JIT_RETURN(cfp)) {
+            if (rb_zjit_cfp_has_pc(cfp)) {
                 pc = rb_zjit_cfp_pc(cfp) - ISEQ_BODY(iseq)->iseq_encoded;
                 iseq_name = RSTRING_PTR(ISEQ_BODY(iseq)->location.label);
                 if (pc >= 0 && (size_t)pc <= ISEQ_BODY(iseq)->iseq_size) {
@@ -341,7 +341,7 @@ box_env_dump(const rb_execution_context_t *ec, const VALUE *env, const rb_contro
         break;
     }
 
-    if (cfp && (cfp->iseq != 0 || CFP_JIT_RETURN(cfp))) {
+    if (cfp && rb_zjit_cfp_has_iseq(cfp)) {
 #define RUBY_VM_IFUNC_P(ptr) IMEMO_TYPE_P(ptr, imemo_ifunc)
         const rb_iseq_t *resolved_iseq = rb_zjit_cfp_iseq(cfp);
         if (RUBY_VM_IFUNC_P(resolved_iseq)) {
@@ -354,7 +354,7 @@ box_env_dump(const rb_execution_context_t *ec, const VALUE *env, const rb_contro
             line = -1;
         }
         else {
-            if (cfp->pc || CFP_JIT_RETURN(cfp)) {
+            if (rb_zjit_cfp_has_pc(cfp)) {
                 iseq = resolved_iseq;
                 pc = rb_zjit_cfp_pc(cfp) - ISEQ_BODY(iseq)->iseq_encoded;
                 iseq_name = RSTRING_PTR(ISEQ_BODY(iseq)->location.label);

zjit.h

@@ -67,6 +67,26 @@ CFP_JIT_RETURN(const rb_control_frame_t *cfp)
     return !!cfp->jit_return;
 }
 
+// Returns true if cfp has an ISEQ, either directly or via JITFrame.
+// When JITFrame is present, it is authoritative (cfp->iseq may be stale).
+// C frames with JITFrame have iseq=NULL, so this returns false for them.
+static inline bool
+rb_zjit_cfp_has_iseq(const rb_control_frame_t *cfp)
+{
+    if (CFP_JIT_RETURN(cfp)) return rb_zjit_jit_return_iseq(cfp->jit_return) != NULL;
+    return !!cfp->iseq;
+}
+
+// Returns true if cfp has a PC, either directly or via JITFrame.
+// When JITFrame is present, it is authoritative (cfp->pc may be stale/poisoned).
+// C frames with JITFrame have pc=NULL, so this returns false for them.
+static inline bool
+rb_zjit_cfp_has_pc(const rb_control_frame_t *cfp)
+{
+    if (CFP_JIT_RETURN(cfp)) return rb_zjit_jit_return_pc(cfp->jit_return) != NULL;
+    return !!cfp->pc;
+}
+
 static inline const VALUE*
 rb_zjit_cfp_pc(const rb_control_frame_t *cfp)
 {

zjit/src/codegen.rs

@@ -2755,9 +2755,8 @@ fn gen_push_frame(asm: &mut Assembler, argc: usize, state: &FrameState, frame: C
         }
         let new_sp = asm.lea(Opnd::mem(64, SP, (ep_offset + 1) * SIZEOF_VALUE_I32));
         asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SP), new_sp);
-        asm.mov(cfp_opnd(RUBY_OFFSET_CFP_JIT_RETURN), 0.into()); // just clear a leftover on stack. TODO: make JIT frame for it
-        asm.mov(cfp_opnd(RUBY_OFFSET_CFP_ISEQ), 0.into()); // TODO: optimize this with JITFrame
-        asm.mov(cfp_opnd(RUBY_OFFSET_CFP_BLOCK_CODE), 0.into()); // TODO: optimize this with JITFrame
+        let jit_frame = JITFrame::new(std::ptr::null(), std::ptr::null(), true);
+        asm.mov(cfp_opnd(RUBY_OFFSET_CFP_JIT_RETURN), Opnd::const_ptr(jit_frame));
     }
 
     asm.mov(cfp_opnd(RUBY_OFFSET_CFP_SELF), frame.recv);