ZJIT: ::RubyVM::ZJIT.induce_side_exit! and induce_compile_failure!

Tells ZJIT to do a side exit or to fail to compile, useful testing and
for bug reports. We are picky about the syntactic form so we can tell
where the call lands early in the compiler pipeline. The `::` prefix
allows us to interpret it without needing to know under what lexical
scope the iseq will run.

Special semantics with ZJIT doesn't interfere with ruby level semantics;
running these methods do nothing in all modes. So it's no problem to
call these without ZJIT.

Author: XrXr

zjit.c

@@ -19,6 +19,7 @@
 #include "vm_insnhelper.h"
 #include "probes.h"
 #include "probes_helper.h"
+#include "constant.h"
 #include "iseq.h"
 #include "ruby/debug.h"
 #include "internal/cont.h"

zjit.rb

@@ -44,6 +44,20 @@ def trace_exit_locations_enabled?
     Primitive.rb_zjit_trace_exit_locations_enabled_p
   end
 
+  # A directive for the compiler to fail to compile the call to this method.
+  # To show this to ZJIT, say `::RubyVM::ZJIT.induce_compile_failure!` verbatim.
+  # Other forms are too dynamic to detect during compilation.
+  #
+  # Actually running this method does nothing, whether ZJIT sees the call or not.
+  def induce_compile_failure! = nil
+
+  # A directive for the compiler to exit out of compiled code at the call site of this method.
+  # To show this to ZJIT, say `::RubyVM::ZJIT.induce_side_exit!` verbatim.
+  # Other forms are too dynamic to detect during compilation.
+  #
+  # Actually running this method does nothing, whether ZJIT sees the call or not.
+  def induce_side_exit! = nil
+
   # 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

zjit/bindgen/src/main.rs

@@ -145,6 +145,7 @@ fn main() {
         .allowlist_function("rb_gc_location")
         .allowlist_function("rb_gc_writebarrier")
         .allowlist_function("rb_gc_writebarrier_remember")
+        .allowlist_function("rb_gc_register_mark_object")
         .allowlist_function("rb_zjit_writebarrier_check_immediate")
 
         // VALUE variables for Ruby class objects
@@ -432,6 +433,7 @@ fn main() {
         .allowlist_function("rb_vm_base_ptr")
         .allowlist_function("rb_ec_stack_check")
         .allowlist_function("rb_vm_top_self")
+        .allowlist_function("rb_const_lookup")
 
         // We define these manually, don't import them
         .blocklist_type("VALUE")

zjit/src/cruby.rs

@@ -1544,14 +1544,18 @@ pub(crate) mod ids {
         name: RUBY_FL_FREEZE
         name: RUBY_ELTS_SHARED
         name: VM_FRAME_FLAG_MODIFIED_BLOCK_PARAM
+        name: RubyVM
+        name: ZJIT
+        name: induce_side_exit_bang       content: b"induce_side_exit!"
+        name: induce_compile_failure_bang content: b"induce_compile_failure!"
     }
 
     /// Get an CRuby `ID` to an interned string, e.g. a particular method name.
     macro_rules! ID {
         ($id_name:ident) => {{
             let id = $crate::cruby::ids::$id_name.load(std::sync::atomic::Ordering::Relaxed);
             debug_assert_ne!(0, id, "ids module should be initialized");
-            ID(id)
+            $crate::cruby::ID(id)
         }}
     }
     pub(crate) use ID;

zjit/src/cruby_bindings.inc.rs

@@ -7,10 +7,12 @@
 #![allow(clippy::match_like_matches_macro)]
 use crate::{
     backend::lir::C_ARG_OPNDS,
-    cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, invariants::{self, has_singleton_class_of}, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState, json::Json
+    cast::IntoUsize, codegen::local_idx_to_ep_offset, cruby::*, invariants::{self, has_singleton_class_of}, payload::{get_or_create_iseq_payload, IseqPayload}, options::{debug, get_option, DumpHIR}, state::ZJITState, json::Json,
+    state,
 };
 use std::{
-    cell::RefCell, collections::{BTreeSet, HashMap, HashSet, VecDeque}, ffi::{c_void, c_uint, c_int, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter
+    cell::RefCell, collections::{BTreeSet, HashMap, HashSet, VecDeque}, ffi::{c_void, c_uint, c_int, CStr}, fmt::Display, mem::{align_of, size_of}, ptr, slice::Iter,
+    sync::atomic::Ordering,
 };
 use crate::hir_type::{Type, types};
 use crate::hir_effect::{Effect, abstract_heaps, effects};
@@ -529,6 +531,7 @@ pub enum SideExitReason {
     SplatKwNotNilOrHash,
     SplatKwPolymorphic,
     SplatKwNotProfiled,
+    DirectiveInduced,
 }
 
 #[derive(Debug, Clone, Copy)]
@@ -6547,6 +6550,7 @@ pub enum ParseError {
     MalformedIseq(u32), // insn_idx into iseq_encoded
     Validation(ValidationError),
     NotAllowed,
+    DirectiveInduced,
 }
 
 /// Return the number of locals in the current ISEQ (includes parameters)
@@ -7077,7 +7081,32 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                 }
                 YARVINSN_opt_getconstant_path => {
                     let ic = get_arg(pc, 0).as_ptr();
-                    state.stack_push(fun.push_insn(block, Insn::GetConstantPath { ic, state: exit_id }));
+                    let get_const_path = fun.push_insn(block, Insn::GetConstantPath { ic, state: exit_id });
+                    state.stack_push(get_const_path);
+
+                    // Check for `::RubyVM::ZJIT` for directives
+                    unsafe {
+                        let mut current_segment = (*ic).segments;
+                        let mut segments = [ID(0); 4];
+                        for segment in segments.iter_mut() {
+                            *segment = current_segment.read();
+                            if *segment == ID(0) {
+                                break;
+                            }
+                            current_segment = current_segment.add(1);
+                        }
+                        if [ID!(NULL), ID!(RubyVM), ID!(ZJIT), ID(0)] == segments {
+                            debug_assert_ne!(ID!(NULL), ID(0));
+                            let ruby_vm_mod = rb_const_lookup(rb_cObject, ID!(RubyVM));
+                            if !ruby_vm_mod.is_null() && (*ruby_vm_mod).value == rb_cRubyVM {
+                                let zjit_module = VALUE(state::ZJIT_MODULE.load(Ordering::Relaxed));
+                                let lookedup_module = rb_const_lookup(rb_cRubyVM, ID!(ZJIT));
+                                if !lookedup_module.is_null() && (*lookedup_module).value == zjit_module {
+                                    fun.insn_types[get_const_path.0] = Type::from_value(zjit_module);
+                                }
+                            }
+                        }
+                    }
                 }
                 YARVINSN_branchunless | YARVINSN_branchunless_without_ints => {
                     if opcode == YARVINSN_branchunless {
@@ -7538,6 +7567,28 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result<Function, ParseError> {
                         break;  // End the block
                     }
                     let argc = unsafe { vm_ci_argc((*cd).ci) };
+                    let mid = unsafe { rb_vm_ci_mid(call_info) };
+
+                    // Check for calls to directives
+                    if argc == 0
+                        && (mid == ID!(induce_side_exit_bang) || mid == ID!(induce_compile_failure_bang))
+                        && fun.type_of(state.stack_top()?)
+                              .ruby_object()
+                              .is_some_and(|obj| obj == VALUE(state::ZJIT_MODULE.load(Ordering::Relaxed)))
+                    {
+
+                        if mid == ID!(induce_side_exit_bang)
+                            && state::zjit_module_method_match_serial(ID!(induce_side_exit_bang), &state::INDUCE_SIDE_EXIT_SERIAL)
+                        {
+                            fun.push_insn(block, Insn::SideExit { state: exit_id, reason: SideExitReason::DirectiveInduced });
+                            break;  // End the block
+                        }
+                        if mid == ID!(induce_compile_failure_bang)
+                            && state::zjit_module_method_match_serial(ID!(induce_compile_failure_bang), &state::INDUCE_COMPILE_FAILURE_SERIAL)
+                        {
+                            return Err(ParseError::DirectiveInduced);
+                        }
+                    }
 
                     {
                         fn new_branch_block(

zjit/src/hir.rs

@@ -4793,7 +4793,132 @@ pub mod hir_build_tests {
           Jump bb8(v67, v78)
         ");
     }
- }
+
+    #[test]
+    fn test_induce_side_exit() {
+        eval("
+          class NonTopLexicalScope
+            RubyVM = 0
+            def test
+              RubyVM::ZJIT.induce_side_exit! # lexical scope dependant -- should not recognize
+              ::RubyVM::ZJIT.induce_side_exit!
+            end
+          end
+        ");
+        assert_snapshot!(hir_string_proc("NonTopLexicalScope.instance_method(:test)"), @"
+        fn test@<compiled>:5:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          v10:BasicObject = GetConstantPath 0x1000
+          v12:BasicObject = Send v10, :induce_side_exit! # SendFallbackReason: Uncategorized(opt_send_without_block)
+          v16:BasicObject = GetConstantPath 0x1000
+          SideExit DirectiveInduced
+        ");
+    }
+
+    #[test]
+    fn test_induce_side_exit_sensitive_to_constant_state() {
+        eval("
+          def test = ::RubyVM::ZJIT.induce_side_exit!
+        ");
+        assert!(hir_string("test").contains("SideExit DirectiveInduced"));
+        eval("
+          class RubyVM
+            remove_const(:ZJIT)
+          end
+        ");
+        let hir_after_removal = hir_string("test");
+        assert_eq!(false, hir_string("test").contains("SideExit DirectiveInduced"), "should not work when the constant lookup would fail");
+        assert_snapshot!(hir_after_removal, @"
+        fn test@<compiled>:2:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          v10:BasicObject = GetConstantPath 0x1000
+          v12:BasicObject = Send v10, :induce_side_exit! # SendFallbackReason: Uncategorized(opt_send_without_block)
+          CheckInterrupts
+          Return v12
+        ");
+    }
+
+    #[test]
+    fn test_induce_side_exit_doesnt_work_when_method_after_undef() {
+        eval("
+          class << RubyVM::ZJIT
+            undef :induce_side_exit!
+          end
+          def test = ::RubyVM::ZJIT.induce_side_exit!
+        ");
+        assert_eq!(false, hir_string("test").contains("SideExit DirectiveInduced"), "should not work after undef");
+    }
+
+    #[test]
+    fn test_induce_compile_failure_does_not_trigger_autoload() {
+        eval("
+          class RubyVM
+            remove_const(:ZJIT)
+            autoload :ZJIT, 'a-file-that-does-not-exist-as-a-trap'
+          end
+          def test = ::RubyVM::ZJIT.induce_compile_failure!
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:6:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          v10:BasicObject = GetConstantPath 0x1000
+          v12:BasicObject = Send v10, :induce_compile_failure! # SendFallbackReason: Uncategorized(opt_send_without_block)
+          CheckInterrupts
+          Return v12
+        ");
+    }
+
+    #[test]
+    fn test_induce_compile_failure_checks_full_const_path() {
+        eval("def test = ::RubyVM::ZJIT::TooDeep.induce_compile_failure!");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:1:
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          v10:BasicObject = GetConstantPath 0x1000
+          v12:BasicObject = Send v10, :induce_compile_failure! # SendFallbackReason: Uncategorized(opt_send_without_block)
+          CheckInterrupts
+          Return v12
+        ");
+    }
+
+    #[test]
+    fn test_induce_compile_failure() {
+        eval("def test = ::RubyVM::ZJIT.induce_compile_failure!");
+        assert_compile_fails("test", ParseError::DirectiveInduced);
+    }
+}
 
  /// Test successor and predecessor set computations.
  #[cfg(test)]

zjit/src/hir/tests.rs

@@ -3,11 +3,14 @@
 use crate::codegen::{gen_entry_trampoline, gen_exit_trampoline, gen_exit_trampoline_with_counter, gen_function_stub_hit_trampoline};
 use crate::cruby::{self, rb_bug_panic_hook, rb_vm_insn_count, src_loc, EcPtr, Qnil, Qtrue, rb_vm_insn_addr2opcode, rb_profile_frames, VALUE, VM_INSTRUCTION_SIZE, size_t, rb_gc_mark, with_vm_lock, rust_str_to_id, rb_funcallv, rb_const_get, rb_cRubyVM};
 use crate::cruby_methods;
+use cruby::{ID, rb_callable_method_entry, get_def_method_serial, rb_gc_register_mark_object};
+use std::sync::atomic::Ordering;
 use crate::invariants::Invariants;
 use crate::asm::CodeBlock;
 use crate::options::{get_option, rb_zjit_prepare_options};
 use crate::stats::{Counters, InsnCounters, SideExitLocations};
 use crate::virtualmem::CodePtr;
+use std::sync::atomic::AtomicUsize;
 use std::collections::HashMap;
 use std::ptr::null;
 
@@ -301,6 +304,25 @@ impl ZJITState {
     }
 }
 
+/// The `::RubyVM::ZJIT` module.
+pub static ZJIT_MODULE: AtomicUsize = AtomicUsize::new(!0);
+/// Serial of the canonical version of `` right after VM boot.
+pub static INDUCE_SIDE_EXIT_SERIAL: AtomicUsize = AtomicUsize::new(!0);
+/// Serial of the canonical version of `` right after VM boot.
+pub static INDUCE_COMPILE_FAILURE_SERIAL: AtomicUsize = AtomicUsize::new(!0);
+
+/// Check if a method, `method_id`, curerntly exists on `ZJIT.singleton_class` and has the `expected_serial`.
+pub fn zjit_module_method_match_serial(method_id: ID, expected_serial: &AtomicUsize) -> bool {
+    let zjit_module_singleton = VALUE(ZJIT_MODULE.load(Ordering::Relaxed)).class_of();
+    let cme = unsafe { rb_callable_method_entry(zjit_module_singleton, method_id) };
+    if cme.is_null() {
+        false
+    } else {
+        let serial = unsafe { get_def_method_serial((*cme).def) };
+        serial == expected_serial.load(std::sync::atomic::Ordering::Relaxed)
+    }
+}
+
 /// Initialize IDs and annotate builtin C method entries.
 /// Must be called at boot before ruby_init_prelude() since the prelude
 /// could redefine core methods (e.g. Kernel.prepend via bundler).
@@ -318,6 +340,25 @@ pub extern "C" fn rb_zjit_init_builtin_cmes() {
     let method_annotations = cruby_methods::init();
 
     unsafe { ZJIT_STATE = Initialized(method_annotations); }
+
+    // Boot time setup for compiler directives
+    unsafe {
+        let zjit_module = rb_const_get(rb_cRubyVM, rust_str_to_id("ZJIT"));
+
+        let cme = rb_callable_method_entry(zjit_module.class_of(), ID!(induce_side_exit_bang));
+        assert!(! cme.is_null(), "RubyVM::ZJIT.induce_side_exit! should exist on boot");
+        let serial = get_def_method_serial((*cme).def) ;
+        INDUCE_SIDE_EXIT_SERIAL.store(serial, Ordering::Relaxed);
+
+        let cme = rb_callable_method_entry(zjit_module.class_of(), ID!(induce_compile_failure_bang));
+        assert!(! cme.is_null(), "RubyVM::ZJIT.induce_compile_failure! should exist on boot");
+        let serial = get_def_method_serial((*cme).def) ;
+        INDUCE_COMPILE_FAILURE_SERIAL.store(serial, Ordering::Relaxed);
+
+        // Root and pin the module since we'll be doing object identity comparisons.
+        ZJIT_MODULE.store(zjit_module.0, Ordering::Relaxed);
+        rb_gc_register_mark_object(zjit_module);
+    }
 }
 
 /// Initialize ZJIT at boot. This is called even if ZJIT is disabled.