Optimize calls and fix bug with Instance preservation of cross-instance calls (#1878)

* push inst to call stack only if it differs

* fix bug with return_call[_indirect] cross-instance calls

* add cross-instance-calls.wast Wasmi test cases

Author: Robbepop

crates/wasmi/src/engine/executor/handler/exec.rs

@@ -413,8 +413,12 @@ execution_handler! {
                 let wasm_func = func.func_body();
                 let callee_instance = *func.instance();
                 let callee_instance = resolve_instance(state.store, &callee_instance).into();
+                let changed_instance = match callee_instance != instance {
+                    true => Some(callee_instance),
+                    false => None,
+                };
                 let (callee_ip, callee_sp) =
-                    return_call_wasm(state, params, wasm_func, Some(instance))?;
+                    return_call_wasm(state, params, wasm_func, changed_instance)?;
                 (callee_ip, callee_sp, callee_instance)
             }
             FuncEntity::Host(host_func) => {
@@ -457,8 +461,12 @@ execution_handler! {
                 let wasm_func = func.func_body();
                 let callee_instance = *func.instance();
                 let callee_instance: Inst = resolve_instance(state.store, &callee_instance).into();
+                let changed_instance = match callee_instance != instance {
+                    true => Some(callee_instance),
+                    false => None,
+                };
                 let (callee_ip, callee_sp) =
-                    return_call_wasm(state, params, wasm_func, Some(instance))?;
+                    return_call_wasm(state, params, wasm_func, changed_instance)?;
                 (callee_ip, callee_sp, callee_instance)
             }
             FuncEntity::Host(host_func) => {

crates/wasmi/src/engine/executor/handler/state.rs

@@ -1225,8 +1225,8 @@ impl CallStack {
             None => debug_assert!(self.frames.is_empty()),
         }
         let prev_instance = match instance {
-            Some(instance) => self.instance.replace(instance),
-            None => self.instance,
+            Some(instance) if self.instance != Some(instance) => self.instance.replace(instance),
+            _ => None,
         };
         let params_offset = usize::from(u16::from(callee_params.span().head()));
         let start = self.top_start().add(params_offset)?;
@@ -1253,21 +1253,30 @@ impl CallStack {
     }
 
     /// Adjusts `self` for a function tail call.
+    ///
+    /// # Note
+    ///
+    /// - If `instance` is `Some` it refers to the _callee_ instance of the tail call
+    ///   which is required to be different from the currently used instance.
+    /// - If `instance` is `None` the callee shares the currently used instance.
     #[inline(always)]
     fn replace(&mut self, callee_ip: Ip, instance: Option<Inst>) -> Result<SpOffset, TrapCode> {
         let Some(caller_frame) = self.frames.last_mut() else {
             unsafe { unreachable_unchecked!("missing caller frame on the call stack") }
         };
-        let prev_instance = match instance {
-            Some(instance) => self.instance.replace(instance),
-            None => self.instance,
-        };
         let start = caller_frame.start;
-        *caller_frame = Frame {
-            start,
-            ip: callee_ip,
-            instance: prev_instance,
-        };
+        caller_frame.ip = callee_ip;
+        if let Some(callee_instance) = instance {
+            debug_assert!(self.instance != Some(callee_instance));
+            // The replaced frame's restoration obligation is carried over unchanged.
+            // However, if the replaced frame has no such obligation yet, its caller
+            // runs in the currently used instance which must be restored when the
+            // new frame returns since the callee continues in a different instance.
+            if caller_frame.instance.is_none() {
+                caller_frame.instance = self.instance;
+            }
+            self.instance = Some(callee_instance);
+        }
         Ok(start)
     }
 }

crates/wast/tests/mod.rs

@@ -200,6 +200,7 @@ macro_rules! foreach_test {
             fn wasmi_audit("wasmi/tests/audit");
             fn wasmi_call("wasmi/tests/call");
             fn wasmi_call_indirect("wasmi/tests/call-indirect");
+            fn wasmi_cross_instance_calls("wasmi/tests/cross-instance-calls");
             fn wasmi_return_call("wasmi/tests/return-call");
             fn wasmi_return_call_indirect("wasmi/tests/return-call-indirect");
             fn wasmi_global_set("wasmi/tests/global-set");