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, 25 insertions, 0 deletions

diff --git a/zjit.rb b/zjit.rb
index 7cdf84cfbe..f5ed347f2c 100644
--- a/zjit.rb
+++ b/zjit.rb
@@ -7,6 +7,8 @@
 # This module may not exist if ZJIT does not support the particular platform
 # for which CRuby is built.
 module RubyVM::ZJIT
+  # Blocks that are called when YJIT is enabled
+  @jit_hooks = []
   # Avoid calling a Ruby method here to avoid interfering with compilation tests
   if Primitive.rb_zjit_print_stats_p
     at_exit { print_stats }
@@ -22,6 +24,18 @@ class << RubyVM::ZJIT
     Primitive.cexpr! 'RBOOL(rb_zjit_enabled_p)'
   end
 
+  # Enable ZJIT compilation.
+  def enable
+    return false if enabled?
+
+    if Primitive.cexpr! 'RBOOL(rb_yjit_enabled_p)'
+      warn("Only one JIT can be enabled at the same time.")
+      return false
+    end
+
+    Primitive.rb_zjit_enable
+  end
+
   # Check if `--zjit-trace-exits` is used
   def trace_exit_locations_enabled?
     Primitive.rb_zjit_trace_exit_locations_enabled_p
@@ -234,6 +248,17 @@ class << RubyVM::ZJIT
   # :stopdoc:
   private
 
+  # Register a block to be called when ZJIT is enabled
+  def add_jit_hook(hook)
+    @jit_hooks << hook
+  end
+
+  # Run ZJIT hooks registered by `#with_jit`
+  def call_jit_hooks
+    @jit_hooks.each(&:call)
+    @jit_hooks.clear
+  end
+
   def print_counters(keys, buf:, stats:, right_align: false, base: nil)
     key_pad = keys.map { |key| key.to_s.sub(/_time_ns\z/, '_time').size }.max + 1
     key_align = '-' unless right_align