YJIT: refactoring to allow for fancier call threshold logic (#8078)

* YJIT: refactoring to allow for fancier call threshold logic

* Avoid potentially compiling functions multiple times.

* Update vm.c

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

---------

Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

6 files changed, 23 insertions, 6 deletions

diff --git a/vm.c b/vm.c
index 88aba3433a..f8cdc5a5ea 100644
--- a/vm.c
+++ b/vm.c
@@ -385,9 +385,14 @@ jit_compile(rb_execution_context_t *ec)
         return 0;
     }
 
+    // Don't try to compile the function if it's already compiled
+    if (body->jit_func) {
+        return body->jit_func;
+    }
+
     // Trigger JIT compilation as needed
     if (yjit_enabled) {
-        if (body->total_calls == rb_yjit_call_threshold())  {
+        if (rb_yjit_threshold_hit(iseq)) {
             rb_yjit_compile_iseq(iseq, ec);
         }
     }
diff --git a/yjit.c b/yjit.c
index 8e0c43f1bd..4db46d59f3 100644
--- a/yjit.c
+++ b/yjit.c
@@ -590,6 +590,12 @@ rb_get_def_bmethod_proc(rb_method_definition_t *def)
     return def->body.bmethod.proc;
 }
 
+unsigned long
+rb_get_iseq_body_total_calls(const rb_iseq_t *iseq)
+{
+    return iseq->body->total_calls;
+}
+
 const rb_iseq_t *
 rb_get_iseq_body_local_iseq(const rb_iseq_t *iseq)
 {
diff --git a/yjit.h b/yjit.h
index a76dc6a850..37a3fb0d49 100644
--- a/yjit.h
+++ b/yjit.h
@@ -27,7 +27,7 @@
 // Expose these as declarations since we are building YJIT.
 bool rb_yjit_enabled_p(void);
 bool rb_yjit_compile_new_iseqs(void);
-unsigned rb_yjit_call_threshold(void);
+bool rb_yjit_threshold_hit(const rb_iseq_t *const iseq);
 void rb_yjit_invalidate_all_method_lookup_assumptions(void);
 void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme);
 void rb_yjit_collect_binding_alloc(void);
@@ -49,7 +49,7 @@ void rb_yjit_tracing_invalidate_all(void);
 
 static inline bool rb_yjit_enabled_p(void) { return false; }
 static inline bool rb_yjit_compile_new_iseqs(void) { return false; }
-static inline unsigned rb_yjit_call_threshold(void) { return UINT_MAX; }
+static inline bool rb_yjit_threshold_hit(const rb_iseq_t *const iseq) { return false; }
 static inline void rb_yjit_invalidate_all_method_lookup_assumptions(void) {}
 static inline void rb_yjit_cme_invalidate(rb_callable_method_entry_t *cme) {}
 static inline void rb_yjit_collect_binding_alloc(void) {}
diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index 30b4c8b5d7..d00816b3d5 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -394,6 +394,7 @@ fn main() {
         .allowlist_function("rb_get_def_iseq_ptr")
         .allowlist_function("rb_get_def_bmethod_proc")
         .allowlist_function("rb_iseq_encoded_size")
+        .allowlist_function("rb_get_iseq_body_total_calls")
         .allowlist_function("rb_get_iseq_body_local_iseq")
         .allowlist_function("rb_get_iseq_body_parent_iseq")
         .allowlist_function("rb_get_iseq_body_iseq_encoded")
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index c27fa1f29d..ed1384571f 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -1251,6 +1251,7 @@ extern "C" {
     pub fn rb_get_mct_func(mct: *const rb_method_cfunc_t) -> *mut ::std::os::raw::c_void;
     pub fn rb_get_def_iseq_ptr(def: *mut rb_method_definition_t) -> *const rb_iseq_t;
     pub fn rb_get_def_bmethod_proc(def: *mut rb_method_definition_t) -> VALUE;
+    pub fn rb_get_iseq_body_total_calls(iseq: *const rb_iseq_t) -> ::std::os::raw::c_ulong;
     pub fn rb_get_iseq_body_local_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
     pub fn rb_get_iseq_body_parent_iseq(iseq: *const rb_iseq_t) -> *const rb_iseq_t;
     pub fn rb_get_iseq_body_local_table_size(iseq: *const rb_iseq_t) -> ::std::os::raw::c_uint;
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index 08440d7076..a453728702 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -45,10 +45,14 @@ pub fn yjit_enabled_p() -> bool {
     YJIT_ENABLED.load(Ordering::Acquire)
 }
 
-/// After how many calls YJIT starts compiling a method
+/// Test whether we are ready to compile an ISEQ or not
 #[no_mangle]
-pub extern "C" fn rb_yjit_call_threshold() -> raw::c_uint {
-    get_option!(call_threshold) as raw::c_uint
+pub extern "C" fn rb_yjit_threshold_hit(iseq: IseqPtr) -> bool {
+
+    let call_threshold = get_option!(call_threshold) as u64;
+    let total_calls = unsafe { rb_get_iseq_body_total_calls(iseq) } as u64;
+
+    return total_calls == call_threshold;
 }
 
 /// This function is called from C code