fix(jit): preserve region progress across side exits

Author: fffonion

src/vm/jit/runtime.rs

@@ -321,11 +321,15 @@ impl Vm {
 
         loop {
             native::clear_bridge_error();
+            let region_edges_before = self.jit_native_region_edge_count;
             let status = unsafe { entry(self as *mut Vm) };
             self.native_trace_exec_count = self.native_trace_exec_count.saturating_add(1);
             if is_region {
                 self.jit_native_region_entry_count =
                     self.jit_native_region_entry_count.saturating_add(1);
+                if self.jit_native_region_edge_count > region_edges_before {
+                    self.jit.record_native_region_progress(current_trace_id);
+                }
             }
             self.jit.mark_trace_executed(current_trace_id);
             let mut trace_exit_key = None;
@@ -823,11 +827,15 @@ impl Vm {
         ) = self.native_trace_state(current_trace_id)?;
         native::clear_bridge_error();
         loop {
+            let region_edges_before = self.jit_native_region_edge_count;
             let status = unsafe { entry(self as *mut Vm) };
             self.native_trace_exec_count = self.native_trace_exec_count.saturating_add(1);
             if is_region {
                 self.jit_native_region_entry_count =
                     self.jit_native_region_entry_count.saturating_add(1);
+                if self.jit_native_region_edge_count > region_edges_before {
+                    self.jit.record_native_region_progress(current_trace_id);
+                }
             }
             self.jit.mark_trace_executed(current_trace_id);
             let mut trace_exit_key = None;

src/vm/jit/trace.rs

@@ -712,6 +712,10 @@ impl TraceJitEngine {
         }
     }
 
+    pub(crate) fn record_native_region_progress(&mut self, trace_id: usize) {
+        self.record_native_loop_back(trace_id);
+    }
+
     pub(crate) fn callable_frame_is_blocked(&self, frame_key: u64) -> bool {
         if frame_key == ROOT_FRAME_KEY {
             return false;
@@ -1609,7 +1613,7 @@ mod tests {
     }
 
     #[test]
-    fn callable_side_exit_backoff_resets_on_native_loopback() {
+    fn callable_side_exit_backoff_resets_on_native_progress() {
         if !native_jit_supported() {
             return;
         }
@@ -1653,6 +1657,10 @@ mod tests {
         for _ in 1..CALLABLE_SIDE_EXIT_BACKOFF_THRESHOLD {
             assert!(!engine.record_native_side_exit(trace_id));
         }
+        engine.record_native_region_progress(trace_id);
+        for _ in 1..CALLABLE_SIDE_EXIT_BACKOFF_THRESHOLD {
+            assert!(!engine.record_native_side_exit(trace_id));
+        }
         assert!(engine.record_native_side_exit(trace_id));
         engine.block_callable_frame(trace_id);
         assert!(engine.callable_frame_is_blocked(0));

tests/jit/jit_tests.rs

@@ -2463,6 +2463,57 @@ fn trace_jit_region_preserves_owned_value_drop_contract() {
     assert_eq!(Arc::strong_count(&output), 1);
 }
 
+#[test]
+fn trace_jit_region_progress_prevents_callable_frame_backoff() {
+    if !native_jit_supported() {
+        return;
+    }
+    let source = r#"
+        fn run(limit, payload) {
+            let mut i = 0;
+            let mut total = 0;
+            while i < limit {
+                if i % 2 == 0 {
+                    total = total + 3;
+                } else {
+                    total = total + 5;
+                }
+                i = i + 1;
+            }
+            total + payload[0]
+        }
+        run(256, [1]);
+    "#;
+    let compiled = compile_source(source).expect("region progress fixture should compile");
+    let mut vm = Vm::new(compiled.program);
+    vm.set_jit_config(JitConfig {
+        enabled: true,
+        hot_loop_threshold: 1,
+        max_trace_len: 256,
+    });
+
+    assert_eq!(vm.run().unwrap(), VmStatus::Halted);
+    assert_eq!(vm.stack(), &[Value::Int(1_025)]);
+    assert_eq!(vm.jit_native_region_count(), 1, "{}", vm.dump_jit_info());
+    let first_execs = vm.jit_native_exec_count();
+    let first_edges = vm.jit_native_internal_region_edge_count();
+
+    vm.reset_for_reuse();
+    assert_eq!(vm.run().unwrap(), VmStatus::Halted);
+    assert_eq!(vm.stack(), &[Value::Int(1_025)]);
+    assert!(
+        vm.jit_native_exec_count() > first_execs,
+        "{}",
+        vm.dump_jit_info()
+    );
+    assert!(
+        vm.jit_native_internal_region_edge_count() > first_edges,
+        "{}",
+        vm.dump_jit_info()
+    );
+    assert_eq!(vm.jit_native_region_count(), 1, "{}", vm.dump_jit_info());
+}
+
 #[test]
 fn trace_jit_region_respects_fuel_and_epoch_interrupts() {
     if !native_jit_supported() {