ZJIT: Add patchpoint for TracePoint (#14420)

ZJIT: Add patchpoint for TracePoint activation

Co-authored-by: Takashi Kokubun <takashikkbn@gmail.com>
Co-authored-by: Alan Wu <XrXr@users.noreply.github.com>

test/ruby/test_zjit.rb

@@ -2256,6 +2256,52 @@ class TestZJIT < Test::Unit::TestCase
     }
   end
 
+  def test_line_tracepoint_on_c_method
+    assert_compiles '"[[:line, true]]"', %q{
+      events = []
+      events.instance_variable_set(
+        :@tp,
+        TracePoint.new(:line) { |tp| events << [tp.event, tp.lineno] if tp.path == __FILE__ }
+      )
+      def events.to_str
+        @tp.enable; ''
+      end
+
+      # Stay in generated code while enabling tracing
+      def events.compiled(obj)
+        String(obj)
+        @tp.disable; __LINE__
+      end
+
+      line = events.compiled(events)
+      events[0][-1] = (events[0][-1] == line)
+
+      events.to_s # can't dump events as it's a singleton object AND it has a TracePoint instance variable, which also can't be dumped
+    }
+  end
+
+  def test_targeted_line_tracepoint_in_c_method_call
+    assert_compiles '"[true]"', %q{
+      events = []
+      events.instance_variable_set(:@tp, TracePoint.new(:line) { |tp| events << tp.lineno })
+      def events.to_str
+        @tp.enable(target: method(:compiled))
+        ''
+      end
+
+      # Stay in generated code while enabling tracing
+      def events.compiled(obj)
+        String(obj)
+        __LINE__
+      end
+
+      line = events.compiled(events)
+      events[0] = (events[0] == line)
+
+      events.to_s # can't dump events as it's a singleton object AND it has a TracePoint instance variable, which also can't be dumped
+    }
+  end
+
   def test_opt_case_dispatch
     assert_compiles '[true, false]', %q{
       def test(x)

zjit/bindgen/src/main.rs

@@ -330,6 +330,7 @@ fn main() {
         .allowlist_type("vm_check_match_type")
         .allowlist_type("vm_opt_newarray_send_type")
         .allowlist_type("rb_iseq_type")
+        .allowlist_type("rb_event_flag_t")
 
         // From zjit.c
         .allowlist_function("rb_object_shape_count")

zjit/src/codegen.rs

@@ -9,7 +9,7 @@ use std::slice;
 
 use crate::asm::Label;
 use crate::backend::current::{Reg, ALLOC_REGS};
-use crate::invariants::{track_bop_assumption, track_cme_assumption, track_single_ractor_assumption, track_stable_constant_names_assumption};
+use crate::invariants::{track_bop_assumption, track_cme_assumption, track_single_ractor_assumption, track_stable_constant_names_assumption, track_no_trace_point_assumption};
 use crate::gc::{append_gc_offsets, get_or_create_iseq_payload, get_or_create_iseq_payload_ptr, IseqPayload, IseqStatus};
 use crate::state::ZJITState;
 use crate::stats::{exit_counter_for_compile_error, incr_counter, incr_counter_by, CompileError};
@@ -593,6 +593,10 @@ fn gen_patch_point(jit: &mut JITState, asm: &mut Assembler, invariant: &Invarian
                 let side_exit_ptr = cb.resolve_label(label);
                 track_stable_constant_names_assumption(idlist, code_ptr, side_exit_ptr, payload_ptr);
             }
+            Invariant::NoTracePoint => {
+                let side_exit_ptr = cb.resolve_label(label);
+                track_no_trace_point_assumption(code_ptr, side_exit_ptr, payload_ptr);
+            }
             Invariant::SingleRactorMode => {
                 let side_exit_ptr = cb.resolve_label(label);
                 track_single_ractor_assumption(code_ptr, side_exit_ptr, payload_ptr);

zjit/src/cruby_bindings.inc.rs

@@ -162,6 +162,7 @@ pub const RMODULE_IS_REFINEMENT: ruby_rmodule_flags = 8192;
 pub type ruby_rmodule_flags = u32;
 pub const ROBJECT_HEAP: ruby_robject_flags = 65536;
 pub type ruby_robject_flags = u32;
+pub type rb_event_flag_t = u32;
 pub const RUBY_ENCODING_INLINE_MAX: ruby_encoding_consts = 127;
 pub const RUBY_ENCODING_SHIFT: ruby_encoding_consts = 22;
 pub const RUBY_ENCODING_MASK: ruby_encoding_consts = 532676608;

zjit/src/invariants.rs

@@ -54,6 +54,9 @@ pub struct Invariants {
     /// Map from constant ID to patch points that assume the constant hasn't been redefined
     constant_state_patch_points: HashMap<ID, HashSet<PatchPoint>>,
 
+    /// Set of patch points that assume that the TracePoint is not enabled
+    no_trace_point_patch_points: HashSet<PatchPoint>,
+
     /// Set of patch points that assume that the interpreter is running with only one ractor
     single_ractor_patch_points: HashSet<PatchPoint>,
 }
@@ -272,6 +275,15 @@ pub extern "C" fn rb_zjit_before_ractor_spawn() {
     });
 }
 
+pub fn track_no_trace_point_assumption(patch_point_ptr: CodePtr, side_exit_ptr: CodePtr, payload_ptr: *mut IseqPayload) {
+    let invariants = ZJITState::get_invariants();
+    invariants.no_trace_point_patch_points.insert(PatchPoint {
+        patch_point_ptr,
+        side_exit_ptr,
+        payload_ptr,
+    });
+}
+
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_tracing_invalidate_all() {
     use crate::gc::{get_or_create_iseq_payload, IseqStatus};
@@ -291,5 +303,12 @@ pub extern "C" fn rb_zjit_tracing_invalidate_all() {
             payload.status = IseqStatus::NotCompiled;
             unsafe { rb_iseq_reset_jit_func(iseq) };
         });
+
+        let cb = ZJITState::get_code_block();
+        let patch_points = mem::take(&mut ZJITState::get_invariants().no_trace_point_patch_points);
+
+        compile_patch_points!(cb, patch_points, "TracePoint is enabled, invalidating no TracePoint assumption");
+
+        cb.mark_all_executable();
     });
 }