move host_coro_state out of the executor

Author: WhaleKit

crates/tinywasm/src/coro.rs

@@ -95,13 +95,13 @@ impl<R, State> PotentialCoroCallResult<R, State> {
             Self::Suspended(suspend, state) => PotentialCoroCallResult::Suspended(suspend, mapper(state)),
         }
     }
-    /// transform result with mapper if there is none - calls "otherwise".
+    /// transform result with res_mapper or state with state_mapper
     /// user_val passed to whichever is called and is guaranteed to be used
-    pub fn map<OutR, Usr, OutS>(
+    pub fn map<OutR, UserData, OutS>(
         self,
-        user_val: Usr,
-        res_mapper: impl FnOnce(R, Usr) -> OutR,
-        state_mapper: impl FnOnce(State, Usr) -> OutS,
+        user_val: UserData,
+        res_mapper: impl FnOnce(R, UserData) -> OutR,
+        state_mapper: impl FnOnce(State, UserData) -> OutS,
     ) -> PotentialCoroCallResult<OutR, OutS> {
         match self {
             Self::Return(res) => PotentialCoroCallResult::Return(res_mapper(res, user_val)),

crates/tinywasm/src/func.rs

@@ -162,7 +162,7 @@ pub(crate) type TypedFuncHandleCallOutcome<R> = crate::coro::PotentialCoroCallRe
 
 #[derive(Debug)]
 struct SuspendedWasmFunc {
-    runtime: interpreter::SuspendedRuntime,
+    runtime: interpreter::SuspendedInterpreterRuntime,
     result_types: Box<[ValType]>,
 }
 impl SuspendedWasmFunc {
@@ -178,8 +178,8 @@ impl SuspendedWasmFunc {
 #[derive(Debug)]
 #[allow(clippy::large_enum_variant)] // Wasm is bigger, but also much more common variant
 enum SuspendedFuncInner {
-    Wasm(SuspendedWasmFunc),
-    Host(SuspendedHostCoroState),
+    Wasm(SuspendedWasmFunc), // could also have host coro-function in the stack, but with some wasm below
+    Host(SuspendedHostCoroState), // in case we directly called host coro-function bypassing wasm stack
 }
 
 /// handle to function that was suspended and can be resumed

crates/tinywasm/src/interpreter/executor.rs

@@ -17,7 +17,7 @@ use crate::*;
 
 pub(crate) enum ReasonToBreak {
     Errored(Error),
-    Suspended(SuspendReason),
+    Suspended(SuspendReason, Option<SuspendedHostCoroState>),
     Finished,
 }
 
@@ -40,18 +40,17 @@ pub(crate) struct SuspendedHostCoroState {
 pub(crate) struct Executor<'store, 'stack> {
     pub(crate) cf: CallFrame,
     pub(crate) module: ModuleInstance,
-    pub(crate) suspended_host_coro: Option<SuspendedHostCoroState>,
     pub(crate) store: &'store mut Store,
     pub(crate) stack: &'stack mut Stack,
 }
 
-pub(crate) type ExecOutcome = coro::CoroStateResumeResult<()>;
+pub(crate) type ExecOutcome = coro::PotentialCoroCallResult<(), Option<SuspendedHostCoroState>>;
 
 impl<'store, 'stack> Executor<'store, 'stack> {
     pub(crate) fn new(store: &'store mut Store, stack: &'stack mut Stack) -> Result<Self> {
         let current_frame = stack.call_stack.pop().expect("no call frame, this is a bug");
         let current_module = store.get_module_instance_raw(current_frame.module_addr());
-        Ok(Self { cf: current_frame, module: current_module, suspended_host_coro: None, stack, store })
+        Ok(Self { cf: current_frame, module: current_module, stack, store })
     }
 
     #[inline(always)]
@@ -60,38 +59,51 @@ impl<'store, 'stack> Executor<'store, 'stack> {
             if let ControlFlow::Break(res) = self.exec_next() {
                 return match res {
                     ReasonToBreak::Errored(e) => Err(e),
-                    ReasonToBreak::Suspended(suspend_reason) => Ok(ExecOutcome::Suspended(suspend_reason)),
+                    ReasonToBreak::Suspended(suspend_reason, coro_state) => {
+                        Ok(ExecOutcome::Suspended(suspend_reason, coro_state))
+                    }
                     ReasonToBreak::Finished => Ok(ExecOutcome::Return(())),
                 };
             }
         }
     }
 
+    // on error running suspended_coro, returns it
+    // so that it could be retried later
     #[inline(always)]
-    pub(crate) fn resume(&mut self, res_arg: ResumeArgument) -> Result<ExecOutcome> {
-        if let Some(coro_state) = self.suspended_host_coro.as_mut() {
+    pub(crate) fn resume(
+        &mut self,
+        res_arg: ResumeArgument,
+        mut suspended_coro: Option<SuspendedHostCoroState>,
+    ) -> Result<ExecOutcome, (Error, Option<SuspendedHostCoroState>)> {
+        if let Some(coro_state) = suspended_coro.as_mut() {
             let ctx = FuncContext { store: self.store, module_addr: self.module.id() };
-            let host_res = coro_state.coro_state.resume(ctx, res_arg)?;
+            let host_res = match coro_state.coro_state.resume(ctx, res_arg) {
+                Ok(val) => val,
+                Err(err) => return Err((err, suspended_coro)),
+            };
+
             let res = match host_res {
                 CoroStateResumeResult::Return(res) => res,
                 CoroStateResumeResult::Suspended(suspend_reason) => {
-                    return Ok(ExecOutcome::Suspended(suspend_reason));
+                    return Ok(ExecOutcome::Suspended(suspend_reason, suspended_coro));
                 }
             };
             self.stack.values.extend_from_wasmvalues(&res);
-            self.suspended_host_coro = None;
 
             // we don't know how much time we spent in host function
-            if let ControlFlow::Break(ReasonToBreak::Suspended(reason)) = self.check_should_suspend() {
-                return Ok(ExecOutcome::Suspended(reason));
+            if let ControlFlow::Break(ReasonToBreak::Suspended(reason, coro_state)) = self.check_should_suspend() {
+                return Ok(ExecOutcome::Suspended(reason, coro_state));
             }
         }
 
         loop {
             if let ControlFlow::Break(res) = self.exec_next() {
                 return match res {
-                    ReasonToBreak::Errored(e) => Err(e),
-                    ReasonToBreak::Suspended(suspend_reason) => Ok(ExecOutcome::Suspended(suspend_reason)),
+                    ReasonToBreak::Errored(e) => Err((e, None)),
+                    ReasonToBreak::Suspended(suspend_reason, coro_state) => {
+                        Ok(ExecOutcome::Suspended(suspend_reason, coro_state))
+                    }
                     ReasonToBreak::Finished => Ok(ExecOutcome::Return(())),
                 };
             }
@@ -102,28 +114,28 @@ impl<'store, 'stack> Executor<'store, 'stack> {
     /// called when execution loops back, because that might happen indefinite amount of times
     /// and before and after function calls, because even without loops or infinite recursion, wasm function calls
     /// can mutliply time spent in execution
-    /// execution may not be suspended in the middle of execution the funcion:
+    /// execution may not be suspended in the middle of execution the innsruction (without rolling back its effects):
     /// so only do it as the last thing or first thing in the intsruction execution
     #[must_use = "If this returns ControlFlow::Break, the caller should propagate it"]
     fn check_should_suspend(&mut self) -> ControlFlow<ReasonToBreak> {
         if let Some(flag) = &self.store.suspend_cond.suspend_flag {
             if flag.load(core::sync::atomic::Ordering::Acquire) {
-                return ReasonToBreak::Suspended(SuspendReason::SuspendedFlag).into();
+                return ReasonToBreak::Suspended(SuspendReason::SuspendedFlag, None).into();
             }
         }
 
         #[cfg(feature = "std")]
         if let Some(when) = &self.store.suspend_cond.timeout_instant {
             if crate::std::time::Instant::now() >= *when {
-                return ReasonToBreak::Suspended(SuspendReason::SuspendedEpoch).into();
+                return ReasonToBreak::Suspended(SuspendReason::SuspendedEpoch, None).into();
             }
         }
 
         if let Some(mut cb) = self.store.suspend_cond.suspend_cb.take() {
             let should_suspend = matches!(cb(self.store), ControlFlow::Break(()));
             self.store.suspend_cond.suspend_cb = Some(cb); // put it back
             if should_suspend {
-                return ReasonToBreak::Suspended(SuspendReason::SuspendedCallback).into();
+                return ReasonToBreak::Suspended(SuspendReason::SuspendedCallback, None).into();
             }
         }
 
@@ -425,10 +437,9 @@ impl<'store, 'stack> Executor<'store, 'stack> {
                 ControlFlow::Continue(())
             }
             PotentialCoroCallResult::Suspended(suspend_reason, state) => {
-                self.suspended_host_coro =
-                    Some(SuspendedHostCoroState { coro_state: state, coro_orig_function: func_ref });
+                let coro_state = Some(SuspendedHostCoroState { coro_state: state, coro_orig_function: func_ref });
                 self.cf.incr_instr_ptr();
-                ReasonToBreak::Suspended(suspend_reason).into()
+                ReasonToBreak::Suspended(suspend_reason, coro_state).into()
             }
         }
     }

crates/tinywasm/src/interpreter/mod.rs

@@ -9,7 +9,7 @@ mod no_std_floats;
 
 use crate::coro;
 use crate::{FuncContext, ModuleInstance, Result, Store};
-use executor::{Executor, SuspendedHostCoroState};
+use executor::{ExecOutcome, Executor, SuspendedHostCoroState};
 use stack::{CallFrame, Stack};
 use tinywasm_types::ResumeArgument;
 pub use values::*;
@@ -21,72 +21,84 @@ pub use values::*;
 pub struct InterpreterRuntime {}
 
 #[derive(Debug)]
-pub(crate) struct SuspendedRuntimeBody {
-    pub(crate) suspended_host_coro: Option<SuspendedHostCoroState>,
-    pub(crate) module: ModuleInstance,
-    pub(crate) frame: CallFrame,
+struct SuspendedInterpreterRuntimeBody {
+    host_coro: Option<SuspendedHostCoroState>,
+    executor_stack: Stack,
+    executor_module: ModuleInstance,
+    executor_frame: CallFrame,
 }
 
-#[derive(Debug)]
-pub(crate) struct SuspendedRuntime {
-    pub(crate) body: Option<(SuspendedRuntimeBody, Stack)>,
-}
-impl SuspendedRuntime {
-    fn make_exec<'store, 'stack>(
-        body: SuspendedRuntimeBody,
-        stack: &'stack mut Stack,
-        store: &'store mut Store,
-    ) -> Executor<'store, 'stack> {
-        Executor { cf: body.frame, suspended_host_coro: body.suspended_host_coro, module: body.module, store, stack }
-    }
-    fn unmake_exec(exec: Executor<'_, '_>) -> SuspendedRuntimeBody {
-        SuspendedRuntimeBody { suspended_host_coro: exec.suspended_host_coro, module: exec.module, frame: exec.cf }
+impl SuspendedInterpreterRuntimeBody {
+    fn new(
+        host_coro: Option<SuspendedHostCoroState>,
+        executor_stack: Stack,
+        executor_module: ModuleInstance,
+        executor_frame: CallFrame,
+    ) -> Self {
+        Self { host_coro, executor_stack, executor_module, executor_frame }
     }
 }
 
-impl<'a> coro::CoroState<stack::Stack, FuncContext<'a>> for SuspendedRuntime {
-    fn resume(
-        &mut self,
-        ctx: FuncContext<'a>,
-        arg: ResumeArgument,
-    ) -> Result<coro::CoroStateResumeResult<stack::Stack>> {
-        // should be put back into self.body unless we're finished
-        let (body, mut stack) = if let Some(body_) = self.body.take() {
-            body_
+#[derive(Debug, Default)]
+pub(crate) struct SuspendedInterpreterRuntime(Option<SuspendedInterpreterRuntimeBody>);
+
+pub(crate) type InterpreterRuntimeExecOutcome = coro::PotentialCoroCallResult<Stack, SuspendedInterpreterRuntime>;
+pub(crate) type InterpreterRuntimeResumeOutcome = coro::CoroStateResumeResult<Stack>;
+
+impl<'_a> coro::CoroState<Stack, FuncContext<'_a>> for SuspendedInterpreterRuntime {
+    fn resume(&mut self, ctx: FuncContext<'_a>, arg: ResumeArgument) -> Result<InterpreterRuntimeResumeOutcome> {
+        let body = if let Some(body) = self.0.take() {
+            body
         } else {
-            return Err(crate::error::Error::InvalidResume);
+            // no suspended state to continue
+            return Result::Err(crate::Error::InvalidResume);
         };
 
-        let mut exec = Self::make_exec(body, &mut stack, ctx.store);
-        let resumed = match exec.resume(arg) {
-            Ok(resumed) => resumed,
-            Err(err) => {
-                self.body = Some((Self::unmake_exec(exec), stack));
-                return Err(err);
+        let SuspendedInterpreterRuntimeBody { host_coro, executor_stack, executor_module, executor_frame } = body;
+
+        let mut stack = executor_stack;
+        let mut exec =
+            executor::Executor { cf: executor_frame, module: executor_module, store: ctx.store, stack: &mut stack };
+        let res = match exec.resume(arg, host_coro) {
+            Ok(val) => val,
+            Err(e) => {
+                let Executor { cf, module, .. } = exec;
+                // pack back in case host_coro isn't corrupted and can be continued
+                self.0 = Some(SuspendedInterpreterRuntimeBody::new(e.1, stack, module, cf));
+                return Err(e.0);
             }
         };
-        match resumed {
-            executor::ExecOutcome::Return(()) => Ok(coro::CoroStateResumeResult::Return(stack)),
-            executor::ExecOutcome::Suspended(suspend) => {
-                self.body = Some((Self::unmake_exec(exec), stack));
-                Ok(coro::CoroStateResumeResult::Suspended(suspend))
+
+        return Ok(match res {
+            ExecOutcome::Return(()) => {
+                // we are finished
+                InterpreterRuntimeResumeOutcome::Return(stack)
+            }
+            ExecOutcome::Suspended(suspend_reason, host_coro) => {
+                // host_coro could be different host_coro than one we provided
+                let Executor { cf, module, .. } = exec;
+                self.0 = Some(SuspendedInterpreterRuntimeBody::new(host_coro, stack, module, cf));
+                InterpreterRuntimeResumeOutcome::Suspended(suspend_reason)
             }
-        }
+        });
     }
 }
 
-pub(crate) type RuntimeExecOutcome = coro::PotentialCoroCallResult<stack::Stack, SuspendedRuntime>;
-
 impl InterpreterRuntime {
-    pub(crate) fn exec(&self, store: &mut Store, stack: stack::Stack) -> Result<RuntimeExecOutcome> {
-        let mut stack = stack;
+    pub(crate) fn exec(&self, store: &mut Store, mut stack: Stack) -> Result<InterpreterRuntimeExecOutcome> {
         let mut executor = executor::Executor::new(store, &mut stack)?;
-        match executor.run_to_suspension()? {
-            coro::CoroStateResumeResult::Return(()) => Ok(RuntimeExecOutcome::Return(stack)),
-            coro::CoroStateResumeResult::Suspended(suspend) => Ok(RuntimeExecOutcome::Suspended(
-                suspend,
-                SuspendedRuntime { body: Some((SuspendedRuntime::unmake_exec(executor), stack)) },
-            )),
-        }
+        let result = executor.run_to_suspension()?;
+        Ok(match result {
+            ExecOutcome::Return(()) => InterpreterRuntimeExecOutcome::Return(stack),
+            ExecOutcome::Suspended(suspend_reason, host_coro) => {
+                let Executor { cf, module, .. } = executor;
+                InterpreterRuntimeExecOutcome::Suspended(
+                    suspend_reason,
+                    SuspendedInterpreterRuntime(Some(SuspendedInterpreterRuntimeBody::new(
+                        host_coro, stack, module, cf,
+                    ))),
+                )
+            }
+        })
     }
 }