From e8c914c2502f6049ebc5bd93e8db0c62dd14400e Mon Sep 17 00:00:00 2001
From: Alan Wu <XrXr@users.noreply.github.com>
Date: Mon, 28 Sep 2020 07:03:28 -0400
Subject: [PATCH] Implement the --disable-ujit command line option

---
 compile.c      |  4 ++--
 ruby.c         | 10 ++++++++++
 ujit_compile.c | 15 ++++++++++-----
 ujit_compile.h |  3 +++
 4 files changed, 25 insertions(+), 7 deletions(-)

diff --git a/compile.c b/compile.c
index cbcda72949..5f6bc8476b 100644
--- a/compile.c
+++ b/compile.c
@@ -869,7 +869,7 @@ rb_iseq_translate_threaded_code(rb_iseq_t *iseq)
     unsigned int next_ujit_idx = 0;
     unsigned int translated_len = 0;
 
-    bool ujit_disabled = false /*get_cmdline_flag()*/;
+    bool ujit_enabled = rb_ujit_enabled_p();
 
     VALUE *translated_insns = ALLOCV_N(VALUE, translated_insns_buf, iseq->body->iseq_size);
     for (insn_idx = 0; insn_idx < iseq->body->iseq_size; /* */) {
@@ -880,7 +880,7 @@ rb_iseq_translate_threaded_code(rb_iseq_t *iseq)
         uint8_t* native_code_ptr = NULL;
 
         // If ujit is enabled and hasn't already compiled this instruction
-        if (!ujit_disabled && insn_idx >= next_ujit_idx)
+        if (ujit_enabled && insn_idx >= next_ujit_idx)
             native_code_ptr = ujit_compile_insn(iseq, insn_idx, &next_ujit_idx);
 
         if (native_code_ptr)
diff --git a/ruby.c b/ruby.c
index 818161710c..ace1e99fcc 100644
--- a/ruby.c
+++ b/ruby.c
@@ -59,6 +59,7 @@
 #include "internal/process.h"
 #include "internal/variable.h"
 #include "mjit.h"
+#include "ujit_compile.h"
 #include "ruby/encoding.h"
 #include "ruby/thread.h"
 #include "ruby/util.h"
@@ -103,6 +104,8 @@ void rb_warning_category_update(unsigned int mask, unsigned int bits);
     X(frozen_string_literal) \
     SEP \
     X(jit) \
+    SEP \
+    X(ujit)
     /* END OF FEATURES */
 #define EACH_DEBUG_FEATURES(X, SEP) \
     X(frozen_string_literal) \
@@ -229,6 +232,7 @@ cmdline_options_init(ruby_cmdline_options_t *opt)
 #ifdef MJIT_FORCE_ENABLE /* to use with: ./configure cppflags="-DMJIT_FORCE_ENABLE" */
     opt->features.set |= FEATURE_BIT(jit);
 #endif
+    opt->features.set |= FEATURE_BIT(ujit);
     return opt;
 }
 
@@ -327,6 +331,7 @@ usage(const char *name, int help, int highlight, int columns)
 	M("rubyopt", "",        "RUBYOPT environment variable (default: enabled)"),
 	M("frozen-string-literal", "", "freeze all string literals (default: disabled)"),
         M("jit", "",            "JIT compiler (default: disabled)"),
+        M("ujit", "",           "in-process JIT compiler (default: enabled)"),
     };
     static const struct message warn_categories[] = {
         M("deprecated", "",       "deprecated features"),
@@ -1435,6 +1440,9 @@ proc_options(long argc, char **argv, ruby_cmdline_options_t *opt, int envopt)
                 rb_warn("MJIT support is disabled.");
 #endif
             }
+            else if (strncmp("ujit", s, 4) == 0) {
+                FEATURE_SET(opt->features, FEATURE_BIT(jit));
+            }
 	    else if (strcmp("yydebug", s) == 0) {
 		if (envopt) goto noenvopt_long;
 		opt->dump |= DUMP_BIT(yydebug);
@@ -1861,6 +1869,8 @@ process_options(int argc, char **argv, ruby_cmdline_options_t *opt)
         /* Using TMP_RUBY_PREFIX created by ruby_init_loadpath(). */
         mjit_init(&opt->mjit);
 #endif
+    if (opt->features.set & FEATURE_BIT(ujit))
+        rb_ujit_init();
 
     Init_ruby_description();
     Init_enc();
diff --git a/ujit_compile.c b/ujit_compile.c
index be74f5f0e0..9097d78b8a 100644
--- a/ujit_compile.c
+++ b/ujit_compile.c
@@ -182,10 +182,8 @@ https://wiki.osdev.org/System_V_ABI#x86-64
 uint8_t *
 ujit_compile_insn(rb_iseq_t *iseq, unsigned int insn_idx, unsigned int* next_ujit_idx)
 {
-    // If not previously done, initialize ujit
-    if (!cb)
-    {
-        ujit_init();
+    if (!cb) {
+        return NULL;
     }
 
     // NOTE: if we are ever deployed in production, we
@@ -379,7 +377,14 @@ void gen_setlocal_wc0(codeblock_t* cb, codeblock_t* ocb, ctx_t* ctx)
     mov(cb, mem_opnd(64, RDX, offs), RCX);
 }
 
-static void ujit_init()
+bool
+rb_ujit_enabled_p(void)
+{
+    return !!cb;
+}
+
+void
+rb_ujit_init(void)
 {
     // Initialize the code blocks
     size_t mem_size = 64 * 1024 * 1024;
diff --git a/ujit_compile.h b/ujit_compile.h
index df9fd8e309..624f8dc5b7 100644
--- a/ujit_compile.h
+++ b/ujit_compile.h
@@ -3,12 +3,15 @@
 
 #include "stddef.h"
 #include "stdint.h"
+#include "stdbool.h"
 
 #ifndef rb_iseq_t
 typedef struct rb_iseq_struct rb_iseq_t;
 #define rb_iseq_t rb_iseq_t
 #endif
 
+void rb_ujit_init(void);
+bool rb_ujit_enabled_p(void);
 uint8_t* ujit_compile_insn(rb_iseq_t *iseq, unsigned int insn_idx, unsigned int* next_ujit_idx);
 
 #endif