ZJIT: Add RubyVM::ZJIT.reset_stats! method (GH-14479)

This allows for more precise tracking of stats programmatically
which is particularly useful for our nightly benchmarking suite.

- Define rust function
- Expose to C
- Wrap with Ruby API
- Add a test

4 files changed, 43 insertions, 0 deletions

diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb
index 280a7503d4..4d8eb30c80 100644
--- a/test/ruby/test_zjit.rb
+++ b/test/ruby/test_zjit.rb
@@ -1878,6 +1878,28 @@ class TestZJIT < Test::Unit::TestCase
     }, 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] },
+      ].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{
diff --git a/zjit.c b/zjit.c
index e1ea6d7e09..5d50775fc7 100644
--- a/zjit.c
+++ b/zjit.c
@@ -342,6 +342,7 @@ rb_zjit_insn_leaf(int insn, const VALUE *opes)
 // 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 target_key);
+VALUE rb_zjit_reset_stats_bang(rb_execution_context_t *ec, VALUE self);
 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);
 
diff --git a/zjit.rb b/zjit.rb
index d70ff1dd47..39b0e5eb22 100644
--- a/zjit.rb
+++ b/zjit.rb
@@ -29,6 +29,11 @@ class << RubyVM::ZJIT
     Primitive.rb_zjit_stats(target_key)
   end
 
+  # Discard statistics collected for `--zjit-stats`.
+  def reset_stats!
+    Primitive.rb_zjit_reset_stats_bang
+  end
+
   # Get the summary of ZJIT statistics as a String
   def stats_string
     buf = +"***ZJIT: Printing ZJIT statistics on exit***\n"
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
index 2e7d743f7b..adfc28b4a8 100644
--- a/zjit/src/stats.rs
+++ b/zjit/src/stats.rs
@@ -237,6 +237,21 @@ pub fn exit_counter_ptr(reason: crate::hir::SideExitReason) -> *mut u64 {
     counter_ptr(counter)
 }
 
+/// Primitive called in zjit.rb. Zero out all the counters.
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_reset_stats_bang(_ec: EcPtr, _self: VALUE) -> VALUE {
+    let counters = ZJITState::get_counters();
+    let exit_counters = ZJITState::get_exit_counters();
+
+    // Reset all counters to zero
+    *counters = Counters::default();
+
+    // Reset exit counters for YARV instructions
+    exit_counters.as_mut_slice().fill(0);
+
+    Qnil
+}
+
 /// Return a Hash object that contains ZJIT statistics
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE, target_key: VALUE) -> VALUE {