YJIT: Fix getconstant exits after opt_ltlt fusion (#10903)

Co-authored-by: Alan Wu <[email protected]>

1 files changed, 21 insertions, 13 deletions

diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index d86f0d1955..c6e0d8e023 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -30,7 +30,6 @@ pub use crate::virtualmem::CodePtr;
 /// Status returned by code generation functions
 #[derive(PartialEq, Debug)]
 enum CodegenStatus {
-    SkipNextInsn,
     KeepCompiling,
     EndBlock,
 }
@@ -197,6 +196,13 @@ impl JITState {
         self.insn_idx + insn_len(self.get_opcode()) as u16
     }
 
+    /// Get the index of the next instruction of the next instruction
+    fn next_next_insn_idx(&self) -> u16 {
+        let next_pc = unsafe { rb_iseq_pc_at_idx(self.iseq, self.next_insn_idx().into()) };
+        let next_opcode: usize = unsafe { rb_iseq_opcode_at_pc(self.iseq, next_pc) }.try_into().unwrap();
+        self.next_insn_idx() + insn_len(next_opcode) as u16
+    }
+
     // Check if we are compiling the instruction at the stub PC
     // Meaning we are compiling the instruction that is next to execute
     pub fn at_current_insn(&self) -> bool {
@@ -1098,7 +1104,16 @@ fn jump_to_next_insn(
     jit: &mut JITState,
     asm: &mut Assembler,
     ocb: &mut OutlinedCb,
-) -> Option<()> {
+) -> Option<CodegenStatus> {
+    end_block_with_jump(jit, asm, ocb, jit.next_insn_idx())
+}
+
+fn end_block_with_jump(
+    jit: &mut JITState,
+    asm: &mut Assembler,
+    ocb: &mut OutlinedCb,
+    continuation_insn_idx: u16,
+) -> Option<CodegenStatus> {
     // Reset the depth since in current usages we only ever jump to
     // chain_depth > 0 from the same instruction.
     let mut reset_depth = asm.ctx;
@@ -1106,20 +1121,20 @@ fn jump_to_next_insn(
 
     let jump_block = BlockId {
         iseq: jit.iseq,
-        idx: jit.next_insn_idx(),
+        idx: continuation_insn_idx,
     };
 
     // We are at the end of the current instruction. Record the boundary.
     if jit.record_boundary_patch_point {
         jit.record_boundary_patch_point = false;
-        let exit_pc = unsafe { jit.pc.offset(insn_len(jit.opcode).try_into().unwrap()) };
+        let exit_pc = unsafe { rb_iseq_pc_at_idx(jit.iseq, continuation_insn_idx.into())};
         let exit_pos = gen_outlined_exit(exit_pc, &reset_depth, ocb);
         record_global_inval_patch(asm, exit_pos?);
     }
 
     // Generate the jump instruction
     gen_direct_jump(jit, &reset_depth, jump_block, asm);
-    Some(())
+    Some(EndBlock)
 }
 
 // Compile a sequence of bytecode instructions for a given basic block version.
@@ -1283,13 +1298,6 @@ pub fn gen_single_block(
         // Move to the next instruction to compile
         insn_idx += insn_len(opcode) as u16;
 
-        // Move past next instruction when instructed
-        if status == Some(SkipNextInsn) {
-            let next_pc = unsafe { rb_iseq_pc_at_idx(iseq, insn_idx.into()) };
-            let next_opcode: usize = unsafe { rb_iseq_opcode_at_pc(iseq, next_pc) }.try_into().unwrap();
-            insn_idx += insn_len(next_opcode) as u16;
-        }
-
         // If the instruction terminates this block
         if status == Some(EndBlock) {
             break;
@@ -1519,7 +1527,7 @@ fn fuse_putobject_opt_ltlt(
 
         asm.stack_pop(1);
         fixnum_left_shift_body(asm, lhs, shift_amt as u64);
-        return Some(SkipNextInsn);
+        return end_block_with_jump(jit, asm, ocb, jit.next_next_insn_idx());
     }
     return None;
 }