Implement resumable calls via register-machine wasmi engine backend (…

…#814)

* implement resumable calling via register-machine engine

This is not tested properly and may not work, yet.

* try to fix fuzzing CI

* fix bug for older rust toolchain (bench CI)

* make special smoldot tail resume tests pass under regmach

* apply rustfmt

* remove outdated comment

Tails calls work when calling host functions in the new register-machine engine.

* add TestData to resumable calls tests

Also use Linker::func_wrap.
New TestData is unused, yet.

* test resumable calls for both engine backends

.github/workflows/rust.yml

@@ -207,7 +207,7 @@ jobs:
       - uses: actions-rs/toolchain@v1
         with:
           profile: minimal
-          toolchain: nightly-2023-06-15 # we must tag pinned nightly version: https://substrate.stackexchange.com/a/9069
+          toolchain: nightly
           override: true
       - name: Set up Cargo cache
         uses: actions/cache@v3

crates/wasmi/src/engine/mod.rs

@@ -409,6 +409,11 @@ impl Engine {
     pub(crate) fn recycle_stack(&self, stack: Stack) {
         self.inner.recycle_stack(stack)
     }
+
+    /// Recycles the given [`Stack`] for reuse in the [`Engine`].
+    pub(crate) fn recycle_stack_2(&self, stack: Stack2) {
+        self.inner.recycle_stack_2(stack)
+    }
 }
 
 /// The internal state of the `wasmi` [`Engine`].
@@ -684,7 +689,42 @@ impl EngineInner {
         results
     }
 
+    /// Executes the given [`Func`] resumably with the given `params` and returns the `results`.
+    ///
+    /// Uses the [`StoreContextMut`] for context information about the Wasm [`Store`].
+    ///
+    /// # Errors
+    ///
+    /// If the Wasm execution traps or runs out of resources.
     fn execute_func_resumable<T, Results>(
+        &self,
+        ctx: StoreContextMut<T>,
+        func: &Func,
+        params: impl CallParams,
+        results: Results,
+    ) -> Result<ResumableCallBase<<Results as CallResults>::Results>, Trap>
+    where
+        Results: CallResults,
+    {
+        match self.config().engine_backend() {
+            EngineBackend::StackMachine => {
+                self.execute_func_resumable_stackmach(ctx, func, params, results)
+            }
+            EngineBackend::RegisterMachine => {
+                self.execute_func_resumable_regmach(ctx, func, params, results)
+            }
+        }
+    }
+
+    /// Executes the given [`Func`] resumably with the given `params` and returns the `results`.
+    ///
+    /// - Uses the `wasmi` stack-machine based engine backend.
+    /// - Uses the [`StoreContextMut`] for context information about the Wasm [`Store`].
+    ///
+    /// # Errors
+    ///
+    /// If the Wasm execution traps or runs out of resources.
+    fn execute_func_resumable_stackmach<T, Results>(
         &self,
         mut ctx: StoreContextMut<T>,
         func: &Func,
@@ -719,12 +759,48 @@ impl EngineInner {
                 *func,
                 host_func,
                 host_trap,
+                None,
                 stack,
             ))),
         }
     }
 
+    /// Resumes the given [`Func`] with the given `params` and returns the `results`.
+    ///
+    /// - Uses the [`StoreContextMut`] for context information about the Wasm [`Store`].
+    ///
+    /// # Errors
+    ///
+    /// If the Wasm execution traps or runs out of resources.
     fn resume_func<T, Results>(
+        &self,
+        ctx: StoreContextMut<T>,
+        invocation: ResumableInvocation,
+        params: impl CallParams,
+        results: Results,
+    ) -> Result<ResumableCallBase<<Results as CallResults>::Results>, Trap>
+    where
+        Results: CallResults,
+    {
+        match self.config().engine_backend() {
+            EngineBackend::StackMachine => {
+                self.resume_func_stackmach(ctx, invocation, params, results)
+            }
+            EngineBackend::RegisterMachine => {
+                self.resume_func_regmach(ctx, invocation, params, results)
+            }
+        }
+    }
+
+    /// Resumes the given [`Func`] with the given `params` and returns the `results`.
+    ///
+    /// - Uses the `wasmi` stack-machine based engine backend.
+    /// - Uses the [`StoreContextMut`] for context information about the Wasm [`Store`].
+    ///
+    /// # Errors
+    ///
+    /// If the Wasm execution traps or runs out of resources.
+    fn resume_func_stackmach<T, Results>(
         &self,
         ctx: StoreContextMut<T>,
         mut invocation: ResumableInvocation,
@@ -736,30 +812,37 @@ impl EngineInner {
     {
         let res = self.res.read();
         let host_func = invocation.host_func();
-        let results = EngineExecutor::new(&res, &mut invocation.stack)
-            .resume_func(ctx, host_func, params, results);
+        let mut stack = invocation.take_stack().into_stackmach();
+        let results =
+            EngineExecutor::new(&res, &mut stack).resume_func(ctx, host_func, params, results);
         match results {
             Ok(results) => {
-                self.stacks.lock().recycle(invocation.take_stack());
+                self.stacks.lock().recycle(stack);
                 Ok(ResumableCallBase::Finished(results))
             }
             Err(TaggedTrap::Wasm(trap)) => {
-                self.stacks.lock().recycle(invocation.take_stack());
+                self.stacks.lock().recycle(stack);
                 Err(trap)
             }
             Err(TaggedTrap::Host {
                 host_func,
                 host_trap,
             }) => {
-                invocation.update(host_func, host_trap);
+                invocation.update(stack, host_func, host_trap);
                 Ok(ResumableCallBase::Resumable(invocation))
             }
         }
     }
 
+    /// Recycles the given [`Stack`] for the stack-machine `wasmi` engine backend.
     fn recycle_stack(&self, stack: Stack) {
         self.stacks.lock().recycle(stack);
     }
+
+    /// Recycles the given [`Stack`] for the register-machine `wasmi` engine backend.
+    fn recycle_stack_2(&self, stack: Stack2) {
+        self.stacks.lock().recycle_2(stack)
+    }
 }
 
 /// Engine resources that are immutable during function execution.

crates/wasmi/src/engine/regmach/executor/instrs.rs

@@ -1,3 +1,4 @@
+pub use self::call::CallKind;
 use self::{call::CallOutcome, return_::ReturnOutcome};
 use crate::{
     core::{TrapCode, UntypedValue},
@@ -45,8 +46,17 @@ 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 +83,7 @@ pub enum WasmOutcome {
     Call {
         results: RegisterSpan,
         host_func: Func,
+        call_kind: CallKind,
     },
 }
 

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,
     },
 }
 
@@ -320,8 +321,6 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
                 Ok(CallOutcome::Continue)
             }
             FuncEntity::Host(host_func) => {
-                // Note: host function calls cannot be implemented as tail calls.
-                //       The Wasm spec is not mandating tail behavior for host calls.
                 let (input_types, output_types) = self
                     .func_types
                     .resolve_func_type(host_func.ty_dedup())
@@ -341,13 +340,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,
                 })
             }
         }

crates/wasmi/src/engine/regmach/executor/mod.rs

@@ -1,6 +1,6 @@
-use self::instrs::{execute_instrs, WasmOutcome};
-use super::stack::CallFrame;
+use self::instrs::{execute_instrs, CallKind, WasmOutcome};
 pub use super::Stack;
+use super::{stack::CallFrame, TaggedTrap};
 use crate::{
     engine::{
         cache::InstanceCache,
@@ -10,10 +10,8 @@ use crate::{
         },
         CallParams,
         CallResults,
-        DedupFuncType,
         EngineResources,
         FuncParams,
-        TaggedTrap,
     },
     func::HostFuncEntity,
     AsContext,
@@ -67,10 +65,8 @@ impl<'engine> EngineExecutor<'engine> {
         Results: CallResults,
     {
         self.stack.reset();
-        let func_type;
         match ctx.as_context().store.inner.resolve_func(func) {
             FuncEntity::Wasm(wasm_func) => {
-                func_type = *wasm_func.ty_dedup();
                 // We reserve space on the stack to write the results of the root function execution.
                 self.stack.values.extend_zeros(results.len_results());
                 let instance = wasm_func.instance();
@@ -92,7 +88,6 @@ impl<'engine> EngineExecutor<'engine> {
                 self.execute_func(ctx.as_context_mut())?;
             }
             FuncEntity::Host(host_func) => {
-                func_type = *host_func.ty_dedup();
                 // The host function signature is required for properly
                 // adjusting, inspecting and manipulating the value stack.
                 let (input_types, output_types) = self
@@ -115,7 +110,44 @@ impl<'engine> EngineExecutor<'engine> {
                 self.dispatch_host_func(ctx.as_context_mut(), host_func, HostFuncCaller::Root)?;
             }
         };
-        let results = self.write_results_back(results, func_type);
+        let results = self.write_results_back(results);
+        Ok(results)
+    }
+
+    /// Resumes the execution of the given [`Func`] using `params`.
+    ///
+    /// Stores the execution result into `results` upon a successful execution.
+    ///
+    /// # Errors
+    ///
+    /// - If the given `params` do not match the expected parameters of `func`.
+    /// - If the given `results` do not match the the length of the expected results of `func`.
+    /// - When encountering a Wasm or host trap during the execution of `func`.
+    pub fn resume_func<T, Results>(
+        &mut self,
+        mut ctx: StoreContextMut<T>,
+        _host_func: Func,
+        params: impl CallParams,
+        caller_results: RegisterSpan,
+        results: Results,
+    ) -> Result<<Results as CallResults>::Results, TaggedTrap>
+    where
+        Results: CallResults,
+    {
+        let caller = self
+            .stack
+            .calls
+            .peek()
+            .expect("must have caller call frame on stack upon function resumption");
+        let mut caller_sp = unsafe { self.stack.values.stack_ptr_at(caller.base_offset()) };
+        let call_params = params.call_params();
+        let len_params = call_params.len();
+        for (result, param) in caller_results.iter(len_params).zip(call_params) {
+            let cell = unsafe { caller_sp.get_mut(result) };
+            *cell = param;
+        }
+        self.execute_func(ctx.as_context_mut())?;
+        let results = self.write_results_back(results);
         Ok(results)
     }
 
@@ -143,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)?;
                 }
             }
         }
@@ -162,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) => {
@@ -174,12 +208,15 @@ 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.
             //
             // This is the default case and we can easily make host function
             // errors return a resumable call handle.
-            result.map_err(|trap| TaggedTrap::host(*func, trap))?;
+            result.map_err(|trap| TaggedTrap::host(*func, trap, results))?;
         } else {
             // Case: No frame is on the call stack. (edge case)
             //
@@ -362,15 +399,11 @@ impl<'engine> EngineExecutor<'engine> {
     ///
     /// - If the `results` buffer length does not match the remaining amount of stack values.
     #[inline]
-    fn write_results_back<Results>(
-        &mut self,
-        results: Results,
-        ty: DedupFuncType,
-    ) -> <Results as CallResults>::Results
+    fn write_results_back<Results>(&mut self, results: Results) -> <Results as CallResults>::Results
     where
         Results: CallResults,
     {
-        let len_results = self.res.func_types.resolve_func_type(&ty).results().len();
+        let len_results = results.len_results();
         results.call_results(&self.stack.values.as_slice()[..len_results])
     }
 }