Insert trampolines to preserve callee-saved registers for backends without register save area information.

lib/runtime-core/src/fault.rs

@@ -1,8 +1,9 @@
-mod raw {
+pub mod raw {
     use std::ffi::c_void;
 
     extern "C" {
         pub fn run_on_alternative_stack(stack_end: *mut u64, stack_begin: *mut u64) -> u64;
+        pub fn register_preservation_trampoline(); // NOT safe to call directly
         pub fn setjmp(env: *mut c_void) -> i32;
         pub fn longjmp(env: *mut c_void, val: i32) -> !;
     }
@@ -39,11 +40,27 @@ struct UnwindInfo {
     payload: Option<Box<dyn Any>>, // out
 }
 
+#[repr(packed)]
+#[derive(Default, Copy, Clone)]
+pub struct BoundaryRegisterPreservation {
+    pub r15: u64,
+    pub r14: u64,
+    pub r13: u64,
+    pub r12: u64,
+    pub rbx: u64,
+}
+
 thread_local! {
     static UNWIND: UnsafeCell<Option<UnwindInfo>> = UnsafeCell::new(None);
     static CURRENT_CTX: UnsafeCell<*mut vm::Ctx> = UnsafeCell::new(::std::ptr::null_mut());
     static CURRENT_CODE_VERSIONS: RefCell<Vec<CodeVersion>> = RefCell::new(vec![]);
     static WAS_SIGINT_TRIGGERED: Cell<bool> = Cell::new(false);
+    static BOUNDARY_REGISTER_PRESERVATION: UnsafeCell<BoundaryRegisterPreservation> = UnsafeCell::new(BoundaryRegisterPreservation::default());
+}
+
+#[no_mangle]
+pub unsafe extern "C" fn get_boundary_register_preservation() -> *mut BoundaryRegisterPreservation {
+    BOUNDARY_REGISTER_PRESERVATION.with(|x| x.get())
 }
 
 struct InterruptSignalMem(*mut u8);

lib/runtime-core/src/state.rs

@@ -104,6 +104,7 @@ pub struct InstanceImage {
 
 #[derive(Debug, Clone)]
 pub struct CodeVersion {
+    pub baseline: bool,
     pub msm: ModuleStateMap,
     pub base: usize,
 }
@@ -372,7 +373,7 @@ impl InstanceImage {
 pub mod x64 {
     use super::*;
     use crate::codegen::BreakpointMap;
-    use crate::fault::{catch_unsafe_unwind, run_on_alternative_stack};
+    use crate::fault::{catch_unsafe_unwind, run_on_alternative_stack, get_boundary_register_preservation};
     use crate::structures::TypedIndex;
     use crate::types::LocalGlobalIndex;
     use crate::vm::Ctx;
@@ -763,8 +764,9 @@ pub mod x64 {
     ) -> ExecutionStateImage {
         let mut known_registers: [Option<u64>; 24] = initially_known_registers;
         let mut results: Vec<WasmFunctionStateDump> = vec![];
+        let mut was_baseline = true;
 
-        for _ in 0.. {
+        for i in 0.. {
             let ret_addr = initial_address.take().unwrap_or_else(|| {
                 let x = *stack;
                 stack = stack.offset(1);
@@ -772,6 +774,7 @@ pub mod x64 {
             });
 
             let mut fsm_state: Option<(&FunctionStateMap, MachineState)> = None;
+            let mut is_baseline: Option<bool> = None;
 
             for version in versions() {
                 match version
@@ -786,6 +789,7 @@ pub mod x64 {
                 {
                     Some(x) => {
                         fsm_state = Some(x);
+                        is_baseline = Some(version.baseline);
                         break;
                     }
                     None => {}
@@ -798,6 +802,22 @@ pub mod x64 {
                 return ExecutionStateImage { frames: results };
             };
 
+            {
+                let is_baseline = is_baseline.unwrap();
+
+                // Are we unwinding through an optimized/baseline boundary?
+                if is_baseline && !was_baseline {
+                    let callee_saved = &*get_boundary_register_preservation();
+                    known_registers[X64Register::GPR(GPR::R15).to_index().0] = Some(callee_saved.r15);
+                    known_registers[X64Register::GPR(GPR::R14).to_index().0] = Some(callee_saved.r14);
+                    known_registers[X64Register::GPR(GPR::R13).to_index().0] = Some(callee_saved.r13);
+                    known_registers[X64Register::GPR(GPR::R12).to_index().0] = Some(callee_saved.r12);
+                    known_registers[X64Register::GPR(GPR::RBX).to_index().0] = Some(callee_saved.rbx);
+                }
+
+                was_baseline = is_baseline;
+            }
+
             let mut wasm_stack: Vec<Option<u64>> = state
                 .wasm_stack
                 .iter()

lib/runtime-core/src/tiering.rs

@@ -111,6 +111,7 @@ pub fn run_tiering<F: Fn(InteractiveShellContext) -> ShellExitOperation>(
         let mut optimized_instances: Vec<Instance> = vec![];
 
         push_code_version(CodeVersion {
+            baseline: true,
             msm: baseline
                 .module
                 .runnable_module
@@ -166,6 +167,7 @@ pub fn run_tiering<F: Fn(InteractiveShellContext) -> ShellExitOperation>(
                 }
 
                 push_code_version(CodeVersion {
+                    baseline: false,
                     msm: optimized
                         .module
                         .runnable_module
@@ -179,6 +181,8 @@ pub fn run_tiering<F: Fn(InteractiveShellContext) -> ShellExitOperation>(
                         .as_ptr() as usize,
                 });
                 n_versions.set(n_versions.get() + 1);
+
+                baseline.context_mut().local_functions = optimized.context_mut().local_functions;
             }
             // TODO: Fix this for optimized version.
             let breakpoints = baseline.module.runnable_module.get_breakpoints();