ZJIT: Update GetBlockHandler to use GetLEP

Author: nirvdrum

zjit/src/codegen.rs

@@ -455,7 +455,7 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio
         &Insn::GuardLess { left, right, state } => gen_guard_less(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state)),
         &Insn::GuardGreaterEq { left, right, state } => gen_guard_greater_eq(jit, asm, opnd!(left), opnd!(right), &function.frame_state(state)),
         &Insn::GuardSuperMethodEntry { lep, cme, state } => no_output!(gen_guard_super_method_entry(jit, asm, opnd!(lep), cme, &function.frame_state(state))),
-        Insn::GetBlockHandler => gen_get_block_handler(jit, asm),
+        Insn::GetBlockHandler { lep } => gen_get_block_handler(asm, opnd!(lep)),
         Insn::PatchPoint { invariant, state } => no_output!(gen_patch_point(jit, asm, invariant, &function.frame_state(*state))),
         Insn::CCall { cfunc, recv, args, name, return_type: _, elidable: _ } => gen_ccall(asm, *cfunc, *name, opnd!(recv), opnds!(args)),
         // Give up CCallWithFrame for 7+ args since asm.ccall() supports at most 6 args (recv + args).
@@ -736,9 +736,8 @@ fn gen_guard_super_method_entry(
 }
 
 /// Get the block handler from ep[VM_ENV_DATA_INDEX_SPECVAL] at the local EP (LEP).
-fn gen_get_block_handler(jit: &JITState, asm: &mut Assembler) -> Opnd {
+fn gen_get_block_handler(asm: &mut Assembler, lep: Opnd) -> Opnd {
     asm_comment!(asm, "get block handler from LEP");
-    let lep = gen_get_lep(jit, asm);
     asm.load(Opnd::mem(64, lep, SIZEOF_VALUE_I32 * VM_ENV_DATA_INDEX_SPECVAL))
 }
 

zjit/src/hir.rs

@@ -1018,7 +1018,7 @@ pub enum Insn {
     /// Used to ensure super calls are made from the expected method context.
     GuardSuperMethodEntry { lep: InsnId, cme: *const rb_callable_method_entry_t, state: InsnId },
     /// Get the block handler from ep[VM_ENV_DATA_INDEX_SPECVAL] at the local EP (LEP).
-    GetBlockHandler,
+    GetBlockHandler { lep: InsnId },
 
     /// Generate no code (or padding if necessary) and insert a patch point
     /// that can be rewritten to a side exit when the Invariant is broken.
@@ -1516,7 +1516,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> {
             Insn::GuardLess { left, right, .. } => write!(f, "GuardLess {left}, {right}"),
             Insn::GuardGreaterEq { left, right, .. } => write!(f, "GuardGreaterEq {left}, {right}"),
             Insn::GuardSuperMethodEntry { lep, cme, .. } => write!(f, "GuardSuperMethodEntry {lep}, {:p}", self.ptr_map.map_ptr(cme)),
-            Insn::GetBlockHandler => write!(f, "GetBlockHandler"),
+            Insn::GetBlockHandler { lep } => write!(f, "GetBlockHandler {lep}"),
             Insn::PatchPoint { invariant, .. } => { write!(f, "PatchPoint {}", invariant.print(self.ptr_map)) },
             Insn::GetConstantPath { ic, .. } => { write!(f, "GetConstantPath {:p}", self.ptr_map.map_ptr(ic)) },
             Insn::IsBlockGiven => { write!(f, "IsBlockGiven") },
@@ -2181,7 +2181,7 @@ impl Function {
             &GuardGreaterEq { left, right, state } => GuardGreaterEq { left: find!(left), right: find!(right), state },
             &GuardLess { left, right, state } => GuardLess { left: find!(left), right: find!(right), state },
             &GuardSuperMethodEntry { lep, cme, state } => GuardSuperMethodEntry { lep: find!(lep), cme, state },
-            &GetBlockHandler => GetBlockHandler,
+            &GetBlockHandler { lep } => GetBlockHandler { lep: find!(lep) },
             &FixnumAdd { left, right, state } => FixnumAdd { left: find!(left), right: find!(right), state },
             &FixnumSub { left, right, state } => FixnumSub { left: find!(left), right: find!(right), state },
             &FixnumMult { left, right, state } => FixnumMult { left: find!(left), right: find!(right), state },
@@ -2476,7 +2476,7 @@ impl Function {
             Insn::AnyToString { .. } => types::String,
             Insn::GetLocal { rest_param: true, .. } => types::ArrayExact,
             Insn::GetLocal { .. } => types::BasicObject,
-            Insn::GetBlockHandler => types::RubyValue,
+            Insn::GetBlockHandler { .. } => types::RubyValue,
             // The type of Snapshot doesn't really matter; it's never materialized. It's used only
             // as a reference for FrameState, which we use to generate side-exit code.
             Insn::Snapshot { .. } => types::Any,
@@ -3429,7 +3429,7 @@ impl Function {
                         });
 
                         // Guard that no block is being passed (implicit or explicit).
-                        let block_handler = self.push_insn(block, Insn::GetBlockHandler);
+                        let block_handler = self.push_insn(block, Insn::GetBlockHandler { lep });
                         self.push_insn(block, Insn::GuardBitEquals {
                             val: block_handler,
                             expected: Const::Value(VALUE(VM_BLOCK_HANDLER_NONE as usize)),
@@ -4373,12 +4373,14 @@ impl Function {
             | &Insn::GetLEP
             | &Insn::LoadSelf
             | &Insn::GetLocal { .. }
-            | &Insn::GetBlockHandler
             | &Insn::PutSpecialObject { .. }
             | &Insn::IsBlockGiven
             | &Insn::IncrCounter(_)
             | &Insn::IncrCounterPtr { .. } =>
                 {}
+            &Insn::GetBlockHandler { lep } => {
+                worklist.push_back(lep);
+            }
             &Insn::PatchPoint { state, .. }
             | &Insn::CheckInterrupts { state }
             | &Insn::GetConstantPath { ic: _, state } => {
@@ -5123,7 +5125,7 @@ impl Function {
             | Insn::EntryPoint { .. }
             | Insn::GuardBlockParamProxy { .. }
             | Insn::GuardSuperMethodEntry { .. }
-            | Insn::GetBlockHandler
+            | Insn::GetBlockHandler { .. }
             | Insn::PatchPoint { .. }
             | Insn::SideExit { .. }
             | Insn::IncrCounter { .. }

zjit/src/hir/opt_tests.rs

@@ -10945,7 +10945,7 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(A@0x1000, foo@0x1008, cme:0x1010)
           v17:CPtr = GetLEP
           GuardSuperMethodEntry v17, 0x1038
-          v19:RubyValue = GetBlockHandler
+          v19:RubyValue = GetBlockHandler v17
           v20:FalseClass = GuardBitEquals v19, Value(false)
           v21:BasicObject = SendWithoutBlockDirect v6, :foo (0x1040)
           CheckInterrupts
@@ -10989,7 +10989,7 @@ mod hir_opt_tests {
           PatchPoint MethodRedefined(A@0x1000, foo@0x1008, cme:0x1010)
           v26:CPtr = GetLEP
           GuardSuperMethodEntry v26, 0x1038
-          v28:RubyValue = GetBlockHandler
+          v28:RubyValue = GetBlockHandler v26
           v29:FalseClass = GuardBitEquals v28, Value(false)
           v30:BasicObject = SendWithoutBlockDirect v8, :foo (0x1040), v9
           v17:Fixnum[1] = Const Value(1)