support suspending start function when instantiating module

Author: WhaleKit

crates/tinywasm/src/coro.rs

@@ -24,6 +24,8 @@ pub enum SuspendReason {
     /// async should_suspend flag was set
     /// host shouldn't provide resume argument when calling resume
     SuspendedFlag,
+
+    // possible others: delimited continuations proposal, debugger breakpoint, out of fuel
 }
 
 /// result of a function that might pause in the middle and yield
@@ -43,6 +45,7 @@ pub enum PotentialCoroCallResult<R, State>
 /// doesn't need to have state, since it's contained in self
 #[derive(Debug)]
 pub enum CoroStateResumeResult<R> {
+    /// CoroState has finished
     /// after this CoroState::resume can't be called again on that CoroState
     Return(R),
 

crates/tinywasm/src/instance.rs

@@ -2,7 +2,7 @@ use alloc::{boxed::Box, format, rc::Rc, string::ToString};
 use tinywasm_types::*;
 
 use crate::func::{FromWasmValueTuple, IntoWasmValueTuple};
-use crate::{Error, FuncHandle, FuncHandleTyped, Imports, MemoryRef, MemoryRefMut, Module, Result, Store};
+use crate::{Error, FuncHandle, FuncHandleTyped, Imports, MemoryRef, MemoryRefMut, Module, PotentialCoroCallResult, Result, Store, SuspendFunc};
 
 /// An instanciated WebAssembly module
 ///
@@ -263,4 +263,18 @@ impl ModuleInstance {
         let _ = func.call(store, &[])?;
         Ok(Some(()))
     }
+
+    /// Invoke the start function of the module
+    /// 
+    /// Returns None if the module has no start function
+    /// If start function suspends, returns SuspededFunc.
+    /// Only when it finishes can this module instance be considered instantiated
+    pub fn start_coro(&self, store: &mut Store) -> Result<Option<PotentialCoroCallResult<(), SuspendFunc>>> {
+        let Some(func) = self.start_func(store)? else {
+            return Ok(None);
+        };
+
+        let res = func.call_coro(store, &[])?;
+        Ok(Some(res.map_result(|_|{()})))
+    }
 }

crates/tinywasm/src/lib.rs

@@ -96,7 +96,7 @@ pub use error::*;
 pub use func::{FuncHandle, FuncHandleTyped, SuspendFunc};
 pub use imports::*;
 pub use instance::ModuleInstance;
-pub use module::Module;
+pub use module::{IncompleteModule, Module};
 pub use reference::*;
 pub use store::*;
 

crates/tinywasm/src/module.rs

@@ -1,5 +1,5 @@
-use crate::{Imports, ModuleInstance, Result, Store};
-use tinywasm_types::TinyWasmModule;
+use crate::{CoroState, Imports, ModuleInstance, PotentialCoroCallResult, Result, Store, SuspendFunc};
+use tinywasm_types::{ResumeArgument, TinyWasmModule};
 
 /// A WebAssembly Module
 ///
@@ -56,4 +56,60 @@ impl Module {
         let _ = instance.start(store)?;
         Ok(instance)
     }
+
+    /// same as [Self::instantiate] but accounts for possibility of start function suspending, in which case it returns
+    /// [PotentialCoroCallResult::Suspended]. You can call [CoroState::resume] on it at any time to resume instantiation
+    pub fn instantiate_coro(
+        self,
+        store: &mut Store,
+        imports: Option<Imports>,
+    ) -> Result<PotentialCoroCallResult<ModuleInstance, IncompleteModule>> {
+        let instance = ModuleInstance::instantiate(store, self, imports)?;
+        let core_res = match instance.start_coro(store)? {
+            Some(res) => res,
+            None => return Ok(PotentialCoroCallResult::Return(instance)),
+        };
+        match core_res {
+            crate::PotentialCoroCallResult::Return(_) => return Ok(PotentialCoroCallResult::Return(instance)),
+            crate::PotentialCoroCallResult::Suspended(suspend_reason, state) => {
+                return Ok(PotentialCoroCallResult::Suspended(
+                    suspend_reason,
+                    IncompleteModule(Some(HitTheFloor(instance, state))),
+                ))
+            }
+        }
+    }
+}
+
+/// a corostate that results in [ModuleInstance] when finished
+#[derive(Debug)]
+pub struct IncompleteModule(Option<HitTheFloor>);
+
+#[derive(Debug)]
+struct HitTheFloor(ModuleInstance, SuspendFunc);
+
+impl<'a> CoroState<ModuleInstance, &'a mut Store> for IncompleteModule {
+    fn resume(&mut self, ctx: &mut Store, arg: ResumeArgument) -> Result<crate::CoroStateResumeResult<ModuleInstance>> {
+        let mut body: HitTheFloor = match self.0.take() {
+            Some(body) => body,
+            None => return Err(crate::Error::InvalidResume),
+        };
+        let coro_res = match body.1.resume(ctx, arg) {
+            Ok(res) => res,
+            Err(e) => {
+                self.0 = Some(body);
+                return Err(e);
+            }
+        };
+        match coro_res {
+            crate::CoroStateResumeResult::Return(_) => {
+                let res = body.0;
+                Ok(crate::CoroStateResumeResult::Return(res))
+            }
+            crate::CoroStateResumeResult::Suspended(suspend_reason) => {
+                self.0 = Some(body); // ...once told me
+                Ok(crate::CoroStateResumeResult::Suspended(suspend_reason))
+            }
+        }
+    }
 }