7 files changed, 265 insertions, 0 deletions

diff --git a/yjit/bindgen/src/main.rs b/yjit/bindgen/src/main.rs
index c3a18c9c6f..212013d70c 100644
--- a/yjit/bindgen/src/main.rs
+++ b/yjit/bindgen/src/main.rs
@@ -259,6 +259,7 @@ fn main() {
         .allowlist_function("rb_yjit_obj_written")
         .allowlist_function("rb_yjit_str_simple_append")
         .allowlist_function("rb_ENCODING_GET")
+        .allowlist_function("rb_yjit_exit_locations_dict")
 
         // from vm_sync.h
         .allowlist_function("rb_vm_barrier")
@@ -293,6 +294,9 @@ fn main() {
         .allowlist_function("rb_class_allocate_instance")
         .allowlist_function("rb_obj_info")
 
+        // From include/ruby/debug.h
+        .allowlist_function("rb_profile_frames")
+
         // We define VALUE manually, don't import it
         .blocklist_type("VALUE")
 
diff --git a/yjit/src/codegen.rs b/yjit/src/codegen.rs
index 4203196e34..4d7e7344ad 100644
--- a/yjit/src/codegen.rs
+++ b/yjit/src/codegen.rs
@@ -424,6 +424,13 @@ fn gen_exit(exit_pc: *mut VALUE, ctx: &Context, cb: &mut CodeBlock) -> CodePtr {
     if get_option!(gen_stats) {
         mov(cb, RDI, const_ptr_opnd(exit_pc as *const u8));
         call_ptr(cb, RSI, rb_yjit_count_side_exit_op as *const u8);
+
+        // If --yjit-trace-exits option is enabled, record the exit stack
+        // while recording the side exits.
+        if get_option!(gen_trace_exits) {
+            mov(cb, C_ARG_REGS[0], const_ptr_opnd(exit_pc as *const u8));
+            call_ptr(cb, REG0, rb_yjit_record_exit_stack as *const u8);
+        }
     }
 
     pop(cb, REG_SP);
diff --git a/yjit/src/cruby_bindings.inc.rs b/yjit/src/cruby_bindings.inc.rs
index 758a8de2fd..b5dd356aef 100644
--- a/yjit/src/cruby_bindings.inc.rs
+++ b/yjit/src/cruby_bindings.inc.rs
@@ -968,12 +968,27 @@ extern "C" {
     pub fn rb_vm_barrier();
 }
 extern "C" {
+    pub fn rb_profile_frames(
+        start: ::std::os::raw::c_int,
+        limit: ::std::os::raw::c_int,
+        buff: *mut VALUE,
+        lines: *mut ::std::os::raw::c_int,
+    ) -> ::std::os::raw::c_int;
+}
+extern "C" {
     pub fn rb_yjit_mark_writable(mem_block: *mut ::std::os::raw::c_void, mem_size: u32);
 }
 extern "C" {
     pub fn rb_yjit_mark_executable(mem_block: *mut ::std::os::raw::c_void, mem_size: u32);
 }
 extern "C" {
+    pub fn rb_yjit_exit_locations_dict(
+        yjit_raw_samples: *mut VALUE,
+        yjit_line_samples: *mut ::std::os::raw::c_int,
+        samples_len: ::std::os::raw::c_int,
+    ) -> VALUE;
+}
+extern "C" {
     pub fn rb_yjit_get_page_size() -> u32;
 }
 extern "C" {
diff --git a/yjit/src/invariants.rs b/yjit/src/invariants.rs
index e457f2ea80..6329c70f87 100644
--- a/yjit/src/invariants.rs
+++ b/yjit/src/invariants.rs
@@ -360,6 +360,11 @@ pub extern "C" fn rb_yjit_constant_state_changed(id: ID) {
 /// See `struct yjijt_root_struct` in C.
 #[no_mangle]
 pub extern "C" fn rb_yjit_root_mark() {
+    // Call rb_gc_mark on exit location's raw_samples to
+    // wrap frames in a GC allocated object. This needs to be called
+    // at the same time as root mark.
+    YjitExitLocations::gc_mark_raw_samples();
+
     // Comment from C YJIT:
     //
     // Why not let the GC move the cme keys in this table?
diff --git a/yjit/src/options.rs b/yjit/src/options.rs
index 07c501a108..704c709bae 100644
--- a/yjit/src/options.rs
+++ b/yjit/src/options.rs
@@ -24,6 +24,9 @@ pub struct Options {
     // Capture and print out stats
     pub gen_stats: bool,
 
+    // Trace locations of exits
+    pub gen_trace_exits: bool,
+
     /// Dump compiled and executed instructions for debugging
     pub dump_insns: bool,
 
@@ -45,6 +48,7 @@ pub static mut OPTIONS: Options = Options {
     no_type_prop: false,
     max_versions: 4,
     gen_stats: false,
+    gen_trace_exits: false,
     dump_insns: false,
     verify_ctx: false,
     global_constant_state: false,
@@ -104,6 +108,7 @@ pub fn parse_option(str_ptr: *const std::os::raw::c_char) -> Option<()> {
         ("greedy-versioning", "") => unsafe { OPTIONS.greedy_versioning = true },
         ("no-type-prop", "") => unsafe { OPTIONS.no_type_prop = true },
         ("stats", "") => unsafe { OPTIONS.gen_stats = true },
+        ("trace-exits", "") => unsafe { OPTIONS.gen_trace_exits = true; OPTIONS.gen_stats = true },
         ("dump-insns", "") => unsafe { OPTIONS.dump_insns = true },
         ("verify-ctx", "") => unsafe { OPTIONS.verify_ctx = true },
         ("global-constant-state", "") => unsafe { OPTIONS.global_constant_state = true },
diff --git a/yjit/src/stats.rs b/yjit/src/stats.rs
index 3b7a3f31d9..e129cc2811 100644
--- a/yjit/src/stats.rs
+++ b/yjit/src/stats.rs
@@ -12,6 +12,105 @@ use crate::yjit::yjit_enabled_p;
 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];
 
+/// Global state needed for collecting backtraces of exits
+pub struct YjitExitLocations {
+    /// Vec to hold raw_samples which represent the control frames
+    /// of method entries.
+    raw_samples: Vec<VALUE>,
+    /// Vec to hold line_samples which represent line numbers of
+    /// the iseq caller.
+    line_samples: Vec<i32>
+}
+
+/// Private singleton instance of yjit exit locations
+static mut YJIT_EXIT_LOCATIONS: Option<YjitExitLocations> = None;
+
+impl YjitExitLocations {
+    /// Initialize the yjit exit locations
+    pub fn init() {
+        // Return if the stats feature is disabled
+        if !cfg!(feature = "stats") {
+            return;
+        }
+
+        // Return if --yjit-trace-exits isn't enabled
+        if !get_option!(gen_trace_exits) {
+            return;
+        }
+
+        let yjit_exit_locations = YjitExitLocations {
+            raw_samples: Vec::new(),
+            line_samples: Vec::new()
+        };
+
+        // Initialize the yjit exit locations instance
+        unsafe {
+            YJIT_EXIT_LOCATIONS = Some(yjit_exit_locations);
+        }
+    }
+
+    /// Get a mutable reference to the yjit exit locations globals instance
+    pub fn get_instance() -> &'static mut YjitExitLocations {
+        unsafe { YJIT_EXIT_LOCATIONS.as_mut().unwrap() }
+    }
+
+    /// Get a mutable reference to the yjit raw samples Vec
+    pub fn get_raw_samples() -> &'static mut Vec<VALUE> {
+        &mut YjitExitLocations::get_instance().raw_samples
+    }
+
+    /// Get a mutable reference to yjit the line samples Vec.
+    pub fn get_line_samples() -> &'static mut Vec<i32> {
+        &mut YjitExitLocations::get_instance().line_samples
+    }
+
+    /// Mark the data stored in YjitExitLocations::get_raw_samples that needs to be used by
+    /// rb_yjit_add_frame. YjitExitLocations::get_raw_samples are an array of
+    /// VALUE pointers, exit instruction, and number of times we've seen this stack row
+    /// as collected by rb_yjit_record_exit_stack.
+    ///
+    /// These need to have rb_gc_mark called so they can be used by rb_yjit_add_frame.
+    pub fn gc_mark_raw_samples() {
+        // Return if YJIT is not enabled
+        if !yjit_enabled_p() {
+            return;
+        }
+
+        // Return if the stats feature is disabled
+        if !cfg!(feature = "stats") {
+            return;
+        }
+
+        // Return if --yjit-trace-exits isn't enabled
+        if !get_option!(gen_trace_exits) {
+            return;
+        }
+
+        let mut idx: size_t = 0;
+        let yjit_raw_samples = YjitExitLocations::get_raw_samples();
+
+        while idx < yjit_raw_samples.len() as size_t {
+            let num = yjit_raw_samples[idx as usize];
+            let mut i = 0;
+            idx += 1;
+
+            // Mark the yjit_raw_samples at the given index. These represent
+            // the data that needs to be GC'd which are the current frames.
+            while i < i32::from(num) {
+                unsafe { rb_gc_mark(yjit_raw_samples[idx as usize]); }
+                i += 1;
+                idx += 1;
+            }
+
+            // Increase index for exit instruction.
+            idx += 1;
+            // Increase index for bookeeping value (number of times we've seen this
+            // row in a stack).
+            idx += 1;
+        }
+    }
+}
+
 // Macro to declare the stat counters
 macro_rules! make_counters {
     ($($counter_name:ident,)+) => {
@@ -168,6 +267,57 @@ pub extern "C" fn rb_yjit_get_stats(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
     with_vm_lock(src_loc!(), || rb_yjit_gen_stats_dict())
 }
 
+/// Primitive called in yjit.rb
+///
+/// Check if trace_exits generation is enabled. Requires the stats feature
+/// to be enabled.
+#[no_mangle]
+pub extern "C" fn rb_yjit_trace_exit_locations_enabled_p(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+    #[cfg(feature = "stats")]
+    if get_option!(gen_trace_exits) {
+        return Qtrue;
+    }
+
+    return Qfalse;
+}
+
+/// Call the C function to parse the raw_samples and line_samples
+/// into raw, lines, and frames hash for RubyVM::YJIT.exit_locations.
+#[no_mangle]
+pub extern "C" fn rb_yjit_get_exit_locations(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+    // Return if YJIT is not enabled
+    if !yjit_enabled_p() {
+        return Qnil;
+    }
+
+    // Return if the stats feature is disabled
+    if !cfg!(feature = "stats") {
+        return Qnil;
+    }
+
+    // Return if --yjit-trace-exits isn't enabled
+    if !get_option!(gen_trace_exits) {
+        return Qnil;
+    }
+
+    // If the stats feature is enabled, pass yjit_raw_samples and yjit_line_samples
+    // to the C function called rb_yjit_exit_locations_dict for parsing.
+    let yjit_raw_samples = YjitExitLocations::get_raw_samples();
+    let yjit_line_samples = YjitExitLocations::get_line_samples();
+
+    // Assert that the two Vec's are the same length. If they aren't
+    // equal something went wrong.
+    assert_eq!(yjit_raw_samples.len(), yjit_line_samples.len());
+
+    // yjit_raw_samples and yjit_line_samples are the same length so
+    // pass only one of the lengths in the C function.
+    let samples_len = yjit_raw_samples.len() as i32;
+
+    unsafe {
+        rb_yjit_exit_locations_dict(yjit_raw_samples.as_mut_ptr(), yjit_line_samples.as_mut_ptr(), samples_len)
+    }
+}
+
 /// Export all YJIT statistics as a Ruby hash.
 fn rb_yjit_gen_stats_dict() -> VALUE {
     // If YJIT is not enabled, return Qnil
@@ -231,6 +381,83 @@ fn rb_yjit_gen_stats_dict() -> VALUE {
     hash
 }
 
+/// Record the backtrace when a YJIT exit occurs. This functionality requires
+/// that the stats feature is enabled as well as the --yjit-trace-exits option.
+///
+/// This function will fill two Vec's in YjitExitLocations to record the raw samples
+/// 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)
+{
+    // Return if YJIT is not enabled
+    if !yjit_enabled_p() {
+        return;
+    }
+
+    // Return if the stats feature is disabled
+    if !cfg!(feature = "stats") {
+        return;
+    }
+
+    // Return if --yjit-trace-exits isn't enabled
+    if !get_option!(gen_trace_exits) {
+        return;
+    }
+
+    // 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.
+    if 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()) };
+
+        // Use the same buffer size as Stackprof.
+        const BUFF_LEN: usize = 2048;
+
+        // Create 2 array buffers to be used to collect frames and lines.
+        let mut frames_buffer = [VALUE(0 as usize); BUFF_LEN];
+        let mut lines_buffer = [0; BUFF_LEN];
+
+        // Records call frame and line information for each method entry into two
+        // temporary buffers. Returns the number of times we added to the buffer (ie
+        // the length of the stack).
+        //
+        // Call frame info is stored in the frames_buffer, line number information
+        // in the lines_buffer. The first argument is the start point and the second
+        // argument is the buffer limit, set at 2048.
+        let num = unsafe { rb_profile_frames(0, BUFF_LEN as i32, frames_buffer.as_mut_ptr(), lines_buffer.as_mut_ptr()) };
+
+        let mut i = num - 1;
+        let yjit_raw_samples = YjitExitLocations::get_raw_samples();
+        let yjit_line_samples = YjitExitLocations::get_line_samples();
+
+        yjit_raw_samples.push(VALUE(num as usize));
+        yjit_line_samples.push(num);
+
+        while i >= 0 {
+            let frame = frames_buffer[i as usize];
+            let line = lines_buffer[i as usize];
+
+            yjit_raw_samples.push(frame);
+            yjit_line_samples.push(line);
+
+            i -= 1;
+        }
+
+        // Push the insn value into the yjit_raw_samples Vec.
+        yjit_raw_samples.push(VALUE(insn as usize));
+
+        // Push the current line onto the yjit_line_samples Vec. This
+        // points to the line in insns.def.
+        let line = yjit_line_samples.len() - 1;
+        yjit_line_samples.push(line as i32);
+
+        yjit_raw_samples.push(VALUE(1 as usize));
+        yjit_line_samples.push(1);
+    }
+}
+
 /// Primitive called in yjit.rb. Zero out all the counters.
 #[no_mangle]
 pub extern "C" fn rb_yjit_reset_stats_bang(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
diff --git a/yjit/src/yjit.rs b/yjit/src/yjit.rs
index 06e428f34a..192e9753d9 100644
--- a/yjit/src/yjit.rs
+++ b/yjit/src/yjit.rs
@@ -3,6 +3,7 @@ use crate::core::*;
 use crate::cruby::*;
 use crate::invariants::*;
 use crate::options::*;
+use crate::stats::YjitExitLocations;
 
 use std::os::raw;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -54,6 +55,7 @@ pub extern "C" fn rb_yjit_init_rust() {
     let result = std::panic::catch_unwind(|| {
         Invariants::init();
         CodegenGlobals::init();
+        YjitExitLocations::init();
 
         // YJIT enabled and initialized successfully
         YJIT_ENABLED.store(true, Ordering::Release);