YJIT: move --yjit-stats at_exit call into Ruby

This change fixes `-v --yjit-stats`. Previously in this situation,
YJIT._print_stats wasn't defined as yjit.rb is not evaluated when there
is only "-v" and no Ruby code to run.

3 files changed, 17 insertions, 14 deletions

diff --git a/test/ruby/test_yjit.rb b/test/ruby/test_yjit.rb
index 218d68ea4b..9014a1e567 100644
--- a/test/ruby/test_yjit.rb
+++ b/test/ruby/test_yjit.rb
@@ -37,6 +37,11 @@ class TestYJIT < Test::Unit::TestCase
     assert_in_out_err('--yjit-greedy-versioning=1', '', [], /warning: argument to --yjit-greedy-versioning is ignored/)
   end
 
+  def test_yjit_stats_and_v_no_error
+    _stdout, stderr, _status = EnvUtil.invoke_ruby(%w(-v --yjit-stats), '', true, true)
+    refute_includes(stderr, "NoMethodError")
+  end
+
   def test_enable_from_env_var
     yjit_child_env = {'RUBY_YJIT_ENABLE' => '1'}
     assert_in_out_err([yjit_child_env, '--version'], '') do |stdout, stderr|
diff --git a/yjit.rb b/yjit.rb
index cbd2ea1e44..66c4eaf702 100644
--- a/yjit.rb
+++ b/yjit.rb
@@ -142,13 +142,18 @@ module YJIT
   end
 
   def self.stats_enabled?
-    Primitive.cexpr! 'rb_yjit_opts.gen_stats ? Qtrue : Qfalse'
+    Primitive.yjit_stats_enabled_p
   end
 
   def self.enabled?
     Primitive.cexpr! 'rb_yjit_enabled_p() ? Qtrue : Qfalse'
   end
 
+  # Avoid calling a method here to not interfere with compilation tests
+  if Primitive.yjit_stats_enabled_p
+    at_exit { _print_stats }
+  end
+
   class << self
     private
 
diff --git a/yjit_iface.c b/yjit_iface.c
index 870a8d1efd..d647cf1113 100644
--- a/yjit_iface.c
+++ b/yjit_iface.c
@@ -684,14 +684,6 @@ yjit_disasm(VALUE self, VALUE code, VALUE from)
 }
 #endif
 
-static VALUE
-at_exit_print_stats(RB_BLOCK_CALL_FUNC_ARGLIST(yieldarg, data))
-{
-    // Defined in yjit.rb
-    rb_funcall(mYjit, rb_intern("_print_stats"), 0);
-    return Qnil;
-}
-
 // Primitive called in yjit.rb. Export all machine code comments as a Ruby array.
 static VALUE
 comments_for(rb_execution_context_t *ec, VALUE self, VALUE start_address, VALUE end_address)
@@ -723,6 +715,12 @@ comments_for(rb_execution_context_t *ec, VALUE self, VALUE start_address, VALUE
     return comment_array;
 }
 
+static VALUE
+yjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self)
+{
+    return RBOOL(YJIT_STATS && rb_yjit_opts.gen_stats);
+}
+
 // Primitive called in yjit.rb. Export all YJIT statistics as a Ruby hash.
 static VALUE
 get_yjit_stats(rb_execution_context_t *ec, VALUE self)
@@ -1144,11 +1142,6 @@ rb_yjit_init(struct rb_yjit_options *options)
 #endif
 #endif
 
-    if (YJIT_STATS && rb_yjit_opts.gen_stats) {
-        // Setup at_exit callback for printing out counters
-        rb_block_call(rb_mKernel, rb_intern("at_exit"), 0, NULL, at_exit_print_stats, Qfalse);
-    }
-
     // Make dependency tables
     method_lookup_dependency = st_init_numtable();
     cme_validity_dependency = st_init_numtable();