Reduce the number of branches in jit_exec (#6722)

* Reduce the number of branches in jit_exec

* Address build failure in some configurations

* Refactor yjit.h

1 files changed, 23 insertions, 30 deletions

diff --git a/vm.c b/vm.c
index 87c6b47825..d540aa1c8b 100644
--- a/vm.c
+++ b/vm.c
@@ -378,7 +378,7 @@ extern VALUE rb_vm_invoke_bmethod(rb_execution_context_t *ec, rb_proc_t *proc, V
                                   const rb_callable_method_entry_t *me);
 static VALUE vm_invoke_proc(rb_execution_context_t *ec, rb_proc_t *proc, VALUE self, int argc, const VALUE *argv, int kw_splat, VALUE block_handler);
 
-#if USE_MJIT
+#if USE_MJIT || USE_YJIT
 # ifdef MJIT_HEADER
 NOINLINE(static COLDFUNC VALUE mjit_check_iseq(rb_execution_context_t *ec, const rb_iseq_t *iseq, struct rb_iseq_constant_body *body));
 # else
@@ -412,46 +412,39 @@ mjit_check_iseq(rb_execution_context_t *ec, const rb_iseq_t *iseq, struct rb_ise
 static inline VALUE
 jit_exec(rb_execution_context_t *ec)
 {
+    // Increment the ISEQ's call counter
     const rb_iseq_t *iseq = ec->cfp->iseq;
     struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
-    bool yjit_enabled = false;
-# ifndef MJIT_HEADER
-    // Don't want to compile with YJIT or use code generated by YJIT
-    // when running inside code generated by MJIT.
-    yjit_enabled = rb_yjit_enabled_p();
-# endif
-
-    if (mjit_call_p || yjit_enabled) {
+    bool yjit_enabled = rb_yjit_enabled_p();
+    if (yjit_enabled || mjit_call_p) {
         body->total_calls++;
     }
+    else {
+        return Qundef;
+    }
 
-# ifndef MJIT_HEADER
-    if (yjit_enabled && !mjit_call_p && body->total_calls == rb_yjit_call_threshold())  {
-        // 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)) {
+    // Trigger JIT compilation as needed
+    jit_func_t func;
+    if (yjit_enabled) {
+        if (body->total_calls == rb_yjit_call_threshold())  {
+            // 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;
+            }
+        }
+        // YJIT tried compiling this function once before and couldn't do
+        // it, so return Qundef so the interpreter handles it.
+        if ((func = body->jit_func) == 0) {
             return Qundef;
         }
     }
-# endif
-
-    if (!(mjit_call_p || yjit_enabled))
-        return Qundef;
-
-    jit_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 (yjit_enabled && func == 0) {
-        return Qundef;
-    }
-
-    if (UNLIKELY((uintptr_t)func <= LAST_JIT_ISEQ_FUNC)) {
+    else if (UNLIKELY((uintptr_t)(func = body->jit_func) <= LAST_JIT_ISEQ_FUNC)) {
         return mjit_check_iseq(ec, iseq, body);
     }
 
-    // Under SystemV x64 calling convention: ec -> RDI, cfp -> RSI
-    return func(ec, ec->cfp);
+    // Call the JIT code
+    return func(ec, ec->cfp); // SystemV x64 calling convention: ec -> RDI, cfp -> RSI
 }
 #endif