ZJIT: Stop tracking EP == BP assumption on JIT entry (#13752)HEADmaster

* ZJIT: Stop tracking EP == BP assumption on JIT entry

* Enable test_method.rb as well

4 files changed, 30 insertions, 22 deletions

diff --git a/.github/workflows/zjit-macos.yml b/.github/workflows/zjit-macos.yml
index 5454f05b09..8e58605fe1 100644
--- a/.github/workflows/zjit-macos.yml
+++ b/.github/workflows/zjit-macos.yml
@@ -118,27 +118,27 @@ jobs:
           ../src/bootstraptest/test_flip.rb \
           ../src/bootstraptest/test_flow.rb \
           ../src/bootstraptest/test_fork.rb \
+          ../src/bootstraptest/test_gc.rb \
           ../src/bootstraptest/test_io.rb \
           ../src/bootstraptest/test_jump.rb \
           ../src/bootstraptest/test_literal.rb \
           ../src/bootstraptest/test_literal_suffix.rb \
           ../src/bootstraptest/test_load.rb \
           ../src/bootstraptest/test_marshal.rb \
+          ../src/bootstraptest/test_method.rb \
           ../src/bootstraptest/test_objectspace.rb \
           ../src/bootstraptest/test_string.rb \
           ../src/bootstraptest/test_struct.rb \
           ../src/bootstraptest/test_syntax.rb \
+          ../src/bootstraptest/test_thread.rb \
           ../src/bootstraptest/test_yjit_30k_ifelse.rb \
           ../src/bootstraptest/test_yjit_30k_methods.rb \
           ../src/bootstraptest/test_yjit_rust_port.rb
         # ../src/bootstraptest/test_eval.rb \
-        # ../src/bootstraptest/test_gc.rb \
         # ../src/bootstraptest/test_insns.rb \
         # ../src/bootstraptest/test_massign.rb \
-        # ../src/bootstraptest/test_method.rb \
         # ../src/bootstraptest/test_proc.rb \
         # ../src/bootstraptest/test_ractor.rb \
-        # ../src/bootstraptest/test_thread.rb \
         # ../src/bootstraptest/test_yjit.rb \
         if: ${{ matrix.test_task == 'btest' }}
 
diff --git a/.github/workflows/zjit-ubuntu.yml b/.github/workflows/zjit-ubuntu.yml
index aa9402546a..443c9c71df 100644
--- a/.github/workflows/zjit-ubuntu.yml
+++ b/.github/workflows/zjit-ubuntu.yml
@@ -140,6 +140,7 @@ jobs:
           ../src/bootstraptest/test_flip.rb \
           ../src/bootstraptest/test_flow.rb \
           ../src/bootstraptest/test_fork.rb \
+          ../src/bootstraptest/test_gc.rb \
           ../src/bootstraptest/test_io.rb \
           ../src/bootstraptest/test_jump.rb \
           ../src/bootstraptest/test_literal.rb \
@@ -147,20 +148,19 @@ jobs:
           ../src/bootstraptest/test_load.rb \
           ../src/bootstraptest/test_marshal.rb \
           ../src/bootstraptest/test_massign.rb \
+          ../src/bootstraptest/test_method.rb \
           ../src/bootstraptest/test_objectspace.rb \
           ../src/bootstraptest/test_string.rb \
           ../src/bootstraptest/test_struct.rb \
           ../src/bootstraptest/test_syntax.rb \
+          ../src/bootstraptest/test_thread.rb \
           ../src/bootstraptest/test_yjit_30k_ifelse.rb \
           ../src/bootstraptest/test_yjit_30k_methods.rb \
           ../src/bootstraptest/test_yjit_rust_port.rb
         # ../src/bootstraptest/test_eval.rb \
-        # ../src/bootstraptest/test_gc.rb \
         # ../src/bootstraptest/test_insns.rb \
-        # ../src/bootstraptest/test_method.rb \
         # ../src/bootstraptest/test_proc.rb \
         # ../src/bootstraptest/test_ractor.rb \
-        # ../src/bootstraptest/test_thread.rb \
         # ../src/bootstraptest/test_yjit.rb \
         if: ${{ matrix.test_task == 'btest' }}
 
diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs
index d7d3cb8aca..33a8af6868 100644
--- a/zjit/src/codegen.rs
+++ b/zjit/src/codegen.rs
@@ -6,7 +6,6 @@ use crate::backend::current::{Reg, ALLOC_REGS};
 use crate::profile::get_or_create_iseq_payload;
 use crate::state::ZJITState;
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
-use crate::invariants::{iseq_escapes_ep, track_no_ep_escape_assumption};
 use crate::backend::lir::{self, asm_comment, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, SP};
 use crate::hir::{iseq_to_hir, Block, BlockId, BranchEdge, CallInfo, RangeType, SELF_PARAM_IDX, SpecialObjectType};
 use crate::hir::{Const, FrameState, Function, Insn, InsnId};
@@ -59,15 +58,6 @@ impl JITState {
             }
         }
     }
-
-    /// Assume that this ISEQ doesn't escape EP. Return false if it's known to escape EP.
-    fn assume_no_ep_escape(iseq: IseqPtr) -> bool {
-        if iseq_escapes_ep(iseq) {
-            return false;
-        }
-        track_no_ep_escape_assumption(iseq);
-        true
-    }
 }
 
 /// CRuby API to compile a given ISEQ
@@ -145,7 +135,7 @@ fn gen_entry(cb: &mut CodeBlock, iseq: IseqPtr, function: &Function, function_pt
     // Set up registers for CFP, EC, SP, and basic block arguments
     let mut asm = Assembler::new();
     gen_entry_prologue(&mut asm, iseq);
-    gen_method_params(&mut asm, iseq, function.block(BlockId(0)));
+    gen_entry_params(&mut asm, iseq, function.block(BlockId(0)));
 
     // Jump to the first block using a call instruction
     asm.ccall(function_ptr.raw_ptr(cb) as *const u8, vec![]);
@@ -501,7 +491,7 @@ fn gen_entry_prologue(asm: &mut Assembler, iseq: IseqPtr) {
 }
 
 /// Assign method arguments to basic block arguments at JIT entry
-fn gen_method_params(asm: &mut Assembler, iseq: IseqPtr, entry_block: &Block) {
+fn gen_entry_params(asm: &mut Assembler, iseq: IseqPtr, entry_block: &Block) {
     let self_param = gen_param(asm, SELF_PARAM_IDX);
     asm.mov(self_param, Opnd::mem(VALUE_BITS, CFP, RUBY_OFFSET_CFP_SELF));
 
@@ -516,7 +506,7 @@ fn gen_method_params(asm: &mut Assembler, iseq: IseqPtr, entry_block: &Block) {
 
         // Assign local variables to the basic block arguments
         for (idx, &param) in params.iter().enumerate() {
-            let local = gen_getlocal(asm, iseq, idx);
+            let local = gen_entry_param(asm, iseq, idx);
             asm.load_into(param, local);
         }
     }
@@ -535,11 +525,13 @@ fn gen_branch_params(jit: &mut JITState, asm: &mut Assembler, branch: &BranchEdg
     Some(())
 }
 
-/// Get the local variable at the given index
-fn gen_getlocal(asm: &mut Assembler, iseq: IseqPtr, local_idx: usize) -> lir::Opnd {
+/// Get a method parameter on JIT entry. As of entry, whether EP is escaped or not solely
+/// depends on the ISEQ type.
+fn gen_entry_param(asm: &mut Assembler, iseq: IseqPtr, local_idx: usize) -> lir::Opnd {
     let ep_offset = local_idx_to_ep_offset(iseq, local_idx);
 
-    if JITState::assume_no_ep_escape(iseq) {
+    // If the ISEQ does not escape EP, we can optimize the local variable access using the SP register.
+    if !iseq_entry_escapes_ep(iseq) {
         // Create a reference to the local variable using the SP register. We assume EP == BP.
         // TODO: Implement the invalidation in rb_zjit_invalidate_ep_is_bp()
         let offs = -(SIZEOF_VALUE_I32 * (ep_offset + 1));
@@ -1057,6 +1049,20 @@ fn side_exit(jit: &mut JITState, state: &FrameState) -> Option<Target> {
     Some(target)
 }
 
+/// Return true if a given ISEQ is known to escape EP to the heap on entry.
+///
+/// As of vm_push_frame(), EP is always equal to BP. However, after pushing
+/// a frame, some ISEQ setups call vm_bind_update_env(), which redirects EP.
+fn iseq_entry_escapes_ep(iseq: IseqPtr) -> bool {
+    match unsafe { get_iseq_body_type(iseq) } {
+        // <main> frame is always associated to TOPLEVEL_BINDING.
+        ISEQ_TYPE_MAIN |
+        // Kernel#eval uses a heap EP when a Binding argument is not nil.
+        ISEQ_TYPE_EVAL => true,
+        _ => false,
+    }
+}
+
 impl Assembler {
     /// Make a C call while marking the start and end positions of it
     fn ccall_with_branch(&mut self, fptr: *const u8, opnds: Vec<Opnd>, branch: &Rc<Branch>) -> Opnd {
diff --git a/zjit/src/invariants.rs b/zjit/src/invariants.rs
index 77fd78d95e..77ccc7d04c 100644
--- a/zjit/src/invariants.rs
+++ b/zjit/src/invariants.rs
@@ -40,6 +40,8 @@ pub extern "C" fn rb_zjit_invalidate_ep_is_bp(iseq: IseqPtr) {
     invariants.ep_escape_iseqs.insert(iseq);
 
     // If the ISEQ has been compiled assuming it doesn't escape EP, invalidate the JIT code.
+    // Note: Nobody calls track_no_ep_escape_assumption() for now, so this is always false.
+    // TODO: Add a PatchPoint that assumes EP == BP in HIR and invalidate it here.
     if invariants.no_ep_escape_iseqs.contains(&iseq) {
         unimplemented!("Invalidation on EP escape is not implemented yet");
     }