Detect and abort tracing if "early return" occurs.

If we start tracing a loop and the function-with-a-control-point returns
whilst tracing is ongoing, we are in dangerous territory. For quite a
while we assumed that "dangerous territory" meant (gongs of doom!) "UB".
Actually, we just hadn't created very good examples: when I did, it
turns out that an `unwrap` in trace_builder fails if early return
happens. We *could*, therefore, detect early returns there, but that
means that we've done a lot of work before realising this.

This commit instead has `mt.rs` detect early returns, and abort tracing,
as soon as we reasonably can. [Note, however, that if a user annotates
`return`s from the function, they could stop things even earlier. Is
this worth it? It is quite complicated to explain and, as we'll see in a
bit, it's not a very common case]

All that said, it turns out early returns aren't that common: you have
to start tracing on the same iteration that you "early return" from.
This is definitely possible -- indeed the two C tests in the PR do that,
one "normally" and one when side-tracing -- but it seems unlikely that
it will happen very often. So although we could force the user to help
us detect this even sooner, I'm not sure it's worth burdening them or us
with that responsibility right now: the cases where it would save
serious resources are borderline pathological. Indeed, the only such
case I can think of is: if you start tracing from the "root" control
point call frame, and then return to the main program (i.e. you will
never go back to the interpreter), then tracing is never turned off. One
day we can fix this, but I don't think we will encounter it very soon.

tests/c/early_return1.c

@@ -0,0 +1,65 @@
+// Run-time:
+//   env-var: YKD_SERIALISE_COMPILATION=1
+//   env-var: YKD_LOG_IR=jit-pre-opt
+//   env-var: YK_LOG=4
+//   stderr:
+//     yk-jit-event: start-tracing
+//     early return
+//     5
+//     yk-jit-event: tracing-aborted
+//     4
+//     yk-jit-event: start-tracing
+//     3
+//     yk-jit-event: stop-tracing
+//     --- Begin jit-pre-opt ---
+//     ...
+//     --- End jit-pre-opt ---
+//     2
+//     yk-jit-event: enter-jit-code
+//     1
+//     yk-jit-event: deoptimise
+//     return
+//     exit
+
+// Check that early return from a recursive interpreter loop aborts tracing,
+// but doesn't stop a location being retraced.
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <yk.h>
+#include <yk_testing.h>
+
+void loop(YkMT *, YkLocation *, int);
+
+void loop(YkMT *mt, YkLocation *loc, int i) {
+  NOOPT_VAL(i);
+  while (i > 0) {
+    yk_mt_control_point(mt, loc);
+    if (i == 6) {
+      loop(mt, loc, i - 1);
+      i--;
+    } else if (i == 5) {
+      fprintf(stderr, "early return\n");
+      return;
+    }
+    fprintf(stderr, "%d\n", i);
+    i--;
+  }
+  fprintf(stderr, "return\n");
+  return;
+}
+
+int main(int argc, char **argv) {
+  YkMT *mt = yk_mt_new(NULL);
+  yk_mt_hot_threshold_set(mt, 1);
+  YkLocation loc = yk_location_new();
+
+  NOOPT_VAL(loc);
+  loop(mt, &loc, 6);
+  fprintf(stderr, "exit\n");
+  yk_location_drop(loc);
+  yk_mt_shutdown(mt);
+  return (EXIT_SUCCESS);
+}

tests/c/early_return2.c

@@ -0,0 +1,118 @@
+
+// Run-time:
+//   env-var: YKD_SERIALISE_COMPILATION=1
+//   env-var: YKD_LOG_IR=jit-pre-opt
+//   env-var: YK_LOG=4
+//   stderr:
+//     yk-jit-event: start-tracing
+//     b3
+//     8
+//     yk-jit-event: stop-tracing
+//     --- Begin jit-pre-opt ---
+//     ...
+//     --- End jit-pre-opt ---
+//     b1
+//     7
+//     yk-jit-event: enter-jit-code
+//     yk-jit-event: deoptimise
+//     yk-jit-event: start-side-tracing
+//     b1
+//     9
+//     yk-jit-event: tracing-aborted
+//     b3
+//     8
+//     yk-jit-event: enter-jit-code
+//     yk-jit-event: deoptimise
+//     yk-jit-event: start-side-tracing
+//     b3
+//     7
+//     yk-jit-event: stop-tracing
+//     --- Begin jit-pre-opt ---
+//     ...
+//     --- End jit-pre-opt ---
+//     b3
+//     6
+//     yk-jit-event: enter-jit-code
+//     yk-jit-event: execute-side-trace
+//     yk-jit-event: deoptimise
+//     yk-jit-event: start-side-tracing
+//     b2
+//     5
+//     yk-jit-event: stop-tracing
+//     --- Begin jit-pre-opt ---
+//     ...
+//     --- End jit-pre-opt ---
+//     b2
+//     4
+//     yk-jit-event: enter-jit-code
+//     yk-jit-event: execute-side-trace
+//     yk-jit-event: execute-side-trace
+//     b2
+//     3
+//     yk-jit-event: execute-side-trace
+//     yk-jit-event: execute-side-trace
+//     b2
+//     2
+//     yk-jit-event: execute-side-trace
+//     yk-jit-event: execute-side-trace
+//     b2
+//     1
+//     yk-jit-event: execute-side-trace
+//     yk-jit-event: execute-side-trace
+//     b2
+//     0
+//     yk-jit-event: deoptimise
+//     yk-jit-event: start-side-tracing
+//     exit
+
+// Check that early return from a recursive interpreter loop aborts tracing,
+// but doesn't stop a location being retraced.
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdbool.h>
+#include <stdlib.h>
+#include <string.h>
+#include <yk.h>
+#include <yk_testing.h>
+
+void loop(YkMT *, YkLocation *, int, bool);
+
+void loop(YkMT *mt, YkLocation *loc, int i, bool inner) {
+  NOOPT_VAL(i);
+  while (i > 0) {
+    yk_mt_control_point(mt, loc);
+    if (i == 10) {
+      loop(mt, loc, i - 1, true);
+      i--;
+    } else if (inner && i <= 8) {
+      fprintf(stderr, "b1\n");
+      i--;
+    } else if (!inner && i <= 6) {
+      fprintf(stderr, "b2\n");
+      i--;
+    } else {
+      fprintf(stderr, "b3\n");
+      i--;
+    }
+    if (inner && i == 6) {
+      return;
+    }
+    fprintf(stderr, "%d\n", i);
+  }
+  return;
+}
+
+int main(int argc, char **argv) {
+  YkMT *mt = yk_mt_new(NULL);
+  yk_mt_hot_threshold_set(mt, 1);
+  yk_mt_sidetrace_threshold_set(mt, 1);
+  YkLocation loc = yk_location_new();
+
+  NOOPT_VAL(loc);
+  loop(mt, &loc, 10, false);
+  fprintf(stderr, "exit\n");
+  yk_location_drop(loc);
+  yk_mt_shutdown(mt);
+  return (EXIT_SUCCESS);
+}

ykrt/src/compile/jitc_yk/codegen/reg_alloc.rs

@@ -50,12 +50,6 @@ pub(crate) enum Register {
     FP(Rx), // floating point
 }
 
-/// Indicates the direction of stack growth.
-pub(crate) enum StackDirection {
-    GrowsUp,
-    GrowsDown,
-}
-
 #[cfg(target_arch = "x86_64")]
 impl VarLocation {
     pub(crate) fn from_yksmp_location(m: &Module, iidx: InstIdx, x: &yksmp::Location) -> Self {

ykrt/src/compile/jitc_yk/codegen/x64/deopt.rs

@@ -94,7 +94,7 @@ fn running_trace(gidxs: &[usize]) -> Arc<X64CompiledTrace> {
 /// * glen - Length for list in `gptr`.
 #[no_mangle]
 pub(crate) extern "C" fn __yk_deopt(
-    frameaddr: *const c_void,
+    frameaddr: *mut c_void,
     gidx: u64,
     gp_regs: &[u64; 16],
     fp_regs: &[u64; 16],
@@ -353,13 +353,13 @@ pub(crate) extern "C" fn __yk_deopt(
 
     // The `clone` should really be `Arc::clone(&ctr)` but that doesn't play well with type
     // inference in this (unusual) case.
-    ctr.mt.guard_failure(ctr.clone(), gidx);
+    ctr.mt.guard_failure(ctr.clone(), gidx, frameaddr);
 
     // Since we won't return from this function, drop `ctr` manually.
     drop(ctr);
 
     // Now overwrite the existing stack with our newly recreated one.
-    unsafe { replace_stack(newframedst as *mut c_void, newstack, memsize) };
+    unsafe { replace_stack(newframedst, newstack, memsize) };
 }
 
 /// Writes the stack frames that we recreated in [__yk_deopt] onto the current stack, overwriting

ykrt/src/compile/jitc_yk/codegen/x64/mod.rs

@@ -26,7 +26,7 @@ use super::{
         jit_ir::{self, BinOp, FloatTy, Inst, InstIdx, Module, Operand, PtrAddInst, Ty},
         CompilationError,
     },
-    reg_alloc::{self, StackDirection, VarLocation},
+    reg_alloc::{self, VarLocation},
     CodeGen,
 };
 #[cfg(any(debug_assertions, test))]
@@ -140,9 +140,6 @@ static RBP_DWARF_NUM: u16 = 6;
 /// The x64 SysV ABI requires a 16-byte aligned stack prior to any call.
 const SYSV_CALL_STACK_ALIGN: usize = 16;
 
-/// On x64 the stack grows down.
-const STACK_DIRECTION: StackDirection = StackDirection::GrowsDown;
-
 /// A function that we can put a debugger breakpoint on.
 /// FIXME: gross hack.
 #[cfg(debug_assertions)]

ykrt/src/compile/jitc_yk/trace_builder.rs

@@ -589,8 +589,10 @@ impl TraceBuilder {
         _aot_inst_idx: usize,
         val: &Option<aot_ir::Operand>,
     ) -> Result<(), CompilationError> {
-        // FIXME: Map return value to AOT call instruction.
+        // If this `unwrap` fails, it means that early return detection in `mt.rs` is not working
+        // as expected.
         let frame = self.frames.pop().unwrap();
+        // FIXME: Map return value to AOT call instruction.
         if let Some(val) = val {
             let op = self.handle_operand(val)?;
             self.local_map.insert(frame.callinst.unwrap(), op);

ykrt/src/lib.rs

@@ -15,6 +15,7 @@ mod location;
 mod log;
 pub(crate) mod mt;
 pub mod promote;
+pub(crate) mod stack;
 pub(crate) mod thread_intercept;
 pub mod trace;