1 files changed, 556 insertions, 0 deletions
diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
new file mode 100644
index 0000000000..a56fea6d51
--- /dev/null
+++ b/test/ruby/test_zjit.rb
@@ -0,0 +1,556 @@
+# 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_enabled
+    assert_runs 'false', <<~RUBY, zjit: false
+      RubyVM::ZJIT.enabled?
+    RUBY
+    assert_runs 'true', <<~RUBY, zjit: true
+      RubyVM::ZJIT.enabled?
+    RUBY
+  end
+
+  def test_stats_enabled
+    assert_runs 'false', <<~RUBY, stats: false
+      RubyVM::ZJIT.stats_enabled?
+    RUBY
+    assert_runs 'true', <<~RUBY, stats: true
+      RubyVM::ZJIT.stats_enabled?
+    RUBY
+  end
+
+  def test_stats_string_no_zjit
+    assert_runs 'nil', <<~RUBY, zjit: false
+      RubyVM::ZJIT.stats_string
+    RUBY
+    assert_runs 'true', <<~RUBY, stats: false
+      RubyVM::ZJIT.stats_string.is_a?(String)
+    RUBY
+    assert_runs 'true', <<~RUBY, stats: true
+      RubyVM::ZJIT.stats_string.is_a?(String)
+    RUBY
+  end
+
+  def test_stats_quiet
+    # Test that --zjit-stats-quiet collects stats but doesn't print them
+    script = <<~RUBY
+      def test = 42
+      test
+      test
+      puts RubyVM::ZJIT.stats_enabled?
+    RUBY
+
+    stats_header = "***ZJIT: Printing ZJIT statistics on exit***"
+
+    # With --zjit-stats, stats should be printed to stderr
+    out, err, status = eval_with_jit(script, stats: true)
+    assert_success(out, err, status)
+    assert_includes(err, stats_header)
+    assert_equal("true\n", out)
+
+    # With --zjit-stats-quiet, stats should NOT be printed but still enabled
+    out, err, status = eval_with_jit(script, stats: :quiet)
+    assert_success(out, err, status)
+    refute_includes(err, stats_header)
+    assert_equal("true\n", out)
+
+    # With --zjit-stats=<path>, stats should be printed to the path
+    Tempfile.create("zjit-stats-") {|tmp|
+      stats_file = tmp.path
+      tmp.puts("Lorem ipsum dolor sit amet, consectetur adipiscing elit, ...")
+      tmp.close
+
+      out, err, status = eval_with_jit(script, stats: stats_file)
+      assert_success(out, err, status)
+      refute_includes(err, stats_header)
+      assert_equal("true\n", out)
+      assert_equal stats_header, File.open(stats_file) {|f| f.gets(chomp: true)}, "should be overwritten"
+    }
+  end
+
+  def test_enable_through_env
+    child_env = {'RUBY_YJIT_ENABLE' => nil, 'RUBY_ZJIT_ENABLE' => '1'}
+    assert_in_out_err([child_env, '-v'], '') do |stdout, stderr|
+      assert_includes(stdout.first, '+ZJIT')
+      assert_equal([], stderr)
+    end
+  end
+
+  def test_zjit_enable
+    # --disable-all is important in case the build/environment has YJIT enabled by
+    # default through e.g. -DYJIT_FORCE_ENABLE. Can't enable ZJIT when YJIT is on.
+    assert_separately(["--disable-all"], <<~'RUBY')
+      refute_predicate RubyVM::ZJIT, :enabled?
+      refute_predicate RubyVM::ZJIT, :stats_enabled?
+      refute_includes RUBY_DESCRIPTION, "+ZJIT"
+
+      RubyVM::ZJIT.enable
+
+      assert_predicate RubyVM::ZJIT, :enabled?
+      refute_predicate RubyVM::ZJIT, :stats_enabled?
+      assert_includes RUBY_DESCRIPTION, "+ZJIT"
+    RUBY
+  end
+
+  def test_zjit_disable
+    assert_separately(["--zjit", "--zjit-disable"], <<~'RUBY')
+      refute_predicate RubyVM::ZJIT, :enabled?
+      refute_includes RUBY_DESCRIPTION, "+ZJIT"
+
+      RubyVM::ZJIT.enable
+
+      assert_predicate RubyVM::ZJIT, :enabled?
+      assert_includes RUBY_DESCRIPTION, "+ZJIT"
+    RUBY
+  end
+
+  def test_zjit_prelude_kernel_prepend
+    # Simulate what bundler/setup can do: prepend a module to Kernel during
+    # the prelude via the BUNDLER_SETUP mechanism in rubygems.rb:
+    #   require ENV["BUNDLER_SETUP"] if ENV["BUNDLER_SETUP"] && !defined?(Bundler)
+    Tempfile.create(["kernel_prepend", ".rb"]) do |f|
+      f.write("Kernel.prepend(Module.new)\n")
+      f.flush
+      assert_separately([{ "BUNDLER_SETUP" => f.path }, "--enable=gems", "--zjit"], "", ignore_stderr: true)
+    end
+  end
+
+  def test_zjit_enable_respects_existing_options
+    assert_separately(['--zjit-disable', '--zjit-stats-quiet'], <<~RUBY)
+      refute_predicate RubyVM::ZJIT, :enabled?
+      assert_predicate RubyVM::ZJIT, :stats_enabled?
+
+      RubyVM::ZJIT.enable
+
+      assert_predicate RubyVM::ZJIT, :enabled?
+      assert_predicate RubyVM::ZJIT, :stats_enabled?
+    RUBY
+  end
+
+  def test_toplevel_binding
+    # Not using assert_compiles, which doesn't use the toplevel frame for `test_script`.
+    out, err, status = eval_with_jit(%q{
+      a = 1
+      b = 2
+      TOPLEVEL_BINDING.local_variable_set(:b, 3)
+      c = 4
+      print [a, b, c]
+    })
+    assert_success(out, err, status)
+    assert_equal "[1, 3, 4]", out
+  end
+
+  def test_send_exit_with_uninitialized_locals
+    assert_runs 'nil', %q{
+      def entry(init)
+        function_stub_exit(init)
+      end
+
+      def function_stub_exit(init)
+        uninitialized_local = 1 if init
+        uninitialized_local
+      end
+
+      entry(true) # profile and set 1 to the local slot
+      entry(false)
+    }, call_threshold: 2, allowed_iseqs: 'entry@-e:2'
+  end
+
+  def test_opt_new_with_custom_allocator
+    assert_compiles '"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"', %q{
+      require "digest"
+      def test = Digest::SHA256.new.hexdigest
+      test; test
+    }, insns: [:opt_new], call_threshold: 2
+  end
+
+  def test_opt_new_with_custom_allocator_raises
+    assert_compiles '[42, 42]', %q{
+      require "digest"
+      class C < Digest::Base; end
+      def test
+        begin
+          Digest::Base.new
+        rescue NotImplementedError
+          42
+        end
+      end
+      [test, test]
+    }, insns: [:opt_new], call_threshold: 2
+  end
+
+  def test_uncached_getconstant_path
+    assert_compiles RUBY_COPYRIGHT.dump, %q{
+      def test = RUBY_COPYRIGHT
+      test
+    }, call_threshold: 1, insns: [:opt_getconstant_path]
+  end
+
+  def test_getconstant_path_autoload
+    # A constant-referencing expression can run arbitrary code through Kernel#autoload.
+    Dir.mktmpdir('autoload') do |tmpdir|
+      autoload_path = File.join(tmpdir, 'test_getconstant_path_autoload.rb')
+      File.write(autoload_path, 'X = RUBY_COPYRIGHT')
+
+      assert_compiles RUBY_COPYRIGHT.dump, %Q{
+        Object.autoload(:X, #{File.realpath(autoload_path).inspect})
+        def test = X
+        test
+      }, call_threshold: 1, insns: [:opt_getconstant_path]
+    end
+  end
+
+  def test_send_backtrace
+    backtrace = [
+      "-e:2:in 'Object#jit_frame1'",
+      "-e:3:in 'Object#entry'",
+      "-e:5:in 'block in <main>'",
+      "-e:6:in '<main>'",
+    ]
+    assert_compiles backtrace.inspect, %q{
+      def jit_frame2 = caller     # 1
+      def jit_frame1 = jit_frame2 # 2
+      def entry = jit_frame1      # 3
+      entry # profile send        # 4
+      entry                       # 5
+    }, call_threshold: 2
+  end
+
+  # tool/ruby_vm/views/*.erb relies on the zjit instructions a) being contiguous and
+  # b) being reliably ordered after all the other instructions.
+  def test_instruction_order
+    insn_names = RubyVM::INSTRUCTION_NAMES
+    zjit, others = insn_names.map.with_index.partition { |name, _| name.start_with?('zjit_') }
+    zjit_indexes = zjit.map(&:last)
+    other_indexes = others.map(&:last)
+    zjit_indexes.product(other_indexes).each do |zjit_index, other_index|
+      assert zjit_index > other_index, "'#{insn_names[zjit_index]}' at #{zjit_index} "\
+        "must be defined after '#{insn_names[other_index]}' at #{other_index}"
+    end
+  end
+
+  def test_require_rubygems
+    assert_runs 'true', %q{
+      require 'rubygems'
+    }, call_threshold: 2
+  end
+
+  def test_require_rubygems_with_auto_compact
+    omit("GC.auto_compact= support is required for this test") unless GC.respond_to?(:auto_compact=)
+    assert_runs 'true', %q{
+      GC.auto_compact = true
+      require 'rubygems'
+    }, call_threshold: 2
+  end
+
+  def test_stats_availability
+    assert_runs '[true, true]', %q{
+      def test = 1
+      test
+      [
+        RubyVM::ZJIT.stats[:zjit_insn_count] > 0,
+        RubyVM::ZJIT.stats(:zjit_insn_count) > 0,
+      ]
+    }, stats: true
+  end
+
+  def test_stats_consistency
+    assert_runs '[]', %q{
+      def test = 1
+      test # increment some counters
+
+      RubyVM::ZJIT.stats.to_a.filter_map do |key, value|
+        # The value may be incremented, but the class should stay the same
+        other_value = RubyVM::ZJIT.stats(key)
+        if value.class != other_value.class
+          [key, value, other_value]
+        end
+      end
+    }, stats: true
+  end
+
+  def test_reset_stats
+    assert_runs 'true', %q{
+      def test = 1
+      100.times { test }
+
+      # Get initial stats and verify they're non-zero
+      initial_stats = RubyVM::ZJIT.stats
+
+      # Reset the stats
+      RubyVM::ZJIT.reset_stats!
+
+      # Get stats after reset
+      reset_stats = RubyVM::ZJIT.stats
+
+      [
+        # After reset, counters should be zero or at least much smaller
+        # (some instructions might execute between reset and reading stats)
+        :zjit_insn_count.then { |s| initial_stats[s] > 0 && reset_stats[s] < initial_stats[s] },
+        :compiled_iseq_count.then { |s| initial_stats[s] > 0 && reset_stats[s] < initial_stats[s] }
+      ].all?
+    }, stats: true
+  end
+
+  def test_zjit_option_uses_array_each_in_ruby
+    omit 'ZJIT wrongly compiles Array#each, so it is disabled for now'
+    assert_runs '"<internal:array>"', %q{
+      Array.instance_method(:each).source_location&.first
+    }
+  end
+
+  def test_line_tracepoint_on_c_method
+    assert_compiles '"[[:line, true]]"', %q{
+      events = []
+      events.instance_variable_set(
+        :@tp,
+        TracePoint.new(:line) { |tp| events << [tp.event, tp.lineno] if tp.path == __FILE__ }
+      )
+      def events.to_str
+        @tp.enable; ''
+      end
+
+      # Stay in generated code while enabling tracing
+      def events.compiled(obj)
+        String(obj)
+        @tp.disable; __LINE__
+      end
+
+      line = events.compiled(events)
+      events[0][-1] = (events[0][-1] == line)
+
+      events.to_s # can't dump events as it's a singleton object AND it has a TracePoint instance variable, which also can't be dumped
+    }
+  end
+
+  def test_targeted_line_tracepoint_in_c_method_call
+    assert_compiles '"[true]"', %q{
+      events = []
+      events.instance_variable_set(:@tp, TracePoint.new(:line) { |tp| events << tp.lineno })
+      def events.to_str
+        @tp.enable(target: method(:compiled))
+        ''
+      end
+
+      # Stay in generated code while enabling tracing
+      def events.compiled(obj)
+        String(obj)
+        __LINE__
+      end
+
+      line = events.compiled(events)
+      events[0] = (events[0] == line)
+
+      events.to_s # can't dump events as it's a singleton object AND it has a TracePoint instance variable, which also can't be dumped
+    }
+  end
+
+  def test_regression_cfp_sp_set_correctly_before_leaf_gc_call
+    assert_compiles ':ok', %q{
+      def check(l, r)
+        return 1 unless l
+        1 + check(*l) + check(*r)
+      end
+
+      def tree(depth)
+        # This duparray is our leaf-gc target.
+        return [nil, nil] unless depth > 0
+
+        # Modify the local and pass it to the following calls.
+        depth -= 1
+        [tree(depth), tree(depth)]
+      end
+
+      def test
+        GC.stress = true
+        2.times do
+          t = tree(11)
+          check(*t)
+        end
+        :ok
+      end
+
+      test
+    }, call_threshold: 14, num_profiles: 5
+  end
+
+  def test_exit_tracing
+    # Smoke test: --zjit-trace-exits writes a Fuchsia trace (.fxt) file to /tmp
+    assert_compiles('true', <<~RUBY, extra_args: ['--zjit-trace-exits'])
+      def test(object) = object.itself
+
+      # induce an exit just for good measure
+      array = []
+      test(array)
+      test(array)
+      def array.itself = :not_itself
+      test(array)
+
+      fxt_files = Dir.glob("/tmp/perfetto-\#{Process.pid}.fxt")
+      result = fxt_files.length == 1 && !File.empty?(fxt_files.first)
+      File.unlink(*fxt_files)
+      result
+    RUBY
+  end
+
+  def test_send_no_profiles_with_disabled_specialized_instruction
+    # Regression test: when specialized_instruction is disabled (as power_assert does),
+    # eval'd code uses `send` instead of `opt_send_without_block`, producing SendNoProfiles.
+    # The `times` call with a literal block is the SendNoProfiles send whose exit profiling
+    # triggers recompilation of `run`. After recompilation, `make`'s eval("proc { }") crashes
+    # in vm_make_env_each because the caller frame's EP[-1] (specval) has a stale value.
+    assert_runs ':ok', <<~RUBY
+      RubyVM::InstructionSequence.compile_option = { specialized_instruction: false }
+      eval <<~'INNERRUBY'
+        def make = eval("proc { }")
+        def run(n) = n.times { make }
+      INNERRUBY
+      run(6)
+      :ok
+    RUBY
+  end
+
+  def test_float_arithmetic
+    assert_compiles '4.0', 'def test = 1.5 + 2.5; test'
+    assert_compiles '6.0', 'def test = 2.0 * 3.0; test'
+    assert_compiles '1.5', 'def test = 3.5 - 2.0; test'
+    assert_compiles '2.5', 'def test = 5.0 / 2.0; test'
+    assert_compiles '4.5', 'def test = 1.5 * 3; test' # Float * Fixnum
+    assert_compiles 'true', 'def test = (Float::NAN + 1.0).nan?; test'
+    assert_compiles 'Infinity', 'def test = Float::INFINITY * 2.0; test'
+    assert_compiles '3', 'def test = 3.7.to_i; test'
+    assert_compiles '-2', 'def test = (-2.9).to_i; test'
+  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, insns: [], **opts)
+    assert_runs(expected, test_script, insns:, assert_compiles: true, **opts)
+  end
+
+  # Assert that `test_script` runs successfully with ZJIT enabled.
+  # Unlike `assert_compiles`, `assert_runs(assert_compiles: false)`
+  # allows ZJIT to skip compiling methods.
+  def assert_runs(expected, test_script, insns: [], assert_compiles: false, **opts)
+    pipe_fd = 3
+    disasm_method = :test
+
+    script = <<~RUBY
+      ret_val = (_test_proc = -> { #{('RubyVM::ZJIT.assert_compiles; ' if assert_compiles)}#{test_script.lstrip} }).call
+      result = {
+        ret_val:,
+        #{ unless insns.empty?
+           "insns: RubyVM::InstructionSequence.of(method(#{disasm_method.inspect})).to_a"
+        end}
+      }
+      IO.open(#{pipe_fd}).write(Marshal.dump(result))
+    RUBY
+
+    out, err, status, result = eval_with_jit(script, pipe_fd:, **opts)
+    assert_success(out, err, status)
+
+    result = Marshal.load(result)
+    assert_equal(expected, result.fetch(:ret_val).inspect)
+
+    unless insns.empty?
+      iseq = result.fetch(:insns)
+      assert_equal(
+        "YARVInstructionSequence/SimpleDataFormat",
+        iseq.first,
+        "Failed to get ISEQ disassembly. " \
+        "Make sure to put code directly under the '#{disasm_method}' method."
+      )
+      iseq_insns = iseq.last
+
+      expected_insns = Set.new(insns)
+      iseq_insns.each do
+        next unless it.is_a?(Array)
+        expected_insns.delete(it.first)
+      end
+      assert(expected_insns.empty?, -> { "Not present in ISeq: #{expected_insns.to_a}" })
+    end
+  end
+
+  # Run a Ruby process with ZJIT options and a pipe for writing test results
+  def eval_with_jit(
+    script,
+    call_threshold: 1,
+    num_profiles: 1,
+    zjit: true,
+    stats: false,
+    debug: true,
+    allowed_iseqs: nil,
+    extra_args: nil,
+    timeout: 1000,
+    pipe_fd: nil
+  )
+    args = ["--disable-gems", *extra_args]
+    if zjit
+      args << "--zjit-call-threshold=#{call_threshold}"
+      args << "--zjit-num-profiles=#{num_profiles}"
+      case stats
+      when true
+        args << "--zjit-stats"
+      when :quiet
+        args << "--zjit-stats-quiet"
+      else
+        args << "--zjit-stats=#{stats}" if stats
+      end
+      args << "--zjit-debug" if debug
+      if allowed_iseqs
+        jitlist = Tempfile.new("jitlist")
+        jitlist.write(allowed_iseqs)
+        jitlist.close
+        args << "--zjit-allowed-iseqs=#{jitlist.path}"
+      end
+    end
+    args << "-e" << script_shell_encode(script)
+    ios = {}
+    if pipe_fd
+      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
+      ios[pipe_fd] = pipe_w
+    end
+    result = EnvUtil.invoke_ruby(args, '', true, true, rubybin: RbConfig.ruby, timeout: timeout, ios:)
+    if pipe_fd
+      pipe_w.close
+      pipe_reader.join(timeout)
+      result << pipe_out
+    end
+    result
+  ensure
+    pipe_reader&.kill
+    pipe_reader&.join(timeout)
+    pipe_r&.close
+    pipe_w&.close
+    jitlist&.unlink
+  end
+
+  def assert_success(out, err, status)
+    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
+  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