2 files changed, 13 insertions, 9 deletions

diff --git a/mjit.c b/mjit.c
index beb59d9646..56e7bb6792 100644
--- a/mjit.c
+++ b/mjit.c
@@ -306,8 +306,8 @@ mjit_add_iseq_to_process(const rb_iseq_t *iseq)
 
 /* Wait for JIT compilation finish for --jit-wait. This should only return a function pointer
    or NOT_COMPILED_JIT_ISEQ_FUNC. */
-mjit_func_t
-mjit_get_iseq_func(struct rb_iseq_constant_body *body)
+VALUE
+mjit_wait_call(rb_execution_context_t *ec, struct rb_iseq_constant_body *body)
 {
     struct timeval tv;
     int tries = 0;
@@ -316,19 +316,23 @@ mjit_get_iseq_func(struct rb_iseq_constant_body *body)
     while (body->jit_func == (mjit_func_t)NOT_READY_JIT_ISEQ_FUNC) {
         tries++;
         if (tries / 1000 > MJIT_WAIT_TIMEOUT_SECONDS || pch_status == PCH_FAILED) {
-            CRITICAL_SECTION_START(3, "in mjit_get_iseq_func to set jit_func");
+            CRITICAL_SECTION_START(3, "in mjit_wait_call to set jit_func");
             body->jit_func = (mjit_func_t)NOT_COMPILED_JIT_ISEQ_FUNC; /* JIT worker seems dead. Give up. */
-            CRITICAL_SECTION_FINISH(3, "in mjit_get_iseq_func to set jit_func");
+            CRITICAL_SECTION_FINISH(3, "in mjit_wait_call to set jit_func");
             mjit_warning("timed out to wait for JIT finish");
             break;
         }
 
-        CRITICAL_SECTION_START(3, "in mjit_get_iseq_func for a client wakeup");
+        CRITICAL_SECTION_START(3, "in mjit_wait_call for a client wakeup");
         rb_native_cond_broadcast(&mjit_worker_wakeup);
-        CRITICAL_SECTION_FINISH(3, "in mjit_get_iseq_func for a client wakeup");
+        CRITICAL_SECTION_FINISH(3, "in mjit_wait_call for a client wakeup");
         rb_thread_wait_for(tv);
     }
-    return body->jit_func;
+
+    if ((uintptr_t)body->jit_func <= (uintptr_t)LAST_JIT_ISEQ_FUNC) {
+        return Qundef;
+    }
+    return body->jit_func(ec, ec->cfp);
 }
 
 extern VALUE ruby_archlibdir_path, ruby_prefix_path;
diff --git a/mjit.h b/mjit.h
index eee3a8badb..0efa04ac8d 100644
--- a/mjit.h
+++ b/mjit.h
@@ -59,7 +59,7 @@ RUBY_EXTERN struct mjit_options mjit_opts;
 RUBY_EXTERN int mjit_call_p;
 
 extern void mjit_add_iseq_to_process(const rb_iseq_t *iseq);
-extern mjit_func_t mjit_get_iseq_func(struct rb_iseq_constant_body *body);
+extern VALUE mjit_wait_call(rb_execution_context_t *ec, struct rb_iseq_constant_body *body);
 RUBY_SYMBOL_EXPORT_END
 
 extern int mjit_compile(FILE *f, const struct rb_iseq_constant_body *body, const char *funcname);
@@ -110,7 +110,7 @@ mjit_exec(rb_execution_context_t *ec)
             if (total_calls == mjit_opts.min_calls && mjit_target_iseq_p(body)) {
                 mjit_add_iseq_to_process(iseq);
                 if (UNLIKELY(mjit_opts.wait)) {
-                    func = mjit_get_iseq_func(body);
+                    return mjit_wait_call(ec, body);
                 }
             }
             return Qundef;