Mark inlined ISeqs during MJIT compilation (#4539)

[Bug #17584]

common.mk

@@ -8408,6 +8408,7 @@ mjit.$(OBJEXT): $(hdrdir)/ruby/ruby.h
 mjit.$(OBJEXT): $(hdrdir)/ruby/version.h
 mjit.$(OBJEXT): $(top_srcdir)/internal/array.h
 mjit.$(OBJEXT): $(top_srcdir)/internal/class.h
+mjit.$(OBJEXT): $(top_srcdir)/internal/compile.h
 mjit.$(OBJEXT): $(top_srcdir)/internal/compilers.h
 mjit.$(OBJEXT): $(top_srcdir)/internal/cont.h
 mjit.$(OBJEXT): $(top_srcdir)/internal/file.h
@@ -8429,6 +8430,7 @@ mjit.$(OBJEXT): {$(VPATH)}backward/2/limits.h
 mjit.$(OBJEXT): {$(VPATH)}backward/2/long_long.h
 mjit.$(OBJEXT): {$(VPATH)}backward/2/stdalign.h
 mjit.$(OBJEXT): {$(VPATH)}backward/2/stdarg.h
+mjit.$(OBJEXT): {$(VPATH)}builtin.h
 mjit.$(OBJEXT): {$(VPATH)}config.h
 mjit.$(OBJEXT): {$(VPATH)}constant.h
 mjit.$(OBJEXT): {$(VPATH)}debug.h
@@ -8439,6 +8441,9 @@ mjit.$(OBJEXT): {$(VPATH)}encoding.h
 mjit.$(OBJEXT): {$(VPATH)}gc.h
 mjit.$(OBJEXT): {$(VPATH)}id.h
 mjit.$(OBJEXT): {$(VPATH)}id_table.h
+mjit.$(OBJEXT): {$(VPATH)}insns.def
+mjit.$(OBJEXT): {$(VPATH)}insns.inc
+mjit.$(OBJEXT): {$(VPATH)}insns_info.inc
 mjit.$(OBJEXT): {$(VPATH)}intern.h
 mjit.$(OBJEXT): {$(VPATH)}internal.h
 mjit.$(OBJEXT): {$(VPATH)}internal/anyargs.h
@@ -8580,6 +8585,7 @@ mjit.$(OBJEXT): {$(VPATH)}internal/value_type.h
 mjit.$(OBJEXT): {$(VPATH)}internal/variable.h
 mjit.$(OBJEXT): {$(VPATH)}internal/warning_push.h
 mjit.$(OBJEXT): {$(VPATH)}internal/xmalloc.h
+mjit.$(OBJEXT): {$(VPATH)}iseq.h
 mjit.$(OBJEXT): {$(VPATH)}method.h
 mjit.$(OBJEXT): {$(VPATH)}missing.h
 mjit.$(OBJEXT): {$(VPATH)}mjit.c

mjit.c

@@ -925,24 +925,32 @@ mjit_mark(void)
         return;
     RUBY_MARK_ENTER("mjit");
 
-    if (compiling_iseq != NULL)
-        rb_gc_mark((VALUE)compiling_iseq);
-
     // We need to release a lock when calling rb_gc_mark to avoid doubly acquiring
     // a lock by by mjit_gc_start_hook inside rb_gc_mark.
     //
     // Because an MJIT worker may modify active_units anytime, we need to convert
     // the linked list to an array to safely loop its ISeqs without keeping a lock.
     CRITICAL_SECTION_START(4, "mjit_mark");
-    int length = active_units.length;
-    rb_iseq_t **iseqs = ALLOCA_N(rb_iseq_t *, length);
+    int length = 0;
+    if (compiling_iseqs != NULL) {
+        while (compiling_iseqs[length]) length++;
+    }
+    length += active_units.length;
+    const rb_iseq_t **iseqs = ALLOCA_N(const rb_iseq_t *, length);
 
     struct rb_mjit_unit *unit = NULL;
     int i = 0;
+    if (compiling_iseqs != NULL) {
+        while (compiling_iseqs[i]) {
+            iseqs[i] = compiling_iseqs[i];
+            i++;
+        }
+    }
     list_for_each(&active_units.head, unit, unode) {
         iseqs[i] = unit->iseq;
         i++;
     }
+    assert(i == length);
     CRITICAL_SECTION_FINISH(4, "mjit_mark");
 
     for (i = 0; i < length; i++) {

mjit_compile.c

@@ -441,6 +441,22 @@ inlinable_iseq_p(const struct rb_iseq_constant_body *body)
     return true;
 }
 
+// Return an iseq pointer if cc has inlinable iseq.
+const rb_iseq_t *
+rb_mjit_inlinable_iseq(const struct rb_callinfo *ci, const struct rb_callcache *cc)
+{
+    const rb_iseq_t *iseq;
+    if (has_valid_method_type(cc) &&
+        !(vm_ci_flag(ci) & VM_CALL_TAILCALL) && // inlining only non-tailcall path
+        vm_cc_cme(cc)->def->type == VM_METHOD_TYPE_ISEQ &&
+        fastpath_applied_iseq_p(ci, cc, iseq = def_iseq_ptr(vm_cc_cme(cc)->def)) &&
+        // CC_SET_FASTPATH in vm_callee_setup_arg
+        inlinable_iseq_p(iseq->body)) {
+        return iseq;
+    }
+    return NULL;
+}
+
 static void
 init_ivar_compile_status(const struct rb_iseq_constant_body *body, struct compile_status *status)
 {
@@ -521,13 +537,9 @@ precompile_inlinable_iseqs(FILE *f, const rb_iseq_t *iseq, struct compile_status
             const struct rb_callinfo *ci = cd->ci;
             const struct rb_callcache *cc = captured_cc_entries(status)[call_data_index(cd, body)]; // use copy to avoid race condition
 
+            extern bool rb_mjit_compiling_iseq_p(const rb_iseq_t *iseq);
             const rb_iseq_t *child_iseq;
-            if (has_valid_method_type(cc) &&
-                !(vm_ci_flag(ci) & VM_CALL_TAILCALL) && // inlining only non-tailcall path
-                vm_cc_cme(cc)->def->type == VM_METHOD_TYPE_ISEQ &&
-                fastpath_applied_iseq_p(ci, cc, child_iseq = def_iseq_ptr(vm_cc_cme(cc)->def)) &&
-                // CC_SET_FASTPATH in vm_callee_setup_arg
-                inlinable_iseq_p(child_iseq->body)) {
+            if ((child_iseq = rb_mjit_inlinable_iseq(ci, cc)) != NULL && rb_mjit_compiling_iseq_p(child_iseq)) {
                 status->inlined_iseqs[pos] = child_iseq->body;
 
                 if (mjit_opts.verbose >= 1) // print beforehand because ISeq may be GCed during copy job.

mjit_worker.c

@@ -93,6 +93,10 @@
 #include "ruby/debug.h"
 #include "ruby/thread.h"
 #include "ruby/version.h"
+#include "builtin.h"
+#include "insns.inc"
+#include "insns_info.inc"
+#include "internal/compile.h"
 
 #ifdef _WIN32
 #include <winsock2.h>
@@ -716,6 +720,51 @@ sprint_funcname(char *funcname, const struct rb_mjit_unit *unit)
     }
 }
 
+static const rb_iseq_t **compiling_iseqs = NULL;
+
+static bool
+set_compiling_iseqs(const rb_iseq_t *iseq)
+{
+    compiling_iseqs = calloc(iseq->body->iseq_size + 2, sizeof(rb_iseq_t *)); // 2: 1 (unit->iseq) + 1 (NULL end)
+    if (compiling_iseqs == NULL)
+        return false;
+
+    compiling_iseqs[0] = iseq;
+    int i = 1;
+
+    unsigned int pos = 0;
+    while (pos < iseq->body->iseq_size) {
+#if OPT_DIRECT_THREADED_CODE || OPT_CALL_THREADED_CODE
+        int insn = rb_vm_insn_addr2insn((void *)iseq->body->iseq_encoded[pos]);
+#else
+        int insn = (int)iseq->body->iseq_encoded[pos];
+#endif
+        if (insn == BIN(opt_send_without_block)) {
+            CALL_DATA cd = (CALL_DATA)iseq->body->iseq_encoded[pos + 1];
+            extern const rb_iseq_t *rb_mjit_inlinable_iseq(const struct rb_callinfo *ci, const struct rb_callcache *cc);
+            const rb_iseq_t *iseq = rb_mjit_inlinable_iseq(cd->ci, cd->cc);
+            if (iseq != NULL) {
+                compiling_iseqs[i] = iseq;
+                i++;
+            }
+        }
+        pos += insn_len(insn);
+    }
+    return true;
+}
+
+bool
+rb_mjit_compiling_iseq_p(const rb_iseq_t *iseq)
+{
+    assert(compiling_iseqs != NULL);
+    int i = 0;
+    while (compiling_iseqs[i]) {
+        if (compiling_iseqs[i] == iseq) return true;
+        i++;
+    }
+    return false;
+}
+
 static const int c_file_access_mode =
 #ifdef O_BINARY
     O_BINARY|
@@ -939,6 +988,11 @@ compile_compact_jit_code(char* c_file)
     // compacted functions (not done yet).
     bool success = true;
     list_for_each(&active_units.head, child_unit, unode) {
+        CRITICAL_SECTION_START(3, "before set_compiling_iseqs");
+        success &= set_compiling_iseqs(child_unit->iseq);
+        CRITICAL_SECTION_FINISH(3, "after set_compiling_iseqs");
+        if (!success) continue;
+
         char funcname[MAXPATHLEN];
         sprint_funcname(funcname, child_unit);
 
@@ -952,6 +1006,11 @@ compile_compact_jit_code(char* c_file)
         if (!iseq_label) iseq_label = sep = "";
         fprintf(f, "\n/* %s%s%s:%ld */\n", iseq_label, sep, iseq_path, iseq_lineno);
         success &= mjit_compile(f, child_unit->iseq, funcname, child_unit->id);
+
+        CRITICAL_SECTION_START(3, "before compiling_iseqs free");
+        free(compiling_iseqs);
+        compiling_iseqs = NULL;
+        CRITICAL_SECTION_FINISH(3, "after compiling_iseqs free");
     }
 
     // release blocking mjit_gc_start_hook
@@ -1091,8 +1150,6 @@ compile_prelude(FILE *f)
 #endif
 }
 
-static rb_iseq_t *compiling_iseq = NULL;
-
 // Compile ISeq in UNIT and return function pointer of JIT-ed code.
 // It may return NOT_COMPILED_JIT_ISEQ_FUNC if something went wrong.
 static mjit_func_t
@@ -1127,7 +1184,7 @@ convert_unit_to_func(struct rb_mjit_unit *unit)
     // We need to check again here because we could've waited on GC above
     in_jit = (unit->iseq != NULL);
     if (in_jit)
-        compiling_iseq = unit->iseq;
+        in_jit &= set_compiling_iseqs(unit->iseq);
     CRITICAL_SECTION_FINISH(3, "before mjit_compile to wait GC finish");
     if (!in_jit) {
         fclose(f);
@@ -1152,7 +1209,8 @@ convert_unit_to_func(struct rb_mjit_unit *unit)
 
     // release blocking mjit_gc_start_hook
     CRITICAL_SECTION_START(3, "after mjit_compile to wakeup client for GC");
-    compiling_iseq = NULL;
+    free(compiling_iseqs);
+    compiling_iseqs = NULL;
     in_jit = false;
     verbose(3, "Sending wakeup signal to client in a mjit-worker for GC");
     rb_native_cond_signal(&mjit_client_wakeup);