ZJIT: Add RubyVM::ZJIT.induce_compile_failure!

Author: XrXr

zjit.c

@@ -321,6 +321,7 @@ VALUE rb_zjit_print_stats_p(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_get_stats_file_path_p(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_trace_exit_locations_enabled_p(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_get_exit_locations(rb_execution_context_t *ec, VALUE self);
+VALUE rb_zjit_induce_compile_failure(rb_execution_context_t *ec, VALUE self);
 VALUE rb_zjit_induce_side_exit(rb_execution_context_t *ec, VALUE self);
 
 // Preprocessed zjit.rb generated during build

zjit.rb

@@ -44,6 +44,14 @@ 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.
+  # Since JIT compilation errors should not meaningfully change the semantics the program,
+  # this method should behave as if it does nothing.
+  def induce_compile_failure!
+    Primitive.attr! :leaf
+    Primitive.rb_zjit_induce_compile_failure
+  end
+
   # A directive for the compiler to exit out of compiled code at the call site to this method.
   # This should not meaningfully change the semantics of the program and should behave as if
   # it does nothing.

zjit/src/codegen.rs

@@ -2634,6 +2634,10 @@ fn compile_iseq(iseq: IseqPtr) -> Result<Function, CompileError> {
         function.optimize();
     }
     function.dump_hir();
+
+    if function.fail_to_compile_at_the_end {
+        return Err(CompileError::DirectiveInduced)
+    }
     Ok(function)
 }
 

zjit/src/cruby_methods.rs

@@ -283,6 +283,7 @@ pub fn init() -> Annotations {
     builtin_funcs.insert(rb_jit_ary_at_end as *mut c_void, FnProperties { inline: inline_ary_at_end, return_type: types::BoolExact, ..Default::default() });
 
     // RubyVM::ZJIT directive methods
+    assert!(builtin_funcs.insert(rb_zjit_induce_compile_failure as *mut c_void, FnProperties { inline: inline_induce_compile_failure, ..Default::default() }).is_none());
     assert!(builtin_funcs.insert(rb_zjit_induce_side_exit as *mut c_void, FnProperties { inline: inline_induce_side_exit, ..Default::default() }).is_none());
 
     Annotations {
@@ -291,6 +292,15 @@ pub fn init() -> Annotations {
     }
 }
 
+// The compiler/linker can give two identical functions the same address, so we do this dance
+// to make sure the function's address is different from rb_zjit_induce_side_exit
+pub static mut NIL: VALUE = Qnil;
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_induce_compile_failure(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
+    std::hint::black_box(&raw const NIL);
+    unsafe { NIL }
+}
+
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_induce_side_exit(_ec: EcPtr, _ruby_self: VALUE) -> VALUE {
     Qnil
@@ -953,6 +963,12 @@ fn inline_ary_at_end(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::I
     Some(result)
 }
 
+/// For `RubyVM::ZJIT.induce_compile_failure!`
+fn inline_induce_compile_failure(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, _args: &[hir::InsnId], _state: hir::InsnId) -> Option<hir::InsnId> {
+    fun.fail_to_compile_at_the_end = true;
+    Some(fun.push_insn(block, hir::Insn::Const { val: hir::Const::Value(Qnil) }))
+}
+
 /// For `RubyVM::ZJIT.induce_side_exit!`
 fn inline_induce_side_exit(fun: &mut hir::Function, block: hir::BlockId, _recv: hir::InsnId, _args: &[hir::InsnId], state: hir::InsnId) -> Option<hir::InsnId> {
     Some(fun.push_insn(block, hir::Insn::SideExit { reason: hir::SideExitReason::DirectiveInduced, state }))

zjit/src/hir.rs

@@ -2302,6 +2302,9 @@ pub struct Function {
     /// fulfilling `(0..=opt_num)` optional parameters.
     jit_entry_blocks: Vec<BlockId>,
     profiles: Option<ProfileOracle>,
+    /// For user directed compilation failures. We still run through the full pipeline before
+    /// failing so to avoid shadowing other failures.
+    pub fail_to_compile_at_the_end: bool,
 }
 
 /// The kind of a value an ISEQ returns
@@ -2401,6 +2404,7 @@ impl Function {
             jit_entry_blocks: vec![],
             param_types: vec![],
             profiles: None,
+            fail_to_compile_at_the_end: false,
         }
     }
 
@@ -6160,6 +6164,9 @@ impl<'a> std::fmt::Display for FunctionPrinter<'a> {
             iseq_name
         };
         writeln!(f, "fn {iseq_name}:")?;
+        if fun.fail_to_compile_at_the_end {
+            writeln!(f, "# Found induce_compile_failure directive")?;
+        }
         for block_id in fun.rpo() {
             if !self.display_snapshot_and_tp_patchpoints && block_id == fun.entries_block {
                 // Unless we're doing --zjit-dump-hir=all, skip the entries superblock — it's an

zjit/src/hir/opt_tests.rs

@@ -14037,4 +14037,37 @@ mod hir_opt_tests {
           SideExit DirectiveInduced
         ");
     }
+
+    #[test]
+    fn test_induced_compile_failure() {
+        eval("
+            def test
+              RubyVM::ZJIT.induce_compile_failure!
+            end
+            test
+            test
+        ");
+        assert_snapshot!(hir_string("test"), @"
+        fn test@<compiled>:3:
+        # Found induce_compile_failure directive
+        bb1():
+          EntryPoint interpreter
+          v1:BasicObject = LoadSelf
+          Jump bb3(v1)
+        bb2():
+          EntryPoint JIT(0)
+          v4:BasicObject = LoadArg :self@0
+          Jump bb3(v4)
+        bb3(v6:BasicObject):
+          PatchPoint SingleRactorMode
+          PatchPoint StableConstantNames(0x1000, RubyVM::ZJIT)
+          v20:ModuleExact[VALUE(0x1008)] = Const Value(VALUE(0x1008))
+          PatchPoint MethodRedefined(Module@0x1010, induce_compile_failure!@0x1018, cme:0x1020)
+          IncrCounter inline_iseq_optimized_send_count
+          v25:NilClass = Const Value(nil)
+          IncrCounter inline_cfunc_optimized_send_count
+          CheckInterrupts
+          Return v25
+        ");
+    }
 }

zjit/src/stats.rs

@@ -330,6 +330,7 @@ make_counters! {
     compile_error_exception_handler,
     compile_error_out_of_memory,
     compile_error_label_linking_failure,
+    compile_error_directive_induced,
     compile_error_jit_to_jit_optional,
     compile_error_register_spill_on_ccall,
     compile_error_register_spill_on_alloc,
@@ -504,6 +505,8 @@ pub enum CompileError {
     /// offsets that don't fit in one instruction. We error in
     /// error that case.
     LabelLinkingFailure,
+    /// Error from `RubyVM::ZJIT.induced_compile_failure!`
+    DirectiveInduced,
 }
 
 /// Return a raw pointer to the exit counter for a given CompileError
@@ -518,6 +521,7 @@ pub fn exit_counter_for_compile_error(compile_error: &CompileError) -> Counter {
         ExceptionHandler        => compile_error_exception_handler,
         OutOfMemory             => compile_error_out_of_memory,
         LabelLinkingFailure     => compile_error_label_linking_failure,
+        DirectiveInduced         => compile_error_directive_induced,
         ParseError(parse_error) => match parse_error {
             StackUnderflow(_)       => compile_error_parse_stack_underflow,
             MalformedIseq(_)        => compile_error_parse_malformed_iseq,