ZJIT: Remove `JITState#iseq` (#16774)

While working on implementing a new inliner for ZJIT, I keep running into an issue where the wrong iseq pointer is being used. Since `JITState#iseq` and `FrameState#iseq` end up being the same value in the absence of inlining, I don't think we've been all that rigorous about using the "correct" one in a given context. I've been migrating to `FrameState#iseq` in my inliner work and adapting on each rebase. But, having looked at this code for a little while, I think we can remove `JITState#iseq`.

This PR switches to using `FrameState#iseq` in cases where we're working with the code associated with a specific frame (e.g., local-table layout, side-exit resumption target) rather than the overall compilation unit. Other cases are cleaned up by either storing the iseq on the instruction or computing the value that required the iseq and storing that in the instruction.

Without inlining, this change set is really just a clean-up of what we have today. With inlining, however, these changes become semantically meaningful. Making the change now removes the decision point of when to use `JITState#iseq` vs `FrameState#iseq` and simplifies keeping the inliner work current with _master_.


* ZJIT: Use `FrameState#iseq` for variable-access codegen

* ZJIT: Use `FrameState#iseq` for `IsMethodCfunc` codegen

* ZJIT: Precompute LEP level for defined?(yield) at HIR construction

Rather than use `JITState#iseq` to compute the LEP level during codegen, we can compute it once and store it at HIR construction.

* ZJIT: Use `FrameState#iseq` for local spills and side-exit reconstruction

Previously, we were using `JITState#iseq`, but these call sites are more about the frame state rather than the outer compilation unit.

* ZJIT: Remove the unused iseq field from JITState

3 files changed, 35 insertions, 40 deletions

diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index 72068c1e5b..4aa0208252 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -52,9 +52,6 @@ const JIT_RETURN_POISON: Option<usize> = if cfg!(feature = "runtime_checks") {
 
 /// Ephemeral code generation state
 struct JITState {
-    /// Instruction sequence for the method being compiled
-    iseq: IseqPtr,
-
     /// ISEQ version that is being compiled, which will be used by PatchPoint
     version: IseqVersionRef,
 
@@ -73,9 +70,8 @@ struct JITState {
 
 impl JITState {
     /// Create a new JITState instance
-    fn new(iseq: IseqPtr, version: IseqVersionRef, num_insns: usize, num_blocks: usize) -> Self {
+    fn new(version: IseqVersionRef, num_insns: usize, num_blocks: usize) -> Self {
         JITState {
-            iseq,
             version,
             opnds: vec![None; num_insns],
             labels: vec![None; num_blocks],
@@ -379,7 +375,7 @@ fn gen_iseq_body(cb: &mut CodeBlock, iseq: IseqPtr, mut version: IseqVersionRef,
 fn gen_function(cb: &mut CodeBlock, iseq: IseqPtr, version: IseqVersionRef, function: &Function) -> Result<(IseqCodePtrs, Vec<CodePtr>, Vec<IseqCallRef>), CompileError> {
     let (mut jit, asm) = trace_compile_phase("codegen", || {
         let num_spilled_params = max_num_params(function).saturating_sub(ALLOC_REGS.len());
-        let mut jit = JITState::new(iseq, version, function.num_insns(), function.num_blocks());
+        let mut jit = JITState::new(version, function.num_insns(), function.num_blocks());
         let mut asm = Assembler::new_with_stack_slots(num_spilled_params);
 
         // Mapping from HIR block IDs to LIR block IDs.
@@ -706,7 +702,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::FixnumMod { left, right, state } => gen_fixnum_mod(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state)),
         &Insn::FixnumAref { recv, index } => gen_fixnum_aref(asm, opnd!(recv), opnd!(index)),
         Insn::IsNil { val } => gen_isnil(asm, opnd!(val)),
-        &Insn::IsMethodCfunc { val, cd, cfunc, state: _ } => gen_is_method_cfunc(jit, asm, opnd!(val), cd, cfunc),
+        &Insn::IsMethodCfunc { val, cd, cfunc, state } => gen_is_method_cfunc(asm, opnd!(val), cd, cfunc, &function.frame_state(state)),
         &Insn::IsBitEqual { left, right } => gen_is_bit_equal(asm, opnd!(left), opnd!(right)),
         &Insn::IsBitNotEqual { left, right } => gen_is_bit_not_equal(asm, opnd!(left), opnd!(right)),
         &Insn::BoxBool { val } => gen_box_bool(asm, opnd!(val)),
@@ -733,7 +729,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::CCallVariadic { cfunc, recv, name, args, cme, state, block, return_type: _, elidable: _ } => {
             gen_ccall_variadic(jit, asm, *cfunc, *name, opnd!(recv), opnds!(args), *cme, *block, &function.frame_state(*state))
         }
-        Insn::GetIvar { self_val, id, ic, state: _ } => gen_getivar(jit, asm, opnd!(self_val), *id, *ic),
+        Insn::GetIvar { self_val, id, ic, state } => gen_getivar(asm, opnd!(self_val), *id, *ic, &function.frame_state(*state)),
         Insn::SetGlobal { id, val, state } => no_output!(gen_setglobal(jit, asm, *id, opnd!(val), &function.frame_state(*state))),
         Insn::GetGlobal { id, state } => gen_getglobal(jit, asm, *id, &function.frame_state(*state)),
         &Insn::IsBlockParamModified { flags } => gen_is_block_param_modified(asm, opnd!(flags)),
@@ -748,7 +744,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         Insn::SideExit { state, reason, recompile } => no_output!(gen_side_exit(jit, asm, reason, *recompile, &function.frame_state(*state))),
         Insn::PutSpecialObject { value_type } => gen_putspecialobject(asm, *value_type),
         Insn::AnyToString { val, str, state } => gen_anytostring(asm, opnd!(val), opnd!(str), &function.frame_state(*state)),
-        Insn::Defined { op_type, obj, pushval, v, state } => gen_defined(jit, asm, *op_type, *obj, *pushval, opnd!(v), &function.frame_state(*state)),
+        Insn::Defined { op_type, obj, pushval, v, lep_level, state } => gen_defined(jit, asm, *op_type, *obj, *pushval, opnd!(v), *lep_level, &function.frame_state(*state)),
         Insn::CheckMatch { target, pattern, flag, state } => gen_checkmatch(jit, asm, opnd!(target), opnd!(pattern), *flag, &function.frame_state(*state)),
         Insn::GetSpecialSymbol { symbol_type, state: _ } => gen_getspecial_symbol(asm, *symbol_type),
         Insn::GetSpecialNumber { nth, state } => gen_getspecial_number(asm, *nth, &function.frame_state(*state)),
@@ -794,13 +790,6 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
     Ok(())
 }
 
-/// Gets the EP of the ISeq of the containing method, or "local level".
-/// Equivalent of GET_LEP() macro.
-fn gen_get_lep(jit: &JITState, asm: &mut Assembler) -> Opnd {
-    let level = get_lvar_level(jit.iseq);
-    gen_get_ep(asm, level)
-}
-
 // Get EP at `level` from CFP
 fn gen_get_ep(asm: &mut Assembler, level: u32) -> Opnd {
     // Load environment pointer EP from CFP into a register
@@ -835,18 +824,12 @@ fn gen_objtostring(jit: &mut JITState, asm: &mut Assembler, val: Opnd, cd: *cons
     ret
 }
 
-fn gen_defined(jit: &JITState, asm: &mut Assembler, op_type: usize, obj: VALUE, pushval: VALUE, tested_value: Opnd, state: &FrameState) -> Opnd {
+fn gen_defined(jit: &JITState, asm: &mut Assembler, op_type: usize, obj: VALUE, pushval: VALUE, tested_value: Opnd, lep_level: u32, state: &FrameState) -> Opnd {
     match op_type as defined_type {
         DEFINED_YIELD => {
-            // `yield` goes to the block handler stowed in the "local" iseq which is
-            // the current iseq or a parent. Only the "method" iseq type can be passed a
-            // block handler. (e.g. `yield` in the top level script is a syntax error.)
-            //
-            // Similar to gen_is_block_given
-            let local_iseq = unsafe { rb_get_iseq_body_local_iseq(jit.iseq) };
-            assert_eq!(unsafe { rb_get_iseq_body_type(local_iseq) }, ISEQ_TYPE_METHOD,
-                       "defined?(yield) in non-method iseq should be handled by HIR construction");
-            let lep = gen_get_lep(jit, asm);
+            // `lep_level` was precomputed at HIR construction so we can materialize the local EP
+            // inline without walking the parent iseq chain here.
+            let lep = gen_get_ep(asm, lep_level);
             let block_handler = asm.load(Opnd::mem(64, lep, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL));
             let pushval = asm.load(pushval.into());
             asm.cmp(block_handler, VM_BLOCK_HANDLER_NONE.into());
@@ -1210,11 +1193,11 @@ fn gen_ccall_variadic(
 }
 
 /// Emit an uncached instance variable lookup
-fn gen_getivar(jit: &mut JITState, asm: &mut Assembler, recv: Opnd, id: ID, ic: *const iseq_inline_iv_cache_entry) -> Opnd {
+fn gen_getivar(asm: &mut Assembler, recv: Opnd, id: ID, ic: *const iseq_inline_iv_cache_entry, state: &FrameState) -> Opnd {
     if ic.is_null() {
         asm_ccall!(asm, rb_ivar_get, recv, id.0.into())
     } else {
-        let iseq = Opnd::Value(jit.iseq.into());
+        let iseq = Opnd::Value(state.iseq.into());
         asm_ccall!(asm, rb_vm_getinstancevariable, iseq, recv, id.0.into(), Opnd::const_ptr(ic))
     }
 }
@@ -1226,19 +1209,19 @@ fn gen_setivar(jit: &mut JITState, asm: &mut Assembler, recv: Opnd, id: ID, ic:
     if ic.is_null() {
         asm_ccall!(asm, rb_ivar_set, recv, id.0.into(), val);
     } else {
-        let iseq = Opnd::Value(jit.iseq.into());
+        let iseq = Opnd::Value(state.iseq.into());
         asm_ccall!(asm, rb_vm_setinstancevariable, iseq, recv, id.0.into(), val, Opnd::const_ptr(ic));
     }
 }
 
 fn gen_getclassvar(jit: &mut JITState, asm: &mut Assembler, id: ID, ic: *const iseq_inline_cvar_cache_entry, state: &FrameState) -> Opnd {
     gen_prepare_non_leaf_call(jit, asm, state);
-    asm_ccall!(asm, rb_vm_getclassvariable, VALUE::from(jit.iseq).into(), CFP, id.0.into(), Opnd::const_ptr(ic))
+    asm_ccall!(asm, rb_vm_getclassvariable, VALUE::from(state.iseq).into(), CFP, id.0.into(), Opnd::const_ptr(ic))
 }
 
 fn gen_setclassvar(jit: &mut JITState, asm: &mut Assembler, id: ID, val: Opnd, ic: *const iseq_inline_cvar_cache_entry, state: &FrameState) {
     gen_prepare_non_leaf_call(jit, asm, state);
-    asm_ccall!(asm, rb_vm_setclassvariable, VALUE::from(jit.iseq).into(), CFP, id.0.into(), val, Opnd::const_ptr(ic));
+    asm_ccall!(asm, rb_vm_setclassvariable, VALUE::from(state.iseq).into(), CFP, id.0.into(), val, Opnd::const_ptr(ic));
 }
 
 /// Look up global variables
@@ -2435,11 +2418,11 @@ fn gen_isnil(asm: &mut Assembler, val: lir::Opnd) -> lir::Opnd {
     asm.csel_e(Opnd::Imm(1), Opnd::Imm(0))
 }
 
-fn gen_is_method_cfunc(jit: &JITState, asm: &mut Assembler, val: lir::Opnd, cd: *const rb_call_data, cfunc: *const u8) -> lir::Opnd {
+fn gen_is_method_cfunc(asm: &mut Assembler, val: lir::Opnd, cd: *const rb_call_data, cfunc: *const u8, state: &FrameState) -> lir::Opnd {
     unsafe extern "C" {
         fn rb_vm_method_cfunc_is(iseq: IseqPtr, cd: *const rb_call_data, recv: VALUE, cfunc: *const u8) -> VALUE;
     }
-    asm_ccall!(asm, rb_vm_method_cfunc_is, VALUE::from(jit.iseq).into(), Opnd::const_ptr(cd), val, Opnd::const_ptr(cfunc))
+    asm_ccall!(asm, rb_vm_method_cfunc_is, VALUE::from(state.iseq).into(), Opnd::const_ptr(cd), val, Opnd::const_ptr(cfunc))
 }
 
 fn gen_is_bit_equal(asm: &mut Assembler, left: lir::Opnd, right: lir::Opnd) -> lir::Opnd {
@@ -2867,7 +2850,7 @@ fn gen_spill_locals(jit: &JITState, asm: &mut Assembler, state: &FrameState) {
     gen_incr_counter(asm, Counter::vm_write_locals_count);
     asm_comment!(asm, "spill locals");
     for (idx, &insn_id) in state.locals().enumerate() {
-        asm.mov(Opnd::mem(64, SP, (-local_idx_to_ep_offset(jit.iseq, idx) - 1) * SIZEOF_VALUE_I32), jit.get_opnd(insn_id));
+        asm.mov(Opnd::mem(64, SP, (-local_idx_to_ep_offset(state.iseq, idx) - 1) * SIZEOF_VALUE_I32), jit.get_opnd(insn_id));
     }
 }
 
@@ -3041,7 +3024,7 @@ fn side_exit(jit: &JITState, state: &FrameState, reason: SideExitReason) -> Targ
 fn side_exit_with_recompile(jit: &JITState, state: &FrameState, reason: SideExitReason, recompile: Option<Recompile>) -> Target {
     let mut exit = build_side_exit(jit, state);
     exit.recompile = recompile.map(|strategy| SideExitRecompile {
-        iseq: Opnd::Value(VALUE::from(jit.iseq)),
+        iseq: Opnd::Value(VALUE::from(state.iseq)),
         insn_idx: state.insn_idx() as u32,
         strategy,
     });
@@ -3064,7 +3047,7 @@ fn build_side_exit(jit: &JITState, state: &FrameState) -> SideExit {
         pc: Opnd::const_ptr(state.pc),
         stack,
         locals,
-        iseq: jit.iseq,
+        iseq: state.iseq,
         recompile: None,
     }
 }
diff --git a/zjit/src/codegen_tests.rs b/zjit/src/codegen_tests.rs
index b2cc401694..ac5f80358d 100644
--- a/zjit/src/codegen_tests.rs
+++ b/zjit/src/codegen_tests.rs
@@ -23,7 +23,6 @@ fn test_breakpoint_hir_codegen() {
     let breakpoint = function.push_insn(function.entries_block, Insn::BreakPoint);
 
     let mut jit = JITState::new(
-        iseq,
         IseqVersion::new(iseq),
         function.num_insns(),
         function.num_blocks(),
diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs
index 935072af40..5314b07c20 100644
--- a/zjit/src/hir.rs
+++ b/zjit/src/hir.rs
@@ -883,7 +883,9 @@ pub enum Insn {
     UnboxFixnum { val: InsnId },
     FixnumAref { recv: InsnId, index: InsnId },
     // TODO(max): In iseq body types that are not ISEQ_TYPE_METHOD, rewrite to Constant false.
-    Defined { op_type: usize, obj: VALUE, pushval: VALUE, v: InsnId, state: InsnId },
+    // `lep_level` is the lexical distance from this insn's iseq up to its local_iseq, used only
+    // for the DEFINED_YIELD op_type to materialize the local EP inline. Zero for other op_types.
+    Defined { op_type: usize, obj: VALUE, pushval: VALUE, v: InsnId, lep_level: u32, state: InsnId },
     GetConstant { klass: InsnId, id: ID, allow_nil: InsnId, state: InsnId },
     GetConstantPath { ic: *const iseq_inline_constant_cache, state: InsnId },
     /// Kernel#block_given? but without pushing a frame. Similar to [`Insn::Defined`] with
@@ -2982,7 +2984,7 @@ impl Function {
                 cfunc, recv: find!(recv), args: find_vec!(args), cme, name, state, return_type, elidable, block
             },
             &CheckMatch { target, pattern, flag, state } => CheckMatch { target: find!(target), pattern: find!(pattern), flag, state: find!(state) },
-            &Defined { op_type, obj, pushval, v, state } => Defined { op_type, obj, pushval, v: find!(v), state: find!(state) },
+            &Defined { op_type, obj, pushval, v, lep_level, state } => Defined { op_type, obj, pushval, v: find!(v), lep_level, state: find!(state) },
             &DefinedIvar { self_val, pushval, id, state } => DefinedIvar { self_val: find!(self_val), pushval, id, state },
             &GetConstant { klass, id, allow_nil, state } => GetConstant { klass: find!(klass), id, allow_nil: find!(allow_nil), state },
             &NewArray { ref elements, state } => NewArray { elements: find_vec!(elements), state: find!(state) },
@@ -7174,7 +7176,18 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         // Similar to gen_is_block_given
                         Insn::Const { val: Const::Value(Qnil) }
                     } else {
-                        Insn::Defined { op_type, obj, pushval, v, state: exit_id }
+                        // For DEFINED_YIELD, codegen materializes the local EP inline (similar to
+                        // gen_is_block_given) to check for a block handler. Precompute the lexical
+                        // distance from this iseq up to local_iseq so codegen does not have to
+                        // walk the parent chain. Any DEFINED_YIELD reaching this branch has a
+                        // method local_iseq by construction — the above branch has already
+                        // diverted the non-method case to Qnil.
+                        let lep_level = if op_type == DEFINED_YIELD as usize {
+                            get_lvar_level(iseq)
+                        } else {
+                            0
+                        };
+                        Insn::Defined { op_type, obj, pushval, v, lep_level, state: exit_id }
                     };
                     state.stack_push(fun.push_insn(block, insn));
                 }