make special smoldot tail resume tests pass under regmach

wasmi-labs · Nov 30, 2023 · 0a4c5da · 0a4c5da
1 parent 9e6e025
commit 0a4c5da
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 7 deletions.
diff --git a/crates/wasmi/src/engine/regmach/executor/instrs.rs b/crates/wasmi/src/engine/regmach/executor/instrs.rs
@@ -1,4 +1,4 @@
-use self::{call::CallOutcome, return_::ReturnOutcome};
+use self::{call::{CallOutcome}, return_::ReturnOutcome};
 use crate::{
     core::{TrapCode, UntypedValue},
     engine::{
@@ -27,6 +27,7 @@ use crate::{
     FuncRef,
     StoreInner,
 };
+pub use self::call::CallKind;
 
 mod binary;
 mod branch;
@@ -45,8 +46,8 @@ mod unary;
 
 macro_rules! forward_call {
     ($expr:expr) => {{
-        if let CallOutcome::Call { results, host_func } = $expr? {
-            return Ok(WasmOutcome::Call { results, host_func });
+        if let CallOutcome::Call { results, host_func, call_kind } = $expr? {
+            return Ok(WasmOutcome::Call { results, host_func, call_kind });
         }
     }};
 }
@@ -73,6 +74,7 @@ pub enum WasmOutcome {
     Call {
         results: RegisterSpan,
         host_func: Func,
+        call_kind: CallKind,
     },
 }
 

diff --git a/crates/wasmi/src/engine/regmach/executor/instrs/call.rs b/crates/wasmi/src/engine/regmach/executor/instrs/call.rs
@@ -39,6 +39,7 @@ pub enum CallOutcome {
     Call {
         results: RegisterSpan,
         host_func: Func,
+        call_kind: CallKind,
     },
 }
 
@@ -341,13 +342,19 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
                 let offset = self.value_stack.extend_zeros(max_inout);
                 let offset_sp = unsafe { self.value_stack.stack_ptr_at(offset) };
                 if matches!(params, CallParams::Some) {
-                    self.ip = self.copy_call_params(offset_sp);
+                    let new_ip = self.copy_call_params(offset_sp);
+                    if matches!(call_kind, CallKind::Nested) {
+                        self.ip = new_ip;
+                    }
+                }
+                if matches!(call_kind, CallKind::Nested) {
+                    self.update_instr_ptr_at(1);
                 }
-                self.update_instr_ptr_at(1);
                 self.cache.reset();
                 Ok(CallOutcome::Call {
                     results,
                     host_func: *func,
+                    call_kind,
                 })
             }
         }

diff --git a/crates/wasmi/src/engine/regmach/executor/mod.rs b/crates/wasmi/src/engine/regmach/executor/mod.rs
@@ -1,4 +1,4 @@
-use self::instrs::{execute_instrs, WasmOutcome};
+use self::instrs::{execute_instrs, WasmOutcome, CallKind};
 pub use super::Stack;
 use super::{stack::CallFrame, TaggedTrap};
 use crate::{
@@ -175,14 +175,15 @@ impl<'engine> EngineExecutor<'engine> {
                 WasmOutcome::Call {
                     results,
                     ref host_func,
+                    call_kind,
                 } => {
                     let instance = *self
                         .stack
                         .calls
                         .peek()
                         .expect("caller must be on the stack")
                         .instance();
-                    self.execute_host_func(&mut ctx, results, host_func, &instance)?;
+                    self.execute_host_func(&mut ctx, results, host_func, &instance, call_kind)?;
                 }
             }
         }
@@ -194,6 +195,7 @@ impl<'engine> EngineExecutor<'engine> {
         results: RegisterSpan,
         func: &Func,
         instance: &Instance,
+        call_kind: CallKind,
     ) -> Result<(), TaggedTrap> {
         let func_entity = match ctx.as_context().store.inner.resolve_func(func) {
             FuncEntity::Wasm(wasm_func) => {
@@ -206,6 +208,9 @@ impl<'engine> EngineExecutor<'engine> {
             func_entity,
             HostFuncCaller::wasm(results, instance),
         );
+        if matches!(call_kind, CallKind::Tail) {
+            self.stack.calls.pop();
+        }
         if self.stack.calls.peek().is_some() {
             // Case: There is a frame on the call stack.
             //