From 06d875b97901b7c6accb36c8b7b525d0884696a7 Mon Sep 17 00:00:00 2001
From: Takashi Kokubun <takashikkbn@gmail.com>
Date: Fri, 14 Feb 2025 17:27:44 -0800
Subject: [PATCH] Backport the latest jit_compile()

---
 vm.c      | 31 +++++++++++++------------------
 vm_core.h | 10 +++++++++-
 2 files changed, 22 insertions(+), 19 deletions(-)

diff --git a/vm.c b/vm.c
index e2cdfda0cb..735e40376a 100644
--- a/vm.c
+++ b/vm.c
@@ -433,29 +433,25 @@ jit_compile(rb_execution_context_t *ec)
 {
     const rb_iseq_t *iseq = ec->cfp->iseq;
     struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
-    bool yjit_enabled = rb_yjit_enabled_p;
-    extern bool rb_zjit_enabled_p;
-    bool zjit_enabled = rb_zjit_enabled_p;
-    if (!(yjit_enabled || zjit_enabled || rb_rjit_call_p)) {
-        return NULL;
-    }
 
     // Increment the ISEQ's call counter and trigger JIT compilation if not compiled
-    if (body->jit_entry == NULL && rb_yjit_enabled_p) {
+#if USE_ZJIT
+    extern bool rb_zjit_enabled_p;
+    if (body->jit_entry == NULL && rb_zjit_enabled_p) {
         body->jit_entry_calls++;
-        if (zjit_enabled) {
+        if (body->jit_entry_calls == 1) {
             extern void rb_zjit_compile_iseq(const rb_iseq_t *iseq, rb_execution_context_t *ec, bool jit_exception);
             rb_zjit_compile_iseq(iseq, ec, false);
         }
-        else if (yjit_enabled) {
-            if (rb_yjit_threshold_hit(iseq, body->jit_entry_calls)) {
-                rb_yjit_compile_iseq(iseq, ec, false);
-            }
-        }
-        else if (body->jit_entry_calls == rb_rjit_call_threshold()) {
-            rb_rjit_compile(iseq);
+    }
+#elif USE_YJIT
+    if (body->jit_entry == NULL && rb_yjit_enabled_p) {
+        body->jit_entry_calls++;
+        if (rb_yjit_threshold_hit(iseq, body->jit_entry_calls)) {
+            rb_yjit_compile_iseq(iseq, ec, false);
         }
     }
+#endif
     return body->jit_entry;
 }
 
@@ -491,15 +487,14 @@ jit_compile_exception(rb_execution_context_t *ec)
     struct rb_iseq_constant_body *body = ISEQ_BODY(iseq);
 
     // Increment the ISEQ's call counter and trigger JIT compilation if not compiled
+#if USE_YJIT
     if (body->jit_exception == NULL && rb_yjit_enabled_p) {
         body->jit_exception_calls++;
         if (body->jit_exception_calls == rb_yjit_call_threshold) {
             rb_yjit_compile_iseq(iseq, ec, true);
         }
-
-
-
     }
+#endif
     return body->jit_exception;
 }
 
diff --git a/vm_core.h b/vm_core.h
index 79e35cc82c..db98cc01a6 100644
--- a/vm_core.h
+++ b/vm_core.h
@@ -534,7 +534,7 @@ struct rb_iseq_constant_body {
 
     const rb_iseq_t *mandatory_only_iseq;
 
-#if USE_YJIT
+#if USE_YJIT || USE_ZJIT
     // Function pointer for JIT code on jit_exec()
     rb_jit_func_t jit_entry;
     // Number of calls on jit_exec()
@@ -543,11 +543,19 @@ struct rb_iseq_constant_body {
     rb_jit_func_t jit_exception;
     // Number of calls on jit_exec_exception()
     long unsigned jit_exception_calls;
+#endif
+
+#if USE_YJIT
     // YJIT stores some data on each iseq.
     void *yjit_payload;
     // Used to estimate how frequently this ISEQ gets called
     uint64_t yjit_calls_at_interv;
 #endif
+
+#if USE_ZJIT
+    // ZJIT stores some data on each iseq.
+    void *zjit_payload;
+#endif
 };
 
 /* T_IMEMO/iseq */