Implement resumable calls for TypedFunc (#605)

* implement resumable calls for TypedFunc

* add tests for TypedFunc::call_resumable

crates/wasmi/src/engine/mod.rs

@@ -26,7 +26,7 @@ pub use self::{
         RelativeDepth,
         TranslationError,
     },
-    resumable::{ResumableCall, ResumableInvocation},
+    resumable::{ResumableCall, ResumableInvocation, TypedResumableCall, TypedResumableInvocation},
     stack::StackLimits,
     traits::{CallParams, CallResults},
 };
@@ -36,6 +36,7 @@ use self::{
     code_map::CodeMap,
     executor::execute_frame,
     func_types::FuncTypeRegistry,
+    resumable::ResumableCallBase,
     stack::{FuncFrame, Stack, ValueStack},
 };
 pub(crate) use self::{
@@ -241,7 +242,7 @@ impl Engine {
         func: Func,
         params: impl CallParams,
         results: Results,
-    ) -> Result<ResumableCall<<Results as CallResults>::Results>, Trap>
+    ) -> Result<ResumableCallBase<<Results as CallResults>::Results>, Trap>
     where
         Results: CallResults,
     {
@@ -277,7 +278,7 @@ impl Engine {
         invocation: ResumableInvocation,
         params: impl CallParams,
         results: Results,
-    ) -> Result<ResumableCall<<Results as CallResults>::Results>, Trap>
+    ) -> Result<ResumableCallBase<<Results as CallResults>::Results>, Trap>
     where
         Results: CallResults,
     {
@@ -411,7 +412,7 @@ impl EngineInner {
         func: Func,
         params: impl CallParams,
         results: Results,
-    ) -> Result<ResumableCall<<Results as CallResults>::Results>, Trap>
+    ) -> Result<ResumableCallBase<<Results as CallResults>::Results>, Trap>
     where
         Results: CallResults,
     {
@@ -426,7 +427,7 @@ impl EngineInner {
         match results {
             Ok(results) => {
                 self.stacks.lock().recycle(stack);
-                Ok(ResumableCall::Finished(results))
+                Ok(ResumableCallBase::Finished(results))
             }
             Err(TaggedTrap::Wasm(trap)) => {
                 self.stacks.lock().recycle(stack);
@@ -435,7 +436,7 @@ impl EngineInner {
             Err(TaggedTrap::Host {
                 host_func,
                 host_trap,
-            }) => Ok(ResumableCall::Resumable(ResumableInvocation::new(
+            }) => Ok(ResumableCallBase::Resumable(ResumableInvocation::new(
                 ctx.as_context().store.engine().clone(),
                 func,
                 host_func,
@@ -451,7 +452,7 @@ impl EngineInner {
         mut invocation: ResumableInvocation,
         params: impl CallParams,
         results: Results,
-    ) -> Result<ResumableCall<<Results as CallResults>::Results>, Trap>
+    ) -> Result<ResumableCallBase<<Results as CallResults>::Results>, Trap>
     where
         Results: CallResults,
     {
@@ -462,7 +463,7 @@ impl EngineInner {
         match results {
             Ok(results) => {
                 self.stacks.lock().recycle(invocation.take_stack());
-                Ok(ResumableCall::Finished(results))
+                Ok(ResumableCallBase::Finished(results))
             }
             Err(TaggedTrap::Wasm(trap)) => {
                 self.stacks.lock().recycle(invocation.take_stack());
@@ -473,7 +474,7 @@ impl EngineInner {
                 host_trap,
             }) => {
                 invocation.update(host_func, host_trap);
-                Ok(ResumableCall::Resumable(invocation))
+                Ok(ResumableCallBase::Resumable(invocation))
             }
         }
     }

crates/wasmi/src/engine/resumable.rs

@@ -1,18 +1,45 @@
 use super::Func;
-use crate::{engine::Stack, AsContextMut, Engine, Error};
-use core::mem::replace;
+use crate::{engine::Stack, func::CallResultsTuple, AsContextMut, Engine, Error, WasmResults};
+use core::{fmt, marker::PhantomData, mem::replace, ops::Deref};
 use wasmi_core::{Trap, Value};
 
-/// Returned by calling a function in a resumable way.
+/// Returned by [`Engine`] methods for calling a function in a resumable way.
+///
+/// # Note
+///
+/// This is the base type for resumable call results and can be converted into
+/// either the dynamically typed [`ResumableCall`] or the statically typed
+/// [`TypedResumableCall`] that act as user facing API. Therefore this type
+/// must provide all the information necessary to be properly converted into
+/// either user facing types.
 #[derive(Debug)]
-pub enum ResumableCall<T> {
+pub(crate) enum ResumableCallBase<T> {
     /// The resumable call has finished properly and returned a result.
     Finished(T),
     /// The resumable call encountered a host error and can be resumed.
     Resumable(ResumableInvocation),
 }
 
-/// State required to resume a function invocation.
+/// Returned by calling a [`Func`] in a resumable way.
+#[derive(Debug)]
+pub enum ResumableCall {
+    /// The resumable call has finished properly and returned a result.
+    Finished,
+    /// The resumable call encountered a host error and can be resumed.
+    Resumable(ResumableInvocation),
+}
+
+impl ResumableCall {
+    /// Creates a [`ResumableCall`] from the [`Engine`]'s base [`ResumableCallBase`].
+    pub(crate) fn new(call: ResumableCallBase<()>) -> Self {
+        match call {
+            ResumableCallBase::Finished(()) => Self::Finished,
+            ResumableCallBase::Resumable(invocation) => Self::Resumable(invocation),
+        }
+    }
+}
+
+/// State required to resume a [`Func`] invocation.
 #[derive(Debug)]
 pub struct ResumableInvocation {
     /// The engine in use for the function invokation.
@@ -120,9 +147,9 @@ impl ResumableInvocation {
         &self.host_error
     }
 
-    /// Resumes the call to the Wasm or host function with the given inputs.
+    /// Resumes the call to the [`Func`] with the given inputs.
     ///
-    /// The result is written back into the `outputs` buffer.
+    /// The result is written back into the `outputs` buffer upon success.
     ///
     /// Returns a resumable handle to the function invocation upon
     /// enountering host errors with which it is possible to handle
@@ -140,7 +167,7 @@ impl ResumableInvocation {
         mut ctx: impl AsContextMut<UserState = T>,
         inputs: &[Value],
         outputs: &mut [Value],
-    ) -> Result<ResumableCall<()>, Error> {
+    ) -> Result<ResumableCall, Error> {
         self.engine
             .resolve_func_type(self.host_func().signature(ctx.as_context()), |func_type| {
                 func_type.match_results(inputs, true)
@@ -155,5 +182,102 @@ impl ResumableInvocation {
             .clone()
             .resume_func(ctx.as_context_mut(), self, inputs, outputs)
             .map_err(Into::into)
+            .map(ResumableCall::new)
+    }
+}
+
+/// Returned by calling a [`TypedFunc`] in a resumable way.
+///
+/// [`TypedFunc`]: [`crate::TypedFunc`]
+#[derive(Debug)]
+pub enum TypedResumableCall<T> {
+    /// The resumable call has finished properly and returned a result.
+    Finished(T),
+    /// The resumable call encountered a host error and can be resumed.
+    Resumable(TypedResumableInvocation<T>),
+}
+
+impl<Results> TypedResumableCall<Results> {
+    /// Creates a [`TypedResumableCall`] from the [`Engine`]'s base [`ResumableCallBase`].
+    pub(crate) fn new(call: ResumableCallBase<Results>) -> Self {
+        match call {
+            ResumableCallBase::Finished(results) => Self::Finished(results),
+            ResumableCallBase::Resumable(invocation) => {
+                Self::Resumable(TypedResumableInvocation::new(invocation))
+            }
+        }
+    }
+}
+
+/// State required to resume a [`TypedFunc`] invocation.
+///
+/// [`TypedFunc`]: [`crate::TypedFunc`]
+pub struct TypedResumableInvocation<Results> {
+    invocation: ResumableInvocation,
+    /// The parameter and result typed encoded in Rust type system.
+    results: PhantomData<fn() -> Results>,
+}
+
+impl<Results> TypedResumableInvocation<Results> {
+    /// Creates a [`TypedResumableInvocation`] wrapper for the given [`ResumableInvocation`].
+    pub(crate) fn new(invocation: ResumableInvocation) -> Self {
+        Self {
+            invocation,
+            results: PhantomData,
+        }
+    }
+
+    /// Resumes the call to the [`TypedFunc`] with the given inputs.
+    ///
+    /// Returns a resumable handle to the function invocation upon
+    /// enountering host errors with which it is possible to handle
+    /// the error and continue the execution as if no error occured.
+    ///
+    /// # Errors
+    ///
+    /// - If the function resumption returned a Wasm [`Trap`].
+    /// - If the types or the number of values in `inputs` does not match
+    ///   the types and number of result values of the errorneous host function.
+    ///
+    /// [`TypedFunc`]: [`crate::TypedFunc`]
+    pub fn resume<T>(
+        self,
+        mut ctx: impl AsContextMut<UserState = T>,
+        inputs: &[Value],
+    ) -> Result<TypedResumableCall<Results>, Error>
+    where
+        Results: WasmResults,
+    {
+        self.engine
+            .resolve_func_type(self.host_func().signature(ctx.as_context()), |func_type| {
+                func_type.match_results(inputs, true)
+            })?;
+        self.engine
+            .clone()
+            .resume_func(
+                ctx.as_context_mut(),
+                self.invocation,
+                inputs,
+                <CallResultsTuple<Results>>::default(),
+            )
+            .map_err(Into::into)
+            .map(TypedResumableCall::new)
+    }
+}
+
+impl<Results> Deref for TypedResumableInvocation<Results> {
+    type Target = ResumableInvocation;
+
+    fn deref(&self) -> &Self::Target {
+        &self.invocation
+    }
+}
+
+impl<Results> fmt::Debug for TypedResumableInvocation<Results> {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        f.debug_struct("TypedResumableInvocation")
+            .field("invocation", &self.invocation)
+            .field("results", &self.results)
+            .finish()
     }
 }

crates/wasmi/src/func/mod.rs

@@ -3,6 +3,7 @@ mod error;
 mod into_func;
 mod typed_func;
 
+pub(crate) use self::typed_func::CallResultsTuple;
 pub use self::{
     caller::Caller,
     error::FuncError,
@@ -328,7 +329,7 @@ impl Func {
         mut ctx: impl AsContextMut<UserState = T>,
         inputs: &[Value],
         outputs: &mut [Value],
-    ) -> Result<ResumableCall<()>, Error> {
+    ) -> Result<ResumableCall, Error> {
         self.verify_and_prepare_inputs_outputs(ctx.as_context(), inputs, outputs)?;
         // Note: Cloning an [`Engine`] is intentionally a cheap operation.
         ctx.as_context()
@@ -337,6 +338,7 @@ impl Func {
             .clone()
             .execute_func_resumable(ctx.as_context_mut(), *self, inputs, outputs)
             .map_err(Into::into)
+            .map(ResumableCall::new)
     }
 
     /// Verify that the `inputs` and `outputs` value types match the function signature.

crates/wasmi/src/func/typed_func.rs

@@ -4,6 +4,7 @@ use crate::{
     AsContext,
     AsContextMut,
     Error,
+    TypedResumableCall,
 };
 use core::{fmt, fmt::Debug, marker::PhantomData};
 use wasmi_core::{Trap, UntypedValue};
@@ -78,7 +79,7 @@ where
         })
     }
 
-    /// Invokes this Wasm or host function with the specified parameters.
+    /// Calls this Wasm or host function with the specified parameters.
     ///
     /// Returns either the results of the call, or a [`Trap`] if one happened.
     ///
@@ -101,6 +102,40 @@ where
             <CallResultsTuple<Results>>::default(),
         )
     }
+
+    /// Calls this Wasm or host function with the specified parameters.
+    ///
+    /// Returns a resumable handle to the function invocation upon
+    /// enountering host errors with which it is possible to handle
+    /// the error and continue the execution as if no error occured.
+    ///
+    /// # Note
+    ///
+    /// This is a non-standard WebAssembly API and might not be available
+    /// at other WebAssembly engines. Please be aware that depending on this
+    /// feature might mean a lock-in to `wasmi` for users.
+    ///
+    /// # Errors
+    ///
+    /// If the function returned a [`Trap`] originating from WebAssembly.
+    pub fn call_resumable(
+        &self,
+        mut ctx: impl AsContextMut,
+        params: Params,
+    ) -> Result<TypedResumableCall<Results>, Trap> {
+        // Note: Cloning an [`Engine`] is intentionally a cheap operation.
+        ctx.as_context()
+            .store
+            .engine()
+            .clone()
+            .execute_func_resumable(
+                ctx.as_context_mut(),
+                self.func,
+                params,
+                <CallResultsTuple<Results>>::default(),
+            )
+            .map(TypedResumableCall::new)
+    }
 }
 
 impl<Params> CallParams for Params

crates/wasmi/src/lib.rs

@@ -121,7 +121,15 @@ pub mod errors {
 }
 
 pub use self::{
-    engine::{Config, Engine, ResumableCall, ResumableInvocation, StackLimits},
+    engine::{
+        Config,
+        Engine,
+        ResumableCall,
+        ResumableInvocation,
+        StackLimits,
+        TypedResumableCall,
+        TypedResumableInvocation,
+    },
     error::Error,
     external::Extern,
     func::{Caller, Func, IntoFunc, TypedFunc, WasmParams, WasmResults, WasmRet, WasmType},