From 2f1c30cd50e464880e44da670d3ad8ebe00fc899 Mon Sep 17 00:00:00 2001
From: Aiden Fox Ivey <aiden.foxivey@shopify.com>
Date: Tue, 30 Sep 2025 11:55:33 -0400
Subject: ZJIT: Add --zjit-trace-exits (#14640)

Add side exit tracing functionality for ZJIT
---
 zjit.rb | 114 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 113 insertions(+), 1 deletion(-)

(limited to 'zjit.rb')

diff --git a/zjit.rb b/zjit.rb
index 4438a10c75..8289846c03 100644
--- a/zjit.rb
+++ b/zjit.rb
@@ -9,7 +9,10 @@
 module RubyVM::ZJIT
   # Avoid calling a Ruby method here to avoid interfering with compilation tests
   if Primitive.rb_zjit_print_stats_p
-    at_exit { print_stats }
+    at_exit {
+      print_stats
+      dump_locations
+    }
   end
 end
 
@@ -19,6 +22,106 @@ class << RubyVM::ZJIT
     Primitive.cexpr! 'RBOOL(rb_zjit_enabled_p)'
   end
 
+  # Check if `--zjit-trace-exits` is used
+  def trace_exit_locations_enabled?
+    Primitive.rb_zjit_trace_exit_locations_enabled_p
+  end
+
+  # If --zjit-trace-exits is enabled parse the hashes from
+  # Primitive.rb_zjit_get_exit_locations into a format readable
+  # by Stackprof. This will allow us to find the exact location of a
+  # side exit in ZJIT based on the instruction that is exiting.
+  def exit_locations
+    return unless trace_exit_locations_enabled?
+
+    results = Primitive.rb_zjit_get_exit_locations
+    raw_samples = results[:raw].dup
+    line_samples = results[:lines].dup
+    frames = results[:frames].dup
+    samples_count = 0
+
+    frames.each do |frame_id, frame|
+      frame[:samples] = 0
+      frame[:edges] = {}
+    end
+
+    # Loop through the instructions and set the frame hash with the data.
+    # We use nonexistent.def for the file name, otherwise insns.def will be displayed
+    # and that information isn't useful in this context.
+    RubyVM::INSTRUCTION_NAMES.each_with_index do |name, frame_id|
+      frame_hash = { samples: 0, total_samples: 0, edges: {}, name: name, file: "nonexistent.def", line: nil }
+      results[:frames][frame_id] = frame_hash
+      frames[frame_id] = frame_hash
+    end
+
+    # Loop through the raw_samples and build the hashes for StackProf.
+    # The loop is based off an example in the StackProf documentation and therefore
+    # this functionality can only work with that library.
+    while raw_samples.length > 0
+      stack_trace = raw_samples.shift(raw_samples.shift + 1)
+      lines = line_samples.shift(line_samples.shift + 1)
+      prev_frame_id = nil
+
+      stack_trace.each_with_index do |frame_id, idx|
+        if prev_frame_id
+          prev_frame = frames[prev_frame_id]
+          prev_frame[:edges][frame_id] ||= 0
+          prev_frame[:edges][frame_id] += 1
+        end
+
+        frame_info = frames[frame_id]
+        frame_info[:total_samples] ||= 0
+        frame_info[:total_samples] += 1
+
+        frame_info[:lines] ||= {}
+        frame_info[:lines][lines[idx]] ||= [0, 0]
+        frame_info[:lines][lines[idx]][0] += 1
+
+        prev_frame_id = frame_id
+      end
+
+      top_frame_id = stack_trace.last
+      top_frame_line = 1
+
+      frames[top_frame_id][:samples] += 1
+      frames[top_frame_id][:lines] ||= {}
+      frames[top_frame_id][:lines][top_frame_line] ||= [0, 0]
+      frames[top_frame_id][:lines][top_frame_line][1] += 1
+
+      samples_count += raw_samples.shift
+      line_samples.shift
+    end
+
+    results[:samples] = samples_count
+    # Set missed_samples and gc_samples to 0 as their values
+    # don't matter to us in this context.
+    results[:missed_samples] = 0
+    results[:gc_samples] = 0
+    results
+  end
+
+  # Marshal dumps exit locations to the given filename.
+  #
+  # Usage:
+  #
+  # In a script call:
+  #
+  #   RubyVM::ZJIT.dump_exit_locations("my_file.dump")
+  #
+  # Then run the file with the following options:
+  #
+  #   ruby --zjit --zjit-stats --zjit-trace-exits test.rb
+  #
+  # Once the code is done running, use Stackprof to read the dump file.
+  # See Stackprof documentation for options.
+  def dump_exit_locations(filename)
+    unless trace_exit_locations_enabled?
+      raise ArgumentError, "--zjit-trace-exits must be enabled to use dump_exit_locations."
+    end
+
+    File.write(filename, Marshal.dump(RubyVM::ZJIT.exit_locations))
+  end
+
   # Check if `--zjit-stats` is used
   def stats_enabled?
     Primitive.rb_zjit_stats_enabled_p
@@ -148,4 +251,13 @@ class << RubyVM::ZJIT
   def print_stats
     $stderr.write stats_string
   end
+
+  def dump_locations # :nodoc:
+    return unless trace_exit_locations_enabled?
+
+    filename = "zjit_exit_locations.dump"
+    dump_exit_locations(filename)
+
+    $stderr.puts("ZJIT exit locations dumped to `#{filename}`.")
+  end
 end
-- 
cgit v1.2.3