2 files changed, 31 insertions, 0 deletions
diff --git a/zjit/src/lib.rs b/zjit/src/lib.rs
index 9e93f0f2be..e68ac93fa5 100644
--- a/zjit/src/lib.rs
+++ b/zjit/src/lib.rs
@@ -89,6 +89,15 @@ fn rb_bug_panic_hook() {
 /// Generate JIT code for a given ISEQ, which takes EC and CFP as its arguments.
 #[unsafe(no_mangle)]
 pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *const u8 {
+    let code_ptr = iseq_gen_entry_point(iseq);
+    if ZJITState::assert_compiles_enabled() && code_ptr == std::ptr::null() {
+        let iseq_location = iseq_get_location(iseq, 0);
+        panic!("Failed to compile: {iseq_location}");
+    }
+    code_ptr
+}
+
+fn iseq_gen_entry_point(iseq: IseqPtr) -> *const u8 {
     // Do not test the JIT code in HIR tests
     if cfg!(test) {
         return std::ptr::null();
@@ -116,3 +125,10 @@ pub extern "C" fn rb_zjit_iseq_gen_entry_point(iseq: IseqPtr, _ec: EcPtr) -> *co
         }
     })
 }
+
+/// Assert that any future ZJIT compilation will return a function pointer (not fail to compile)
+#[unsafe(no_mangle)]
+pub extern "C" fn rb_zjit_assert_compiles(_ec: EcPtr, _self: VALUE) -> VALUE {
+    ZJITState::enable_assert_compiles();
+    Qnil
+}
diff --git a/zjit/src/state.rs b/zjit/src/state.rs
index 92efc3a48e..c176d58b20 100644
--- a/zjit/src/state.rs
+++ b/zjit/src/state.rs
@@ -12,6 +12,9 @@ pub struct ZJITState {
 
     /// Assumptions that require invalidation
     invariants: Invariants,
+
+    /// Assert successful compilation if set to true
+    assert_compiles: bool,
 }
 
 /// Private singleton instance of the codegen globals
@@ -64,6 +67,7 @@ impl ZJITState {
             code_block: cb,
             options,
             invariants: Invariants::default(),
+            assert_compiles: false,
         };
         unsafe { ZJIT_STATE = Some(zjit_state); }
     }
@@ -92,4 +96,15 @@ impl ZJITState {
     pub fn get_invariants() -> &'static mut Invariants {
         &mut ZJITState::get_instance().invariants
     }
+
+    /// Return true if successful compilation should be asserted
+    pub fn assert_compiles_enabled() -> bool {
+        ZJITState::get_instance().assert_compiles
+    }
+
+    /// Start asserting successful compilation
+    pub fn enable_assert_compiles() {
+        let instance = ZJITState::get_instance();
+        instance.assert_compiles = true;
+    }
 }