Save exit Wasm FP and PC in component-to-host trampolines (#4601)

* Wasmtime: Add a pointer to `VMRuntimeLimits` in component contexts

* Save exit Wasm FP and PC in component-to-host trampolines

Fixes #4535

* Add comment about why we deref the trampoline's FP

* Update some tests to use new `vmruntime_limits_*` methods

crates/cranelift/src/compiler.rs

@@ -30,8 +30,8 @@ use std::mem;
 use std::sync::Mutex;
 use wasmtime_environ::{
     AddressMapSection, CompileError, FilePos, FlagValue, FunctionBodyData, FunctionInfo,
-    InstructionAddressMap, Module, ModuleTranslation, ModuleTypes, StackMapInformation, Trampoline,
-    TrapCode, TrapEncodingBuilder, TrapInformation, Tunables, VMOffsets,
+    InstructionAddressMap, Module, ModuleTranslation, ModuleTypes, PtrSize, StackMapInformation,
+    Trampoline, TrapCode, TrapEncodingBuilder, TrapInformation, Tunables, VMOffsets,
 };
 
 #[cfg(feature = "component-model")]
@@ -196,7 +196,7 @@ impl wasmtime_environ::Compiler for Compiler {
         });
         let stack_limit = context.func.create_global_value(ir::GlobalValueData::Load {
             base: interrupts_ptr,
-            offset: i32::try_from(func_env.offsets.vmruntime_limits_stack_limit())
+            offset: i32::try_from(func_env.offsets.ptr.vmruntime_limits_stack_limit())
                 .unwrap()
                 .into(),
             global_type: isa.pointer_type(),

crates/cranelift/src/compiler/component.rs

@@ -13,7 +13,7 @@ use wasmtime_environ::component::{
     AlwaysTrapInfo, CanonicalOptions, Component, ComponentCompiler, ComponentTypes, FunctionInfo,
     LowerImport, LoweredIndex, RuntimeAlwaysTrapIndex, VMComponentOffsets,
 };
-use wasmtime_environ::{PrimaryMap, SignatureIndex, Trampoline, TrapCode, WasmFuncType};
+use wasmtime_environ::{PrimaryMap, PtrSize, SignatureIndex, Trampoline, TrapCode, WasmFuncType};
 
 impl ComponentCompiler for Compiler {
     fn compile_lowered_trampoline(
@@ -46,6 +46,44 @@ impl ComponentCompiler for Compiler {
             self.wasm_to_host_spill_args(ty, &mut builder, block0);
         let vmctx = builder.func.dfg.block_params(block0)[0];
 
+        // Save the exit FP and return address for stack walking purposes.
+        //
+        // First we need to get the `VMRuntimeLimits`.
+        let limits = builder.ins().load(
+            pointer_type,
+            MemFlags::trusted(),
+            vmctx,
+            i32::try_from(offsets.limits()).unwrap(),
+        );
+        // Then save the exit Wasm FP to the limits. We dereference the current
+        // FP to get the previous FP because the current FP is the trampoline's
+        // FP, and we want the Wasm function's FP, which is the caller of this
+        // trampoline.
+        let trampoline_fp = builder.ins().get_frame_pointer(pointer_type);
+        let wasm_fp = builder.ins().load(
+            pointer_type,
+            MemFlags::trusted(),
+            trampoline_fp,
+            // The FP always points to the next older FP for all supported
+            // targets. See assertion in
+            // `crates/runtime/src/traphandlers/backtrace.rs`.
+            0,
+        );
+        builder.ins().store(
+            MemFlags::trusted(),
+            wasm_fp,
+            limits,
+            offsets.ptr.vmruntime_limits_last_wasm_exit_fp(),
+        );
+        // Finally save the Wasm return address to the limits.
+        let wasm_pc = builder.ins().get_return_address(pointer_type);
+        builder.ins().store(
+            MemFlags::trusted(),
+            wasm_pc,
+            limits,
+            offsets.ptr.vmruntime_limits_last_wasm_exit_pc(),
+        );
+
         // Below this will incrementally build both the signature of the host
         // function we're calling as well as the list of arguments since the
         // list is somewhat long.

crates/cranelift/src/func_environ.rs

@@ -528,7 +528,7 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
     ) -> (ir::Value, ir::immediates::Offset32) {
         (
             builder.use_var(self.vmruntime_limits_ptr),
-            i32::from(self.offsets.vmruntime_limits_fuel_consumed()).into(),
+            i32::from(self.offsets.ptr.vmruntime_limits_fuel_consumed()).into(),
         )
     }
 
@@ -628,12 +628,15 @@ impl<'module_environment> FuncEnvironment<'module_environment> {
 
     fn epoch_load_deadline_into_var(&mut self, builder: &mut FunctionBuilder<'_>) {
         let interrupts = builder.use_var(self.vmruntime_limits_ptr);
-        let deadline = builder.ins().load(
-            ir::types::I64,
-            ir::MemFlags::trusted(),
-            interrupts,
-            ir::immediates::Offset32::new(self.offsets.vmruntime_limits_epoch_deadline() as i32),
-        );
+        let deadline =
+            builder.ins().load(
+                ir::types::I64,
+                ir::MemFlags::trusted(),
+                interrupts,
+                ir::immediates::Offset32::new(
+                    self.offsets.ptr.vmruntime_limits_epoch_deadline() as i32
+                ),
+            );
         builder.def_var(self.epoch_deadline_var, deadline);
     }