ZJIT: Allow querying a single ZJIT stat (#14309)

* ZJIT: Add RubyVM::ZJIT.stats_enabled?

* ZJIT: Allow querying a single ZJIT stat

4 files changed, 43 insertions, 14 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 9b4271ffc0..8c4ddd598e 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -18,6 +18,15 @@ class TestZJIT < Test::Unit::TestCase
     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_enable_through_env
     child_env = {'RUBY_YJIT_ENABLE' => nil, 'RUBY_ZJIT_ENABLE' => '1'}
     assert_in_out_err([child_env, '-v'], '') do |stdout, stderr|
@@ -1486,10 +1495,13 @@ class TestZJIT < Test::Unit::TestCase
   end
 
   def test_stats
-    assert_runs 'true', %q{
+    assert_runs '[true, true]', %q{
       def test = 1
       test
-      RubyVM::ZJIT.stats[:zjit_insns_count] > 0
+      [
+        RubyVM::ZJIT.stats[:zjit_insns_count] > 0,
+        RubyVM::ZJIT.stats(:zjit_insns_count) > 0,
+      ]
     }, stats: true
   end
 
diff --git a/zjit.c b/zjit.c
index 54d9f7ed86..9cc8b51423 100644
--- a/zjit.c
+++ b/zjit.c
@@ -358,7 +358,7 @@ rb_zjit_singleton_class_p(VALUE klass)
 
 // Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them.
 VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self);
-VALUE rb_zjit_stats(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);
 
 // Preprocessed zjit.rb generated during build
diff --git a/zjit.rb b/zjit.rb
index 7f98c5adc7..871b1a91ca 100644
--- a/zjit.rb
+++ b/zjit.rb
@@ -19,10 +19,15 @@ class << RubyVM::ZJIT
     Primitive.cexpr! 'RBOOL(rb_zjit_enabled_p)'
   end
 
+  # Check if `--zjit-stats` is used
+  def stats_enabled?
+    Primitive.rb_zjit_stats_enabled_p
+  end
+
   # Return ZJIT statistics as a Hash
-  def stats
-    stats = Primitive.rb_zjit_stats
-    return nil if stats.nil?
+  def stats(key = nil)
+    stats = Primitive.rb_zjit_stats(key)
+    return stats if stats.nil? || !key.nil?
 
     if stats.key?(:vm_insns_count) && stats.key?(:zjit_insns_count)
       stats[:total_insns_count] = stats[:vm_insns_count] + stats[:zjit_insns_count]
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
index fa8b741eea..ce185597c4 100644
--- a/zjit/src/stats.rs
+++ b/zjit/src/stats.rs
@@ -72,29 +72,41 @@ fn incr_counter(counter: Counter, amount: u64) {
 
 /// Return a Hash object that contains ZJIT statistics
 #[unsafe(no_mangle)]
-pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE) -> VALUE {
+pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE, target_key: VALUE) -> VALUE {
     if !zjit_enabled_p() {
         return Qnil;
     }
 
-    fn set_stat(hash: VALUE, key: &str, value: u64) {
-        unsafe { rb_hash_aset(hash, rust_str_to_sym(key), VALUE::fixnum_from_usize(value as usize)); }
+    macro_rules! set_stat {
+        ($hash:ident, $key:expr, $value:expr) => {
+            let key = rust_str_to_sym($key);
+            // Evaluate $value only when it's needed
+            if key == target_key {
+                return VALUE::fixnum_from_usize($value as usize);
+            } else if $hash != Qnil {
+                #[allow(unused_unsafe)]
+                unsafe { rb_hash_aset($hash, key, VALUE::fixnum_from_usize($value as usize)); }
+            }
+        }
     }
 
-    let hash = unsafe { rb_hash_new() };
+    let hash = if target_key.nil_p() {
+        unsafe { rb_hash_new() }
+    } else {
+        Qnil
+    };
     let counters = ZJITState::get_counters();
 
     for &counter in DEFAULT_COUNTERS {
-        let counter_val = unsafe { *counter_ptr(counter) };
-        set_stat(hash, &counter.name(), counter_val);
+        set_stat!(hash, &counter.name(), unsafe { *counter_ptr(counter) });
     }
 
     // Set counters that are enabled when --zjit-stats is enabled
     if get_option!(stats) {
-        set_stat(hash, "zjit_insns_count", counters.zjit_insns_count);
+        set_stat!(hash, "zjit_insns_count", counters.zjit_insns_count);
 
         if unsafe { rb_vm_insns_count } > 0 {
-            set_stat(hash, "vm_insns_count", unsafe { rb_vm_insns_count });
+            set_stat!(hash, "vm_insns_count", unsafe { rb_vm_insns_count });
         }
     }