ZJIT: Add --zjit-stats=quiet option to collect stats without printing (#14467)

Similar to YJIT's --yjit-stats=quiet, this option allows ZJIT to collect
statistics and make them available via the Ruby API without printing them
at exit. This is useful for programmatic access to stats without the
output noise.

- Added print_stats field to Options struct
- Modified option parsing to support --zjit-stats=quiet
- Added rb_zjit_print_stats_p primitive to check if stats should be printed
- Updated zjit.rb to only register at_exit handler when print_stats is true
- Update the help text shown by `ruby --help` to indicate that
  --zjit-stats now accepts an optional =quiet parameter.
- Added test for --zjit-stats=quiet option

4 files changed, 48 insertions, 3 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index b751d482e2..9bfe2c0c00 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -27,6 +27,30 @@ class TestZJIT < Test::Unit::TestCase
     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)
+  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|
@@ -2490,7 +2514,7 @@ class TestZJIT < Test::Unit::TestCase
     if zjit
       args << "--zjit-call-threshold=#{call_threshold}"
       args << "--zjit-num-profiles=#{num_profiles}"
-      args << "--zjit-stats" if stats
+      args << "--zjit-stats#{"=#{stats}" unless stats == true}" if stats
       args << "--zjit-debug" if debug
       if allowed_iseqs
         jitlist = Tempfile.new("jitlist")
diff --git a/zjit.c b/zjit.c
index 313cced2aa..e1ea6d7e09 100644
--- a/zjit.c
+++ b/zjit.c
@@ -343,6 +343,7 @@ rb_zjit_insn_leaf(int insn, const VALUE *opes)
 VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self, VALUE target_key);
 VALUE rb_zjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self);
+VALUE rb_zjit_print_stats_p(rb_execution_context_t *ec, VALUE self);
 
 // Preprocessed zjit.rb generated during build
 #include "zjit.rbinc"
diff --git a/zjit.rb b/zjit.rb
index 7e5807876c..d70ff1dd47 100644
--- a/zjit.rb
+++ b/zjit.rb
@@ -8,7 +8,7 @@
 # for which CRuby is built.
 module RubyVM::ZJIT
   # Avoid calling a Ruby method here to avoid interfering with compilation tests
-  if Primitive.rb_zjit_stats_enabled_p
+  if Primitive.rb_zjit_print_stats_p
     at_exit { print_stats }
   end
 end
diff --git a/zjit/src/options.rs b/zjit/src/options.rs
index 3efcb933bb..dbb6ee8ebb 100644
--- a/zjit/src/options.rs
+++ b/zjit/src/options.rs
@@ -41,6 +41,9 @@ pub struct Options {
     /// Enable YJIT statsitics
     pub stats: bool,
 
+    /// Print stats on exit (when stats is also true)
+    pub print_stats: bool,
+
     /// Enable debug logging
     pub debug: bool,
 
@@ -78,6 +81,7 @@ impl Default for Options {
             exec_mem_bytes: 64 * 1024 * 1024,
             num_profiles: DEFAULT_NUM_PROFILES,
             stats: false,
+            print_stats: false,
             debug: false,
             disable_hir_opt: false,
             dump_hir_init: None,
@@ -103,7 +107,7 @@ pub const ZJIT_OPTIONS: &[(&str, &str)] = &[
                      "Number of calls to trigger JIT (default: 2)."),
     ("--zjit-num-profiles=num",
                      "Number of profiled calls before JIT (default: 1, max: 255)."),
-    ("--zjit-stats", "Enable collecting ZJIT statistics."),
+    ("--zjit-stats[=quiet]", "Enable collecting ZJIT statistics (=quiet to suppress output)."),
     ("--zjit-perf",  "Dump ISEQ symbols into /tmp/perf-{}.map for Linux perf."),
     ("--zjit-log-compiled-iseqs=path",
                      "Log compiled ISEQs to the file. The file will be truncated."),
@@ -220,6 +224,11 @@ fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
 
         ("stats", "") => {
             options.stats = true;
+            options.print_stats = true;
+        }
+        ("stats", "quiet") => {
+            options.stats = true;
+            options.print_stats = false;
         }
 
         ("debug", "") => options.debug = true,
@@ -344,3 +353,14 @@ pub extern "C" fn rb_zjit_stats_enabled_p(_ec: EcPtr, _self: VALUE) -> VALUE {
         Qfalse
     }
 }
+
+/// Return Qtrue if stats should be printed at exit.
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_print_stats_p(_ec: EcPtr, _self: VALUE) -> VALUE {
+    // Builtin zjit.rb calls this even if ZJIT is disabled, so OPTIONS may not be set.
+    if unsafe { OPTIONS.as_ref() }.is_some_and(|opts| opts.stats && opts.print_stats) {
+        Qtrue
+    } else {
+        Qfalse
+    }
+}