Always use `ret` to return to the interpreter

Always using `ret` to return to the interpreter means that we never have
to check the VM_FRAME_FLAG_FINISH flag.

In the case that we return `Qundef`, the interpreter will execute the
cfp.  We can take advantage of this by setting the PC to the instruction
we can't handle, and let the interpreter pick up the ball from there.

If we return a value other than Qundef, the interpreter will take that
value as the "return value" from the JIT and push that to the SP of the
caller

The leave instruction puts the return value on the top of the calling
frame's stack.  YJIT does the same thing for leave instructions.
However, when we're returning back to the interpreter, the leave
instruction _should not_ put the return value on the top of the stack,
but put it in RAX and use RET.  This commit pops the last value from the
stack pointer and puts it in RAX so that the interpreter is happy with
SP.

5 files changed, 43 insertions, 39 deletions

diff --git a/mjit.h b/mjit.h
index e6bf6be8dc..1694fcba88 100644
--- a/mjit.h
+++ b/mjit.h
@@ -150,17 +150,27 @@ mjit_exec(rb_execution_context_t *ec)
 
 #ifndef MJIT_HEADER
     if (rb_yjit_enabled_p() && !mjit_call_p && body->total_calls == rb_yjit_call_threshold())  {
-        rb_yjit_compile_iseq(iseq, ec);
-        return Qundef;
+        // If we couldn't generate any code for this iseq, then return
+        // Qundef so the interpreter will handle the call.
+        if (!rb_yjit_compile_iseq(iseq, ec)) {
+          return Qundef;
+        }
     }
 #endif
 
-    if (!mjit_call_p)
+    if (!(mjit_call_p || rb_yjit_enabled_p()))
         return Qundef;
 
     RB_DEBUG_COUNTER_INC(mjit_exec);
 
     mjit_func_t func = body->jit_func;
+
+    // YJIT tried compiling this function once before and couldn't do
+    // it, so return Qundef so the interpreter handles it.
+    if (rb_yjit_enabled_p() && func == 0) {
+      return Qundef;
+    }
+
     if (UNLIKELY((uintptr_t)func <= LAST_JIT_ISEQ_FUNC)) {
 #  ifdef MJIT_HEADER
         RB_DEBUG_COUNTER_INC(mjit_frame_JT2VM);
@@ -176,6 +186,8 @@ mjit_exec(rb_execution_context_t *ec)
     RB_DEBUG_COUNTER_INC(mjit_frame_VM2JT);
 #  endif
     RB_DEBUG_COUNTER_INC(mjit_exec_call_func);
+    // ec -> RDI
+    // cfp -> RSI
     return func(ec, ec->cfp);
 }
 
diff --git a/yjit.h b/yjit.h
index cb41e76ecf..0e7a7b6c9b 100644
--- a/yjit.h
+++ b/yjit.h
@@ -62,7 +62,7 @@ void rb_yjit_cme_invalidate(VALUE cme);
 void rb_yjit_collect_vm_usage_insn(int insn);
 void rb_yjit_collect_binding_alloc(void);
 void rb_yjit_collect_binding_set(void);
-void rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec);
+bool rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec);
 void rb_yjit_init(struct rb_yjit_options *options);
 void rb_yjit_bop_redefined(VALUE klass, const rb_method_entry_t *me, enum ruby_basic_operators bop);
 void rb_yjit_constant_state_changed(void);
diff --git a/yjit_codegen.c b/yjit_codegen.c
index f126407dc8..25d265432d 100644
--- a/yjit_codegen.c
+++ b/yjit_codegen.c
@@ -261,21 +261,6 @@ yjit_gen_exit(jitstate_t *jit, ctx_t *ctx, codeblock_t *cb)
 
     VALUE *exit_pc = jit->pc;
 
-    // YJIT only ever patches the first instruction in an iseq
-    if (jit->insn_idx == 0) {
-        // Table mapping opcodes to interpreter handlers
-        const void *const *handler_table = rb_vm_get_insns_address_table();
-
-        // Write back the old instruction at the exit PC
-        // Otherwise the interpreter may jump right back to the
-        // JITted code we're trying to exit
-        int exit_opcode = yjit_opcode_at_pc(jit->iseq, exit_pc);
-        void* handler_addr = (void*)handler_table[exit_opcode];
-        mov(cb, REG0, const_ptr_opnd(exit_pc));
-        mov(cb, REG1, const_ptr_opnd(handler_addr));
-        mov(cb, mem_opnd(64, REG0, 0), REG1);
-    }
-
     // Generate the code to exit to the interpreters
     // Write the adjusted SP back into the CFP
     if (ctx->sp_offset != 0) {
@@ -299,7 +284,8 @@ yjit_gen_exit(jitstate_t *jit, ctx_t *ctx, codeblock_t *cb)
     }
 #endif
 
-    cb_write_post_call_bytes(cb);
+    mov(cb, RAX, imm_opnd(Qundef));
+    ret(cb);
 
     return code_ptr;
 }
@@ -316,10 +302,13 @@ yjit_gen_leave_exit(codeblock_t *cb)
     // Every exit to the interpreter should be counted
     GEN_COUNTER_INC(cb, leave_interp_return);
 
-    // Put PC into the return register, which the post call bytes dispatches to
-    mov(cb, RAX, member_opnd(REG_CFP, rb_control_frame_t, pc));
+    // GEN_COUNTER_INC clobbers RAX, so put the top of the stack
+    // in to RAX and return.
+    mov(cb, RAX, mem_opnd(64, REG_SP, -SIZEOF_VALUE));
+    sub(cb, member_opnd(REG_CFP, rb_control_frame_t, sp), imm_opnd(SIZEOF_VALUE));
+    mov(cb, REG_SP, member_opnd(REG_CFP, rb_control_frame_t, sp));
 
-    cb_write_post_call_bytes(cb);
+    ret(cb);
 
     return code_ptr;
 }
@@ -348,9 +337,14 @@ yjit_entry_prologue(void)
     cb_align_pos(cb, 64);
 
     uint8_t *code_ptr = cb_get_ptr(cb, cb->write_pos);
+    ADD_COMMENT(cb, "yjit prolog");
 
-    // Write the interpreter entry prologue
-    cb_write_pre_call_bytes(cb);
+    // Fix registers for YJIT.  The MJIT callback puts the ec in RDI
+    // and the CFP in RSI, but REG_CFP == RDI and REG_EC == RSI
+    mov(cb, REG0, RDI); // EC
+    mov(cb, REG1, RSI); // CFP
+    mov(cb, REG_EC, REG0);
+    mov(cb, REG_CFP, REG1);
 
     // Load the current SP from the CFP into REG_SP
     mov(cb, REG_SP, member_opnd(REG_CFP, rb_control_frame_t, sp));
@@ -363,7 +357,8 @@ yjit_entry_prologue(void)
     return code_ptr;
 }
 
-// Generate code to check for interrupts and take a side-exit
+// Generate code to check for interrupts and take a side-exit.
+// Warning: this function clobbers REG0
 static void
 yjit_check_ints(codeblock_t* cb, uint8_t* side_exit)
 {
@@ -3374,15 +3369,10 @@ gen_leave(jitstate_t* jit, ctx_t* ctx)
     uint8_t* side_exit = yjit_side_exit(jit, ctx);
 
     // Load environment pointer EP from CFP
-    mov(cb, REG0, member_opnd(REG_CFP, rb_control_frame_t, ep));
-
-    // if (flags & VM_FRAME_FLAG_FINISH) != 0
-    ADD_COMMENT(cb, "check for finish frame");
-    x86opnd_t flags_opnd = mem_opnd(64, REG0, sizeof(VALUE) * VM_ENV_DATA_INDEX_FLAGS);
-    test(cb, flags_opnd, imm_opnd(VM_FRAME_FLAG_FINISH));
-    jnz_ptr(cb, COUNTED_EXIT(side_exit, leave_se_finish_frame));
+    mov(cb, REG1, member_opnd(REG_CFP, rb_control_frame_t, ep));
 
     // Check for interrupts
+    ADD_COMMENT(cb, "check for interrupts");
     yjit_check_ints(cb, COUNTED_EXIT(side_exit, leave_se_interrupt));
 
     // Load the return value
diff --git a/yjit_iface.c b/yjit_iface.c
index d347cbd5ef..c6fdb7124e 100644
--- a/yjit_iface.c
+++ b/yjit_iface.c
@@ -465,9 +465,10 @@ yjit_block_assumptions_free(block_t *block)
     }
 }
 
-void
+bool
 rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec)
 {
+    bool success = true;
 #if OPT_DIRECT_THREADED_CODE || OPT_CALL_THREADED_CODE
     RB_VM_LOCK_ENTER();
     // TODO: I think we need to stop all other ractors here
@@ -478,14 +479,16 @@ rb_yjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec)
 
     if (code_ptr)
     {
-        // Map the code address to the corresponding opcode
-        int first_opcode = yjit_opcode_at_pc(iseq, &encoded[0]);
-        map_addr2insn(code_ptr, first_opcode);
-        encoded[0] = (VALUE)code_ptr;
+        iseq->body->jit_func = code_ptr;
+    }
+    else {
+        iseq->body->jit_func = 0;
+        success = false;
     }
 
     RB_VM_LOCK_LEAVE();
 #endif
+    return success;
 }
 
 struct yjit_block_itr {
diff --git a/yjit_iface.h b/yjit_iface.h
index 30bee18ed1..4f63cdb974 100644
--- a/yjit_iface.h
+++ b/yjit_iface.h
@@ -46,7 +46,6 @@ YJIT_DECLARE_COUNTERS(
     send_se_cf_overflow,
     send_se_protected_check_failed,
 
-    leave_se_finish_frame,
     leave_se_interrupt,
     leave_interp_return,