ZJIT: Add RubyVM::ZJIT.reset_stats! method

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

Author: rwstauner

test/ruby/test_zjit.rb

@@ -1847,6 +1847,28 @@ def test = 1
     }, 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
+      initial_count = initial_stats[:zjit_insn_count]
+
+      # Reset the stats
+      RubyVM::ZJIT.reset_stats!
+
+      # Get stats after reset
+      reset_stats = RubyVM::ZJIT.stats
+      reset_count = reset_stats[:zjit_insn_count]
+
+      # After reset, counters should be zero or at least much smaller
+      # (some instructions might execute between reset and reading stats)
+      initial_count > 0 && reset_count < initial_count / 10
+    }, 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{

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);
 

zjit.rb

@@ -29,6 +29,11 @@ def stats(target_key = nil)
     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"

zjit/src/stats.rs

@@ -237,6 +237,23 @@ 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
+    for i in 0..(VM_INSTRUCTION_SIZE as usize) {
+        exit_counters[i] = 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 {