ZJIT: Add patchpoint for tracing activation

Author: st0012

test/ruby/test_zjit.rb

@@ -2240,6 +2240,30 @@ def foo = 1
     }
   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
+    }
+  end
+
   def test_opt_case_dispatch
     assert_compiles '[true, false]', %q{
       def test(x)

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_tracing_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::NoTracing => {
+                let side_exit_ptr = cb.resolve_label(label);
+                track_no_tracing_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/hir.rs

@@ -134,6 +134,8 @@ pub enum Invariant {
     StableConstantNames {
         idlist: *const ID,
     },
+    /// TracePoint is enabled
+    NoTracing,
     /// There is one ractor running. If a non-root ractor gets spawned, this is invalidated.
     SingleRactorMode,
 }
@@ -247,6 +249,7 @@ impl<'a> std::fmt::Display for InvariantPrinter<'a> {
                 }
                 write!(f, ")")
             }
+            Invariant::NoTracing => write!(f, "NoTracing"),
             Invariant::SingleRactorMode => write!(f, "SingleRactorMode"),
         }
     }
@@ -3399,6 +3402,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let args = state.stack_pop_n(argc as usize)?;
                     let recv = state.stack_pop()?;
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoTracing, state: exit_id });
                     let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: exit_id });
                     state.stack_push(send);
                 }
@@ -3422,6 +3426,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
 
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
                     let recv = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) });
+                    fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoTracing, state: exit_id });
                     let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: exit_id });
                     state.stack_push(send);
                 }
@@ -3478,6 +3483,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                     let args = state.stack_pop_n(argc as usize)?;
                     let recv = state.stack_pop()?;
                     let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state });
+                    fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoTracing, state: exit_id });
                     let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: exit_id });
                     state.stack_push(send);
                 }
@@ -3575,6 +3581,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         .get_builtin_properties(&bf)
                         .map(|props| props.return_type);
 
+                    fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoTracing, state: exit_id });
                     let insn_id = fun.push_insn(block, Insn::InvokeBuiltin {
                         bf,
                         args,
@@ -3601,6 +3608,7 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         .get_builtin_properties(&bf)
                         .map(|props| props.return_type);
 
+                    fun.push_insn(block, Insn::PatchPoint { invariant: Invariant::NoTracing, state: exit_id });
                     let insn_id = fun.push_insn(block, Insn::InvokeBuiltin {
                         bf,
                         args,

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 C function call/return tracing is disabled
+    no_tracing_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_tracing_assumption(patch_point_ptr: CodePtr, side_exit_ptr: CodePtr, payload_ptr: *mut IseqPayload) {
+    let invariants = ZJITState::get_invariants();
+    invariants.no_tracing_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,13 @@ 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_tracing_patch_points);
+
+        // Invalidate all patch points for no tracing
+        compile_patch_points!(cb, patch_points, "TracePoint is enabled, invalidating no tracing assumption");
+
+        cb.mark_all_executable();
     });
 }