ZJIT: add support for lazy `RubyVM::ZJIT.enable`

This implements Shopify#854:

- Splits boot-time and enable-time initialization,
  tracks progress with `InitializationState` enum

- Introduces `RubyVM::ZJIT.enable` Ruby method for
  enabling the JIT lazily, if not already enabled

- Introduces `--zjit-disable` flag, which can be
  used alongside the other `--zjit-*` flags but
  prevents enabling the JIT at boot time

- Adds ZJIT infra to support JIT hooks, but this
  is not currently exercised (Shopify/ruby#667)

Left for future enhancements:

- Support kwargs for overriding the CLI flags in
  `RubyVM::ZJIT.enable`

Closes Shopify#854

1 files changed, 11 insertions, 8 deletions

diff --git a/ruby.c b/ruby.c
index 872a317e3b..f1089ca41e 100644
--- a/ruby.c
+++ b/ruby.c
@@ -1842,10 +1842,8 @@ ruby_opt_init(ruby_cmdline_options_t *opt)
     rb_yjit_init(opt->yjit);
 #endif
 #if USE_ZJIT
-    if (opt->zjit) {
-        extern void rb_zjit_init(void);
-        rb_zjit_init();
-    }
+    extern void rb_zjit_init(bool);
+    rb_zjit_init(opt->zjit);
 #endif
 
     ruby_set_script_name(opt->script_name);
@@ -2368,6 +2366,12 @@ process_options(int argc, char **argv, ruby_cmdline_options_t *opt)
 #if USE_ZJIT
         if (!FEATURE_USED_P(opt->features, zjit) && env_var_truthy("RUBY_ZJIT_ENABLE")) {
             FEATURE_SET(opt->features, FEATURE_BIT(zjit));
+
+            // When the --zjit flag is specified, we would have call setup_zjit_options(""),
+            // which would have called rb_zjit_prepare_options() internally. This ensures we
+            // go through the same set up but with less overhead than setup_zjit_options("").
+            extern void rb_zjit_prepare_options();
+            rb_zjit_prepare_options();
         }
 #endif
     }
@@ -2383,10 +2387,9 @@ process_options(int argc, char **argv, ruby_cmdline_options_t *opt)
     }
 #endif
 #if USE_ZJIT
-    if (FEATURE_SET_P(opt->features, zjit) && !opt->zjit) {
-        extern void rb_zjit_prepare_options(void);
-        rb_zjit_prepare_options();
-        opt->zjit = true;
+    if (FEATURE_SET_P(opt->features, zjit)) {
+        bool rb_zjit_option_enable(void);
+        opt->zjit = rb_zjit_option_enable(); // set opt->zjit for Init_ruby_description() and calling rb_zjit_init()
     }
 #endif