ZJIT: Assert no side exits in assert_compiles()

Most of the time, we want to assert that we compile and the compiled
code runs without exiting. A small number of tests trigger side exits,
and those are changed to use assert_compiles_allowing_exits().

```console
$ rg -F 'assert_compiles(' -o | wc
     289     289   11862
$ rg -F 'assert_compiles_allowing_exits(' -o | wc
      38      38    2196
```

5 files changed, 65 insertions, 39 deletions

diff --git a/zjit/src/backend/lir.rs b/zjit/src/backend/lir.rs
index bbe8b5a4ec..bb8d1e1e73 100644
--- a/zjit/src/backend/lir.rs
+++ b/zjit/src/backend/lir.rs
@@ -2750,12 +2750,12 @@ impl Assembler
                 }).unwrap_or(false);
 
                 // If enabled, instrument exits first, and then jump to a shared exit.
-                let counted_exit = if get_option!(stats) || should_record_exit {
+                let counted_exit = if get_option!(stats) || should_record_exit || cfg!(test) {
                     let counted_exit = self.new_label("counted_exit");
                     self.write_label(counted_exit.clone());
                     asm_comment!(self, "Counted Exit: {reason}");
 
-                    if get_option!(stats) {
+                    if get_option!(stats) || cfg!(test) {
                         asm_comment!(self, "increment a side exit counter");
                         self.incr_counter(Opnd::const_ptr(exit_counter_ptr(reason)), 1.into());
 
diff --git a/zjit/src/codegen_tests.rs b/zjit/src/codegen_tests.rs
index 55149c22ca..474cbee5b5 100644
--- a/zjit/src/codegen_tests.rs
+++ b/zjit/src/codegen_tests.rs
@@ -2119,7 +2119,7 @@ fn test_opt_empty_p() {
         def test(x) = x.empty?
     ");
     assert_contains_opcode("test", YARVINSN_opt_empty_p);
-    assert_snapshot!(assert_compiles("[test([1]), test(\"1\"), test({})]"), @"[false, false, true]");
+    assert_snapshot!(assert_compiles_allowing_exits("[test([1]), test(\"1\"), test({})]"), @"[false, false, true]");
 }
 
 #[test]
@@ -2128,7 +2128,7 @@ fn test_opt_succ() {
         def test(obj) = obj.succ
     ");
     assert_contains_opcode("test", YARVINSN_opt_succ);
-    assert_snapshot!(assert_compiles(r#"[test(-1), test("A")]"#), @r#"[0, "B"]"#);
+    assert_snapshot!(assert_compiles_allowing_exits(r#"[test(-1), test("A")]"#), @r#"[0, "B"]"#);
 }
 
 #[test]
@@ -2137,7 +2137,7 @@ fn test_opt_and() {
         def test(x, y) = x & y
     ");
     assert_contains_opcode("test", YARVINSN_opt_and);
-    assert_snapshot!(assert_compiles("[test(0b1101, 3), test([3, 2, 1, 4], [8, 1, 2, 3])]"), @"[1, [3, 2, 1]]");
+    assert_snapshot!(assert_compiles_allowing_exits("[test(0b1101, 3), test([3, 2, 1, 4], [8, 1, 2, 3])]"), @"[1, [3, 2, 1]]");
 }
 
 #[test]
@@ -2146,7 +2146,7 @@ fn test_opt_or() {
         def test(x, y) = x | y
     ");
     assert_contains_opcode("test", YARVINSN_opt_or);
-    assert_snapshot!(assert_compiles("[test(0b1000, 3), test([3, 2, 1], [1, 2, 3])]"), @"[11, [3, 2, 1]]");
+    assert_snapshot!(assert_compiles_allowing_exits("[test(0b1000, 3), test([3, 2, 1], [1, 2, 3])]"), @"[11, [3, 2, 1]]");
 }
 
 #[test]
@@ -2170,7 +2170,7 @@ fn test_fixnum_and_side_exit() {
         def test(a, b) = a & b
     ");
     assert_contains_opcode("test", YARVINSN_opt_and);
-    assert_snapshot!(assert_compiles("
+    assert_snapshot!(assert_compiles_allowing_exits("
         [
                   test(2, 2),
                   test(0b011, 0b110),
@@ -2200,7 +2200,7 @@ fn test_fixnum_or_side_exit() {
         def test(a, b) = a | b
     ");
     assert_contains_opcode("test", YARVINSN_opt_or);
-    assert_snapshot!(assert_compiles("
+    assert_snapshot!(assert_compiles_allowing_exits("
         [
                   test(1, 2),
                   test(2, 2),
@@ -2273,7 +2273,7 @@ fn test_opt_not() {
         def test(obj) = !obj
     ");
     assert_contains_opcode("test", YARVINSN_opt_not);
-    assert_snapshot!(assert_compiles("[test(nil), test(false), test(0)]"), @"[true, true, false]");
+    assert_snapshot!(assert_compiles_allowing_exits("[test(nil), test(false), test(0)]"), @"[true, true, false]");
 }
 
 #[test]
@@ -2369,7 +2369,7 @@ fn test_opt_newarray_send_include_p_redefined() {
         end
     ");
     assert_contains_opcode("test", YARVINSN_opt_newarray_send);
-    assert_snapshot!(assert_compiles("
+    assert_snapshot!(assert_compiles_allowing_exits("
         def test(x)
           [:y, 1, Object.new].include?(x)
         end
@@ -2404,7 +2404,7 @@ fn test_opt_duparray_send_include_p_redefined() {
         end
     ");
     assert_contains_opcode("test", YARVINSN_opt_duparray_send);
-    assert_snapshot!(assert_compiles("
+    assert_snapshot!(assert_compiles_allowing_exits("
         def test(x)
           [:y, 1].include?(x)
         end
@@ -2441,7 +2441,7 @@ fn test_opt_newarray_send_pack_redefined() {
         end
     "#);
     assert_contains_opcode("test", YARVINSN_opt_newarray_send);
-    assert_snapshot!(assert_compiles(r#"
+    assert_snapshot!(assert_compiles_allowing_exits(r#"
         [test(65), test(66), test(67)]
     "#), @r#"["override:A", "override:B", "override:C"]"#);
 }
@@ -2476,7 +2476,7 @@ fn test_opt_newarray_send_pack_buffer_redefined() {
         end
     "#);
     assert_contains_opcode("test", YARVINSN_opt_newarray_send);
-    assert_snapshot!(assert_compiles(r#"
+    assert_snapshot!(assert_compiles_allowing_exits(r#"
         def test(num, buffer)
           [num].pack('C', buffer:)
         end
@@ -2509,7 +2509,7 @@ fn test_opt_newarray_send_hash_redefined() {
         test(20)
     ");
     assert_contains_opcode("test", YARVINSN_opt_newarray_send);
-    assert_snapshot!(assert_compiles("test(20)"), @"42");
+    assert_snapshot!(assert_compiles_allowing_exits("test(20)"), @"42");
 }
 
 #[test]
@@ -2534,7 +2534,7 @@ fn test_opt_newarray_send_max_redefined() {
         def test(a,b) = [a,b].max
     ");
     assert_contains_opcode("test", YARVINSN_opt_newarray_send);
-    assert_snapshot!(assert_compiles("
+    assert_snapshot!(assert_compiles_allowing_exits("
         def test(a,b) = [a,b].max
         test(15, 30)
         [test(15, 30), test(45, 35)]
@@ -2694,7 +2694,7 @@ fn test_opt_hash_freeze_rewritten() {
         test
     ");
     assert_contains_opcode("test", YARVINSN_opt_hash_freeze);
-    assert_snapshot!(assert_compiles("test"), @"5");
+    assert_snapshot!(assert_compiles_allowing_exits("test"), @"5");
 }
 
 #[test]
@@ -2799,7 +2799,7 @@ fn test_opt_ary_freeze_rewritten() {
         test
     ");
     assert_contains_opcode("test", YARVINSN_opt_ary_freeze);
-    assert_snapshot!(assert_compiles("test"), @"5");
+    assert_snapshot!(assert_compiles_allowing_exits("test"), @"5");
 }
 
 #[test]
@@ -2828,7 +2828,7 @@ fn test_opt_str_freeze_rewritten() {
         test
     ");
     assert_contains_opcode("test", YARVINSN_opt_str_freeze);
-    assert_snapshot!(assert_compiles("test"), @"5");
+    assert_snapshot!(assert_compiles_allowing_exits("test"), @"5");
 }
 
 #[test]
@@ -2857,7 +2857,7 @@ fn test_opt_str_uminus_rewritten() {
         test
     ");
     assert_contains_opcode("test", YARVINSN_opt_str_uminus);
-    assert_snapshot!(assert_compiles("test"), @"5");
+    assert_snapshot!(assert_compiles_allowing_exits("test"), @"5");
 }
 
 #[test]
@@ -2928,7 +2928,7 @@ fn test_array_fixnum_aref_out_of_bounds_positive() {
         test(10)
     ");
     assert_contains_opcode("test", YARVINSN_opt_aref);
-    assert_snapshot!(assert_compiles("test(10)"), @"nil");
+    assert_snapshot!(assert_compiles_allowing_exits("test(10)"), @"nil");
 }
 
 #[test]
@@ -2938,7 +2938,7 @@ fn test_array_fixnum_aref_out_of_bounds_negative() {
         test(-10)
     ");
     assert_contains_opcode("test", YARVINSN_opt_aref);
-    assert_snapshot!(assert_compiles("test(-10)"), @"nil");
+    assert_snapshot!(assert_compiles_allowing_exits("test(-10)"), @"nil");
 }
 
 #[test]
@@ -3666,7 +3666,7 @@ fn test_getivar_t_data_then_string() {
       end
       OBJ.test; OBJ.test # profile and compile for Thread (T_DATA)
     "#);
-    assert_snapshot!(assert_compiles("[STR.test, STR.test]"), @"[1000, 1000]");
+    assert_snapshot!(assert_compiles_allowing_exits("[STR.test, STR.test]"), @"[1000, 1000]");
 }
 
 #[test]
@@ -3694,7 +3694,7 @@ fn test_getivar_t_object_then_string() {
       end
       OBJ.test; OBJ.test # profile and compile for MyObject
     "#);
-    assert_snapshot!(assert_compiles("[STR.test, STR.test]"), @"[1000, 1000]");
+    assert_snapshot!(assert_compiles_allowing_exits("[STR.test, STR.test]"), @"[1000, 1000]");
 }
 
 #[test]
@@ -3725,7 +3725,7 @@ fn test_getivar_t_class_then_string() {
       p MyClass.test; p MyClass.test # profile and compile for MyClass
       p STR.test
     "#);
-    assert_snapshot!(assert_compiles("[STR.test, STR.test]"), @"[1000, 1000]");
+    assert_snapshot!(assert_compiles_allowing_exits("[STR.test, STR.test]"), @"[1000, 1000]");
 }
 
 
@@ -3806,7 +3806,7 @@ fn test_expandarray_splat() {
         test [3, 4]
     ");
     assert_contains_opcode("test", YARVINSN_expandarray);
-    assert_snapshot!(assert_compiles("test [3, 4]"), @"[3, [4]]");
+    assert_snapshot!(assert_compiles_allowing_exits("test [3, 4]"), @"[3, [4]]");
 }
 
 #[test]
@@ -3819,7 +3819,7 @@ fn test_expandarray_splat_post() {
         test [3, 4, 5]
     ");
     assert_contains_opcode("test", YARVINSN_expandarray);
-    assert_snapshot!(assert_compiles("test [3, 4, 5]"), @"[3, [4], 5]");
+    assert_snapshot!(assert_compiles_allowing_exits("test [3, 4, 5]"), @"[3, [4], 5]");
 }
 
 #[test]
@@ -3876,7 +3876,7 @@ fn test_dupn() {
         test([1, 1])
     ");
     assert_contains_opcode("test", YARVINSN_dupn);
-    assert_snapshot!(assert_compiles("
+    assert_snapshot!(assert_compiles_allowing_exits("
         one = [1, 1]
         start_empty = []
         [test(one), one, test(start_empty), start_empty]
@@ -4438,7 +4438,7 @@ fn test_nil_value_nil_opt_with_guard_side_exit() {
         test(nil)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(1)"), @"false");
+    assert_snapshot!(assert_compiles_allowing_exits("test(1)"), @"false");
 }
 
 #[test]
@@ -4459,7 +4459,7 @@ fn test_true_nil_opt_with_guard_side_exit() {
         test(true)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(nil)"), @"true");
+    assert_snapshot!(assert_compiles_allowing_exits("test(nil)"), @"true");
 }
 
 #[test]
@@ -4480,7 +4480,7 @@ fn test_false_nil_opt_with_guard_side_exit() {
         test(false)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(nil)"), @"true");
+    assert_snapshot!(assert_compiles_allowing_exits("test(nil)"), @"true");
 }
 
 #[test]
@@ -4501,7 +4501,7 @@ fn test_integer_nil_opt_with_guard_side_exit() {
         test(2)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(nil)"), @"true");
+    assert_snapshot!(assert_compiles_allowing_exits("test(nil)"), @"true");
 }
 
 #[test]
@@ -4522,7 +4522,7 @@ fn test_float_nil_opt_with_guard_side_exit() {
         test(2.0)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(nil)"), @"true");
+    assert_snapshot!(assert_compiles_allowing_exits("test(nil)"), @"true");
 }
 
 #[test]
@@ -4543,7 +4543,7 @@ fn test_symbol_nil_opt_with_guard_side_exit() {
         test(:bar)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(nil)"), @"true");
+    assert_snapshot!(assert_compiles_allowing_exits("test(nil)"), @"true");
 }
 
 #[test]
@@ -4553,7 +4553,7 @@ fn test_class_nil_opt_with_guard() {
         test(String)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(Integer)"), @"false");
+    assert_snapshot!(assert_compiles_allowing_exits("test(Integer)"), @"false");
 }
 
 #[test]
@@ -4564,7 +4564,7 @@ fn test_class_nil_opt_with_guard_side_exit() {
         test(Integer)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(nil)"), @"true");
+    assert_snapshot!(assert_compiles_allowing_exits("test(nil)"), @"true");
 }
 
 #[test]
@@ -4574,7 +4574,7 @@ fn test_module_nil_opt_with_guard() {
         test(Enumerable)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(Kernel)"), @"false");
+    assert_snapshot!(assert_compiles_allowing_exits("test(Kernel)"), @"false");
 }
 
 #[test]
@@ -4585,7 +4585,7 @@ fn test_module_nil_opt_with_guard_side_exit() {
         test(Kernel)
     ");
     assert_contains_opcode("test", YARVINSN_opt_nil_p);
-    assert_snapshot!(assert_compiles("test(nil)"), @"true");
+    assert_snapshot!(assert_compiles_allowing_exits("test(nil)"), @"true");
 }
 
 #[test]
@@ -4923,7 +4923,7 @@ fn test_allocating_in_hir_c_method_is() {
         second
     ");
     assert_contains_opcode("test", YARVINSN_opt_new);
-    assert_snapshot!(assert_compiles("a(Foo)"), @":k");
+    assert_snapshot!(assert_compiles_allowing_exits("a(Foo)"), @":k");
 }
 
 #[test]
@@ -5050,7 +5050,7 @@ fn test_fixnum_div_zero() {
         test(0)
     ");
     assert_contains_opcode("test", YARVINSN_opt_div);
-    assert_snapshot!(assert_compiles(r#"test(0)"#), @r#""divided by 0""#);
+    assert_snapshot!(assert_compiles_allowing_exits(r#"test(0)"#), @r#""divided by 0""#);
 }
 
 #[test]
diff --git a/zjit/src/cruby.rs b/zjit/src/cruby.rs
index 6ea041d72f..bdea8a2404 100644
--- a/zjit/src/cruby.rs
+++ b/zjit/src/cruby.rs
@@ -1280,11 +1280,23 @@ pub mod test_utils {
     }
 
     /// Like inspect, but also asserts that all compilations triggered by this program succeed.
+    pub fn assert_compiles_allowing_exits(program: &str) -> String {
+        use crate::state::ZJITState;
+        ZJITState::enable_assert_compiles();
+        let result = inspect(program);
+        ZJITState::disable_assert_compiles();
+        result
+    }
+
+    /// Like inspect, but also asserts that all compilations triggered by this program succeed and
+    /// no side exits occurr during the program.
     pub fn assert_compiles(program: &str) -> String {
         use crate::state::ZJITState;
+        let exits_before = crate::stats::total_exit_count();
         ZJITState::enable_assert_compiles();
         let result = inspect(program);
         ZJITState::disable_assert_compiles();
+        assert_eq!(exits_before, crate::stats::total_exit_count(), "Program side-exited");
         result
     }
 
diff --git a/zjit/src/hir/opt_tests.rs b/zjit/src/hir/opt_tests.rs
index 9be09ba21d..f4755d5fac 100644
--- a/zjit/src/hir/opt_tests.rs
+++ b/zjit/src/hir/opt_tests.rs
@@ -7511,6 +7511,15 @@ mod hir_opt_tests {
     }
 
     #[test]
+    fn test_no_side_exit_assertion() {
+        eval("
+          def side_exit = ::RubyVM::ZJIT.induce_side_exit!
+          side_exit
+        ");
+        std::panic::catch_unwind(|| assert_compiles("side_exit")).expect_err("Should panic because the program should side exit");
+    }
+
+    #[test]
     fn test_optimize_getivar_on_class_embedded() {
         eval("
             class C
@@ -7519,6 +7528,7 @@ mod hir_opt_tests {
             end
             C.test
         ");
+        assert_snapshot!(assert_compiles("C.test"), @"42");
         assert_snapshot!(hir_string_proc("C.method(:test)"), @"
         fn test@<compiled>:4:
         bb1():
diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs
index 38d69c7533..2e43706e85 100644
--- a/zjit/src/stats.rs
+++ b/zjit/src/stats.rs
@@ -983,6 +983,10 @@ pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE, target_key: VALUE) ->
     hash
 }
 
+pub fn total_exit_count() -> u64 {
+    EXIT_COUNTERS.iter().fold(0, |sum, counter| sum + unsafe { *counter_ptr(*counter) })
+}
+
 /// Measure the time taken by func() and add that to zjit_compile_time.
 pub fn with_time_stat<F, R>(counter: Counter, func: F) -> R where F: FnOnce() -> R {
     let start = Instant::now();