1 files changed, 25 insertions, 49 deletions

diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index ba84b7a549..b23cd91618 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -1,9 +1,8 @@
 //! Everything related to the collection of runtime stats in YJIT
 //! See the --yjit-stats command-line option
 
-use std::alloc::{GlobalAlloc, Layout, System};
 use std::ptr::addr_of_mut;
-use std::sync::atomic::{AtomicUsize, Ordering};
+use std::sync::atomic::Ordering;
 use std::time::Instant;
 use std::collections::HashMap;
 
@@ -12,6 +11,10 @@ use crate::cruby::*;
 use crate::options::*;
 use crate::yjit::{yjit_enabled_p, YJIT_INIT_TIME};
 
+#[cfg(feature = "stats_allocator")]
+#[path = "../../jit/src/lib.rs"]
+mod jit;
+
 /// Running total of how many ISeqs are in the system.
 #[no_mangle]
 pub static mut rb_yjit_live_iseq_count: u64 = 0;
@@ -20,43 +23,9 @@ pub static mut rb_yjit_live_iseq_count: u64 = 0;
 #[no_mangle]
 pub static mut rb_yjit_iseq_alloc_count: u64 = 0;
 
-/// A middleware to count Rust-allocated bytes as yjit_alloc_size.
-#[global_allocator]
-static GLOBAL_ALLOCATOR: StatsAlloc = StatsAlloc { alloc_size: AtomicUsize::new(0) };
-
-pub struct StatsAlloc {
-    alloc_size: AtomicUsize,
-}
-
-unsafe impl GlobalAlloc for StatsAlloc {
-    unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
-        self.alloc_size.fetch_add(layout.size(), Ordering::SeqCst);
-        System.alloc(layout)
-    }
-
-    unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
-        self.alloc_size.fetch_sub(layout.size(), Ordering::SeqCst);
-        System.dealloc(ptr, layout)
-    }
-
-    unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
-        self.alloc_size.fetch_add(layout.size(), Ordering::SeqCst);
-        System.alloc_zeroed(layout)
-    }
-
-    unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
-        if new_size > layout.size() {
-            self.alloc_size.fetch_add(new_size - layout.size(), Ordering::SeqCst);
-        } else if new_size < layout.size() {
-            self.alloc_size.fetch_sub(layout.size() - new_size, Ordering::SeqCst);
-        }
-        System.realloc(ptr, layout, new_size)
-    }
-}
-
 /// The number of bytes YJIT has allocated on the Rust heap.
 pub fn yjit_alloc_size() -> usize {
-    GLOBAL_ALLOCATOR.alloc_size.load(Ordering::SeqCst)
+    jit::GLOBAL_ALLOCATOR.alloc_size.load(Ordering::SeqCst)
 }
 
 /// Mapping of C function / ISEQ name to integer indices
@@ -121,7 +90,9 @@ pub extern "C" fn incr_iseq_counter(idx: usize) {
     iseq_call_count[idx] += 1;
 }
 
-// YJIT exit counts for each instruction type
+/// YJIT exit counts for each instruction type.
+/// Note that `VM_INSTRUCTION_SIZE` is an upper bound and the actual number
+/// of VM opcodes may be different in the build. See [`rb_vm_instruction_size()`]
 const VM_INSTRUCTION_SIZE_USIZE: usize = VM_INSTRUCTION_SIZE as usize;
 static mut EXIT_OP_COUNT: [u64; VM_INSTRUCTION_SIZE_USIZE] = [0; VM_INSTRUCTION_SIZE_USIZE];
 
@@ -277,6 +248,7 @@ pub const DEFAULT_COUNTERS: &'static [Counter] = &[
     Counter::compiled_blockid_count,
     Counter::compiled_block_count,
     Counter::deleted_defer_block_count,
+    Counter::exceptional_entry_escaped_env,
     Counter::compiled_branch_count,
     Counter::compile_time_ns,
     Counter::compilation_failure,
@@ -349,7 +321,6 @@ macro_rules! ptr_to_counter {
         }
     };
 }
-pub(crate) use ptr_to_counter;
 
 // Declare all the counters we track
 make_counters! {
@@ -384,6 +355,7 @@ make_counters! {
     send_iseq_arity_error,
     send_iseq_block_arg_type,
     send_iseq_clobbering_block_arg,
+    send_iseq_block_arg_gc_unsafe,
     send_iseq_complex_discard_extras,
     send_iseq_leaf_builtin_block_arg_block_param,
     send_iseq_kw_splat_non_nil,
@@ -431,6 +403,7 @@ make_counters! {
     invokesuper_megamorphic,
     invokesuper_no_cme,
     invokesuper_no_me,
+    invokesuper_bmethod_zsuper,
     invokesuper_not_iseq_or_cfunc,
     invokesuper_refinement,
     invokesuper_singleton_class,
@@ -518,9 +491,7 @@ make_counters! {
     opt_aset_not_array,
     opt_aset_not_fixnum,
     opt_aset_not_hash,
-
-    opt_aref_with_qundef,
-    opt_aset_with_qundef,
+    opt_aset_frozen,
 
     opt_case_dispatch_megamorphic,
 
@@ -531,6 +502,7 @@ make_counters! {
     expandarray_postarg,
     expandarray_not_array,
     expandarray_to_ary,
+    expandarray_method_missing,
     expandarray_chain_max_depth,
 
     // getblockparam
@@ -630,6 +602,8 @@ make_counters! {
     iseq_stack_too_large,
     iseq_too_long,
 
+    exceptional_entry_escaped_env,
+
     temp_reg_opnd,
     temp_mem_opnd,
     temp_spill,
@@ -792,8 +766,8 @@ fn rb_yjit_gen_stats_dict(key: VALUE) -> VALUE {
         set_stat_usize!(hash, "context_cache_bytes", crate::core::CTX_ENCODE_CACHE_BYTES + crate::core::CTX_DECODE_CACHE_BYTES);
 
         // VM instructions count
-        if rb_vm_insns_count > 0 {
-            set_stat_usize!(hash, "vm_insns_count", rb_vm_insns_count as usize);
+        if rb_vm_insn_count > 0 {
+            set_stat_usize!(hash, "vm_insns_count", rb_vm_insn_count as usize);
         }
 
         set_stat_usize!(hash, "live_iseq_count", rb_yjit_live_iseq_count as usize);
@@ -838,7 +812,8 @@ fn rb_yjit_gen_stats_dict(key: VALUE) -> VALUE {
 
         // For each entry in exit_op_count, add a stats entry with key "exit_INSTRUCTION_NAME"
         // and the value is the count of side exits for that instruction.
-        for op_idx in 0..VM_INSTRUCTION_SIZE_USIZE {
+        use crate::utils::IntoUsize;
+        for op_idx in 0..rb_vm_instruction_size().as_usize() {
             let op_name = insn_name(op_idx);
             let key_string = "exit_".to_owned() + &op_name;
             let count = EXIT_OP_COUNT[op_idx];
@@ -864,8 +839,8 @@ fn rb_yjit_gen_stats_dict(key: VALUE) -> VALUE {
         set_stat_double!(hash, "avg_len_in_yjit", avg_len_in_yjit);
 
         // Proportion of instructions that retire in YJIT
-        if rb_vm_insns_count > 0 {
-            let total_insns_count = retired_in_yjit + rb_vm_insns_count;
+        if rb_vm_insn_count > 0 {
+            let total_insns_count = retired_in_yjit + rb_vm_insn_count;
             set_stat_usize!(hash, "total_insns_count", total_insns_count as usize);
 
             let ratio_in_yjit: f64 = 100.0 * retired_in_yjit as f64 / total_insns_count as f64;
@@ -927,7 +902,7 @@ fn rb_yjit_gen_stats_dict(key: VALUE) -> VALUE {
 /// and line samples. Their length should be the same, however the data stored in
 /// them is different.
 #[no_mangle]
-pub extern "C" fn rb_yjit_record_exit_stack(_exit_pc: *const VALUE)
+pub extern "C" fn rb_yjit_record_exit_stack(exit_pc: *const VALUE)
 {
     // Return if YJIT is not enabled
     if !yjit_enabled_p() {
@@ -951,10 +926,11 @@ pub extern "C" fn rb_yjit_record_exit_stack(_exit_pc: *const VALUE)
     // rb_vm_insn_addr2opcode won't work in cargo test --all-features
     // because it's a C function. Without insn call, this function is useless
     // so wrap the whole thing in a not test check.
+    let _ = exit_pc;
     #[cfg(not(test))]
     {
         // Get the opcode from the encoded insn handler at this PC
-        let insn = unsafe { rb_vm_insn_addr2opcode((*_exit_pc).as_ptr()) };
+        let insn = unsafe { rb_vm_insn_addr2opcode((*exit_pc).as_ptr()) };
 
         // Use the same buffer size as Stackprof.
         const BUFF_LEN: usize = 2048;