Assert everything is compiled in test_zjit (https://github.com/Shopify/zjit/pull/40)

* Assert everything is compiled in test_zjit

* Update a comment on rb_zjit_assert_compiles

Co-authored-by: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>

* Add a comment about assert_compiles

* Actually use pipe_fd

---------

Co-authored-by: Maxime Chevalier-Boisvert <maxime.chevalierboisvert@shopify.com>

.document

@@ -30,6 +30,7 @@ thread_sync.rb
 trace_point.rb
 warning.rb
 yjit.rb
+zjit.rb
 
 # Errno::*
 known_errors.inc

.github/workflows/zjit-macos.yml

@@ -31,15 +31,9 @@ jobs:
           - test_task: 'zjit-test'
             configure: '--enable-zjit=dev'
 
-          - test_task: 'btest'
-            zjit_opts: '--zjit-call-threshold=1'
+          - test_task: 'test-all'
             configure: '--enable-zjit=dev'
-            btests: '../src/bootstraptest/test_zjit.rb'
-
-          - test_task: 'btest'
-            zjit_opts: '--zjit-call-threshold=2'
-            configure: '--enable-zjit=dev'
-            btests: '../src/bootstraptest/test_zjit.rb'
+            tests: '../src/test/ruby/test_zjit.rb'
 
           # Test without ZJIT for now
           - test_task: 'check'
@@ -101,14 +95,13 @@ jobs:
           make -s ${{ matrix.test_task }} ${TESTS:+TESTS="$TESTS"}
           RUN_OPTS="$RUN_OPTS"
           SPECOPTS="$SPECOPTS"
-          BTESTS="$BTESTS"
         timeout-minutes: 60
         env:
           RUBY_TESTOPTS: '-q --tty=no'
           TEST_BUNDLED_GEMS_ALLOW_FAILURES: 'typeprof'
           SYNTAX_SUGGEST_TIMEOUT: '5'
           PRECHECK_BUNDLED_GEMS: 'no'
-          BTESTS: ${{ matrix.btests }}
+          TESTS: ${{ matrix.tests }}
         continue-on-error: ${{ matrix.continue-on-test_task || false }}
 
   result:

.github/workflows/zjit-ubuntu.yml

@@ -36,15 +36,9 @@ jobs:
           - test_task: 'zjit-test'
             configure: '--enable-zjit=dev'
 
-          - test_task: 'btest'
-            zjit_opts: '--zjit-call-threshold=1'
+          - test_task: 'test-all'
             configure: '--enable-zjit=dev'
-            btests: '../src/bootstraptest/test_zjit.rb'
-
-          - test_task: 'btest'
-            zjit_opts: '--zjit-call-threshold=2'
-            configure: '--enable-zjit=dev'
-            btests: '../src/bootstraptest/test_zjit.rb'
+            tests: '../src/test/ruby/test_zjit.rb'
 
           # Test without ZJIT for now
           - test_task: 'check'
@@ -125,7 +119,6 @@ jobs:
           make -s ${{ matrix.test_task }} ${TESTS:+TESTS="$TESTS"}
           RUN_OPTS="$RUN_OPTS" MSPECOPT=--debug SPECOPTS="$SPECOPTS"
           YJIT_BENCH_OPTS="$YJIT_BENCH_OPTS" YJIT_BINDGEN_DIFF_OPTS="$YJIT_BINDGEN_DIFF_OPTS"
-          BTESTS="$BTESTS"
         timeout-minutes: 90
         env:
           RUBY_TESTOPTS: '-q --tty=no'
@@ -134,7 +127,7 @@ jobs:
           SYNTAX_SUGGEST_TIMEOUT: '5'
           YJIT_BINDGEN_DIFF_OPTS: '--exit-code'
           LIBCLANG_PATH: ${{ matrix.libclang_path }}
-          BTESTS: ${{ matrix.btests }}
+          TESTS: ${{ matrix.tests }}
         continue-on-error: ${{ matrix.continue-on-test_task || false }}
 
   result:

bootstraptest/test_zjit.rb

@@ -1,33 +0,0 @@
-# Tests of Ruby methods that ZJIT can currently compile.
-# make btest BTESTS=bootstraptest/test_zjit.rb RUN_OPTS="--zjit"
-
-assert_equal 'nil', %q{
-  def test = nil
-  test; test.inspect
-}
-
-assert_equal '1', %q{
-  def test = 1
-  test; test
-}
-
-assert_equal '3', %q{
-  def test = 1 + 2
-  test; test
-}
-
-assert_equal '[6, 3]', %q{
-  def test(a, b) = a + b
-  [test(2, 4), test(1, 2)]
-}
-
-# Test argument ordering
-assert_equal '2', %q{
-  def test(a, b) = a - b
-  test(6, 4)
-}
-
-assert_equal '6', %q{
-  def test(a, b, c) = a + b + c
-  test(1, 2, 3)
-}

common.mk

@@ -1223,6 +1223,7 @@ BUILTIN_RB_SRCS = \
 		$(srcdir)/gem_prelude.rb \
 		$(srcdir)/yjit.rb \
 		$(srcdir)/yjit_hook.rb \
+		$(srcdir)/zjit.rb \
 		$(empty)
 BUILTIN_RB_INCS = $(BUILTIN_RB_SRCS:.rb=.rbinc)
 
@@ -10717,6 +10718,7 @@ miniinit.$(OBJEXT): {$(VPATH)}vm_opts.h
 miniinit.$(OBJEXT): {$(VPATH)}warning.rb
 miniinit.$(OBJEXT): {$(VPATH)}yjit.rb
 miniinit.$(OBJEXT): {$(VPATH)}yjit_hook.rb
+miniinit.$(OBJEXT): {$(VPATH)}zjit.rb
 node.$(OBJEXT): $(CCAN_DIR)/check_type/check_type.h
 node.$(OBJEXT): $(CCAN_DIR)/container_of/container_of.h
 node.$(OBJEXT): $(CCAN_DIR)/list/list.h
@@ -21179,4 +21181,5 @@ zjit.$(OBJEXT): {$(VPATH)}vm_opts.h
 zjit.$(OBJEXT): {$(VPATH)}vm_sync.h
 zjit.$(OBJEXT): {$(VPATH)}zjit.c
 zjit.$(OBJEXT): {$(VPATH)}zjit.h
+zjit.$(OBJEXT): {$(VPATH)}zjit.rbinc
 # AUTOGENERATED DEPENDENCIES END

inits.c

@@ -103,6 +103,7 @@ rb_call_builtin_inits(void)
     BUILTIN(thread_sync);
     BUILTIN(nilclass);
     BUILTIN(marshal);
+    BUILTIN(zjit);
     Init_builtin_prelude();
 }
 #undef CALL

test/lib/jit_support.rb

@@ -16,4 +16,10 @@ module JITSupport
   def yjit_force_enabled?
     "#{RbConfig::CONFIG['CFLAGS']} #{RbConfig::CONFIG['CPPFLAGS']}".match?(/(\A|\s)-D ?YJIT_FORCE_ENABLE\b/)
   end
+
+  def zjit_supported?
+    return @zjit_supported if defined?(@zjit_supported)
+    # nil in mswin
+    @zjit_supported = ![nil, 'no'].include?(RbConfig::CONFIG['ZJIT_SUPPORT'])
+  end
 end

test/ruby/test_zjit.rb

@@ -0,0 +1,117 @@
+# frozen_string_literal: true
+#
+# This set of tests can be run with:
+# make test-all TESTS=test/ruby/test_zjit.rb
+
+require 'test/unit'
+require 'envutil'
+require_relative '../lib/jit_support'
+return unless JITSupport.zjit_supported?
+
+class TestZJIT < Test::Unit::TestCase
+  def test_nil
+    assert_compiles nil, %q{
+      def test = nil
+      test
+    }
+  end
+
+  def test_putobject
+    assert_compiles 1, %q{
+      def test = 1
+      test
+    }
+  end
+
+  def test_opt_plus_const
+    assert_compiles 3, %q{
+      def test = 1 + 2
+      test # profile opt_plus
+      test
+    }, call_threshold: 2
+  end
+
+  def test_opt_plus_fixnum
+    assert_compiles 3, %q{
+      def test(a, b) = a + b
+      test(0, 1) # profile opt_plus
+      test(1, 2)
+    }, call_threshold: 2
+  end
+
+  def test_opt_plus_chain
+    assert_compiles 6, %q{
+      def test(a, b, c) = a + b + c
+      test(0, 1, 2) # profile opt_plus
+      test(1, 2, 3)
+    }, call_threshold: 2
+  end
+
+  # Test argument ordering
+  def test_opt_minus
+    omit 'FixnumSub is not implemented yet'
+    assert_compiles 2, %q{
+      def test(a, b) = a - b
+      test(2, 1) # profile opt_minus
+      test(6, 4)
+    }, call_threshold: 2
+  end
+
+  private
+
+  # Assert that every method call in `test_script` can be compiled by ZJIT
+  # at a given call_threshold
+  def assert_compiles(expected, test_script, call_threshold: 1)
+    pipe_fd = 3
+
+    script = <<~RUBY
+      _test_proc = -> {
+        RubyVM::ZJIT.assert_compiles
+        #{test_script}
+      }
+      result = _test_proc.call
+      IO.open(#{pipe_fd}).write(Marshal.dump(result))
+    RUBY
+
+    status, out, err, pipe_out = eval_with_jit(script, call_threshold:, pipe_fd:)
+
+    message = "exited with status #{status.to_i}"
+    message << "\nstdout:\n```\n#{out}```\n" unless out.empty?
+    message << "\nstderr:\n```\n#{err}```\n" unless err.empty?
+    assert status.success?, message
+
+    actual = Marshal.load(pipe_out)
+    assert_equal expected, actual
+  end
+
+  # Run a Ruby process with ZJIT options and a pipe for writing test results
+  def eval_with_jit(script, call_threshold: 1, timeout: 1000, pipe_fd:)
+    args = [
+      "--disable-gems",
+      "--zjit-call-threshold=#{call_threshold}",
+    ]
+    args << "-e" << script_shell_encode(script)
+    pipe_r, pipe_w = IO.pipe
+    # Separate thread so we don't deadlock when
+    # the child ruby blocks writing the output to pipe_fd
+    pipe_out = nil
+    pipe_reader = Thread.new do
+      pipe_out = pipe_r.read
+      pipe_r.close
+    end
+    out, err, status = EnvUtil.invoke_ruby(args, '', true, true, rubybin: RbConfig.ruby, timeout: timeout, ios: { pipe_fd => pipe_w })
+    pipe_w.close
+    pipe_reader.join(timeout)
+    [status, out, err, pipe_out]
+  ensure
+    pipe_reader&.kill
+    pipe_reader&.join(timeout)
+    pipe_r&.close
+    pipe_w&.close
+  end
+
+  def script_shell_encode(s)
+    # We can't pass utf-8-encoded characters directly in a shell arg. But we can use Ruby \u constants.
+    s.chars.map { |c| c.ascii_only? ? c : "\\u%x" % c.codepoints[0] }.join
+  end
+end

zjit.c

@@ -677,3 +677,9 @@ rb_iseq_set_zjit_payload(const rb_iseq_t *iseq, void *payload)
     RUBY_ASSERT_ALWAYS(NULL == iseq->body->zjit_payload);
     iseq->body->zjit_payload = payload;
 }
+
+// Primitives used by zjit.rb
+VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
+
+// Preprocessed zjit.rb generated during build
+#include "zjit.rbinc"

zjit.rb

@@ -0,0 +1,6 @@
+module RubyVM::ZJIT
+  # Assert that any future ZJIT compilation will return a function pointer
+  def self.assert_compiles
+    Primitive.rb_zjit_assert_compiles
+  end
+end

zjit/src/lib.rs

@@ -89,6 +89,15 @@ fn rb_bug_panic_hook() {
 /// Generate JIT code for a given ISEQ, which takes EC and CFP as its arguments.
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *const u8 {
+    let code_ptr = iseq_gen_entry_point(iseq);
+    if ZJITState::assert_compiles_enabled() && code_ptr == std::ptr::null() {
+        let iseq_location = iseq_get_location(iseq, 0);
+        panic!("Failed to compile: {iseq_location}");
+    }
+    code_ptr
+}
+
+fn iseq_gen_entry_point(iseq: IseqPtr) -> *const u8 {
     // Do not test the JIT code in HIR tests
     if cfg!(test) {
         return std::ptr::null();
@@ -116,3 +125,10 @@ pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *co
         }
     })
 }
+
+/// Assert that any future ZJIT compilation will return a function pointer (not fail to compile)
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_assert_compiles(_ec: EcPtr, _self: VALUE) -> VALUE {
+    ZJITState::enable_assert_compiles();
+    Qnil
+}

zjit/src/state.rs

@@ -12,6 +12,9 @@ pub struct ZJITState {
 
     /// Assumptions that require invalidation
     invariants: Invariants,
+
+    /// Assert successful compilation if set to true
+    assert_compiles: bool,
 }
 
 /// Private singleton instance of the codegen globals
@@ -64,6 +67,7 @@ impl ZJITState {
             code_block: cb,
             options,
             invariants: Invariants::default(),
+            assert_compiles: false,
         };
         unsafe { ZJIT_STATE = Some(zjit_state); }
     }
@@ -92,4 +96,15 @@ impl ZJITState {
     pub fn get_invariants() -> &'static mut Invariants {
         &mut ZJITState::get_instance().invariants
     }
+
+    /// Return true if successful compilation should be asserted
+    pub fn assert_compiles_enabled() -> bool {
+        ZJITState::get_instance().assert_compiles
+    }
+
+    /// Start asserting successful compilation
+    pub fn enable_assert_compiles() {
+        let instance = ZJITState::get_instance();
+        instance.assert_compiles = true;
+    }
 }