test(jco): add basic async num returning test

crates/js-component-bindgen/src/core.rs

@@ -64,6 +64,7 @@
     },
 }
 
+#[derive(Debug)]
 pub enum AugmentedImport<'a> {
     CoreDef(&'a CoreDef),
     Memory { mem: &'a CoreDef, op: AugmentedOp },
@@ -534,7 +535,7 @@
     ($( @$proposal:ident $op:ident $({ $($arg:ident: $argty:ty),* })? => $visit:ident ($($ann:tt)*))*) => {
         $(
             #[allow(unreachable_code)]
             fn $visit(&mut self $( $( ,$arg: $argty)* )?) {
                 define_visit!(augment self $op $($($arg)*)?);
             }
         )*

crates/js-component-bindgen/src/esm_bindgen.rs

@@ -396,6 +396,9 @@ impl EsmBindgen {
                             maybe_quote_member(specifier)
                         );
                         for (external_name, local_name) in bound_external_names {
+                            // For imports that are functions, ensure that they are noted as host provided
+                            uwriteln!(output, "{local_name}._isHostProvided = true;");
+
                             uwriteln!(
                                 output,
                                 r#"
@@ -436,6 +439,8 @@ impl EsmBindgen {
                     } else {
                         uwriteln!(output, "{local_name} from '{specifier}';");
                     }
+                    uwriteln!(output, "{local_name}._isHostProvided = true;");
+
                     for other_local_name in &binding_local_names[1..] {
                         uwriteln!(output, "const {other_local_name} = {local_name};");
                     }
@@ -473,9 +478,13 @@ impl EsmBindgen {
             }
             uwriteln!(output, "}} = {iface_local_name};");
 
-            // Ensure that the imports we destructured were defined
-            // (if they were not, the user is likely missing an import @ instantiation time)
+            // Process all external host-provided imports
             for (member_name, local_name) in generated_member_names {
+                // For imports that are functions, ensure that they are noted as host provided
+                uwriteln!(output, "{local_name}._isHostProvided = true;");
+
+                // Ensure that the imports we destructured were defined
+                // (if they were not, the user is likely missing an import @ instantiation time)
                 uwriteln!(
                     output,
                     r#"

crates/js-component-bindgen/src/function_bindgen.rs

@@ -322,15 +322,19 @@ impl FunctionBindgen<'_> {
     }
 
     /// Start the current task
+    ///
+    /// The code generated by this function *may* also start a subtask
+    /// where appropriate.
     fn start_current_task(&mut self, instr: &Instruction) {
         let is_async = self.is_async;
         let fn_name = self.callee;
         let err_handling = self.err.to_js_string();
-        let callback_fn_name = self
+        let callback_fn_js = self
             .canon_opts
             .callback
             .as_ref()
-            .map(|v| format!("callback_{}", v.as_u32()));
+            .map(|v| format!("callback_{}", v.as_u32()))
+            .unwrap_or_else(|| "null".into());
         let prefix = match instr {
             Instruction::CallWasm { .. } => "_wasm_call_",
             Instruction::CallInterface { .. } => "_interface_call_",
@@ -344,18 +348,16 @@ impl FunctionBindgen<'_> {
 
         uwriteln!(
             self.src,
-            "
-            const [_, {prefix}currentTaskID] = {start_current_task_fn}({{
-                componentIdx: {component_instance_idx},
-                isAsync: {is_async},
-                entryFnName: '{fn_name}',
-                getCallbackFn: () => {callback_fn_name},
-                callbackFnName: '{callback_fn_name}',
-                errHandling: '{err_handling}',
-            }});
-            ",
-            // NOTE: callback functions are missing on async imports that are host defined
-            callback_fn_name = callback_fn_name.unwrap_or_else(|| "null".into()),
+            r#"
+              const [task, {prefix}currentTaskID] = {start_current_task_fn}({{
+                  componentIdx: {component_instance_idx},
+                  isAsync: {is_async},
+                  entryFnName: '{fn_name}',
+                  getCallbackFn: () => {callback_fn_js},
+                  callbackFnName: '{callback_fn_js}',
+                  errHandling: '{err_handling}',
+              }});
+            "#,
         );
     }
 
@@ -1256,7 +1258,12 @@ impl Bindgen for FunctionBindgen<'_> {
                     has_post_return = self.post_return.is_some(),
                 );
 
+                // Write out whether the caller was host provided
+                // (if we're calling into wasm then we know it was not)
+                uwriteln!(self.src, "const hostProvided = false;");
+
                 // Inject machinery for starting a 'current' task
+                // (this will define the 'task' variable)
                 self.start_current_task(inst);
 
                 // TODO: trap if this component is already on the call stack (re-entrancy)
@@ -1307,31 +1314,98 @@ impl Bindgen for FunctionBindgen<'_> {
 
             // Call to an interface, usually but not always an externally imported interface
             Instruction::CallInterface { func, async_ } => {
-                // let get_or_create_async_state_fn = self.intrinsic(Intrinsic::Component(
-                //     ComponentIntrinsic::GetOrCreateAsyncState,
-                // ));
-                // let current_task_get_fn =
-                //     self.intrinsic(Intrinsic::AsyncTask(AsyncTaskIntrinsic::GetCurrentTask));
-                // let component_instance_idx = self.canon_opts.instance.as_u32();
-
                 let debug_log_fn = self.intrinsic(Intrinsic::DebugLog);
+                let start_current_task_fn =
+                    self.intrinsic(Intrinsic::AsyncTask(AsyncTaskIntrinsic::StartCurrentTask));
+                let current_task_get_fn =
+                    self.intrinsic(Intrinsic::AsyncTask(AsyncTaskIntrinsic::GetCurrentTask));
+
                 uwriteln!(
                     self.src,
-                    "{debug_log_fn}('{prefix} [Instruction::CallInterface] (async? {async_}, @ enter)');",
+                    "{debug_log_fn}('{prefix} [Instruction::CallInterface] ({async_}, @ enter)');",
                     prefix = self.tracing_prefix,
                     async_ = async_.then_some("async").unwrap_or("sync"),
                 );
 
-                // // Inject machinery for starting a 'current' task
-                // self.start_current_task(
-                //     inst,
-                //     *async_,
-                //     &func.name,
-                //     self.canon_opts
-                //         .callback
-                //         .as_ref()
-                //         .map(|v| format!("callback_{}", v.as_u32())),
-                // );
+                // Determine the callee function and arguments
+                let (fn_js, args_js) = if self.callee_resource_dynamic {
+                    (
+                        format!("{}.{}", operands[0], self.callee),
+                        operands[1..].join(", "),
+                    )
+                } else {
+                    (self.callee.into(), operands.join(", "))
+                };
+
+                // Write out whether the caller was host provided
+                uwriteln!(self.src, "const hostProvided = {fn_js}._isHostProvided;");
+
+                let component_instance_idx = self.canon_opts.instance.as_u32();
+
+                // Start the necessary subtasks and/or host task
+                //
+                // We must create a subtask in the case of an async host import.
+                //
+                // If there's no parent task, we're not executing in a subtask situation,
+                // so we can just create the new task and immediately continue execution.
+                //
+                // If there *is* a parent task, then we are likely about to create new task that
+                // matches/belongs to an existing subtask in the parent task.
+                //
+                // If we're dealing with a function that has been marked as a host import, then
+                // we expect that `Trampoline::LowerImport` and relevant intrinsics were called before
+                // this, and a subtask has been set up.
+                //
+                // TODO: we use getLatestSubtask(), but could ordering change? Not under threads (assuming same thread)?
+                uwriteln!(
+                    self.src,
+                    r#"
+                    let parentTask;
+                    let task;
+                    let subtask;
+
+                    const createTask = () => {{
+                        const results = {start_current_task_fn}({{
+                            componentIdx: {component_instance_idx},
+                            isAsync: {is_async},
+                            entryFnName: '{fn_name}',
+                            getCallbackFn: () => {callback_fn_js},
+                            callbackFnName: '{callback_fn_js}',
+                            errHandling: '{err_handling}',
+                        }});
+                        task = results[0];
+                    }};
+
+                    taskCreation: {{
+                        parentTask = {current_task_get_fn}({component_instance_idx})?.task;
+                        if (!parentTask) {{
+                            createTask();
+                            break taskCreation;
+                        }}
+
+                        createTask();
+
+                        const isHostAsyncImport = hostProvided && {is_async};
+                        if (isHostAsyncImport) {{
+                            subtask = parentTask.getLatestSubtask();
+                            if (!subtask) {{
+                                throw new Error("Missing subtask for host import call, has the import been lowered?");
+                            }}
+                            subtask.setChildTask(task);
+                            task.setParentSubtask(subtask);
+                        }}
+                    }}
+                    "#,
+                    is_async = self.is_async,
+                    fn_name = self.callee,
+                    err_handling = self.err.to_js_string(),
+                    callback_fn_js = self
+                        .canon_opts
+                        .callback
+                        .as_ref()
+                        .map(|v| format!("callback_{}", v.as_u32()))
+                        .unwrap_or_else(|| "null".into()),
+                );
 
                 let results_length = if func.result.is_none() { 0 } else { 1 };
                 let maybe_await = if self.requires_async_porcelain | async_ {
@@ -1340,17 +1414,8 @@ impl Bindgen for FunctionBindgen<'_> {
                     ""
                 };
 
-                // Build the call
-                let call = if self.callee_resource_dynamic {
-                    format!(
-                        "{maybe_await} {}.{}({})",
-                        operands[0],
-                        self.callee,
-                        operands[1..].join(", ")
-                    )
-                } else {
-                    format!("{maybe_await} {}({})", self.callee, operands.join(", "))
-                };
+                // Build the JS expression that calls the callee
+                let call = format!("{maybe_await} {fn_js}({args_js})",);
 
                 match self.err {
                     // If configured to do *no* error handling at all or throw
@@ -1379,13 +1444,13 @@ impl Bindgen for FunctionBindgen<'_> {
                         uwriteln!(
                             self.src,
                             r#"
-                          let ret;
-                          try {{
-                              ret = {{ tag: 'ok', val: {call} }};
-                          }} catch (e) {{
-                              ret = {{ tag: 'err', val: {err_payload}(e) }};
-                          }}
-                       "#,
+                            let ret;
+                            try {{
+                                ret = {{ tag: 'ok', val: {call} }};
+                            }} catch (e) {{
+                                ret = {{ tag: 'err', val: {err_payload}(e) }};
+                            }}
+                            "#,
                         );
                         results.push("ret".to_string());
                     }
@@ -1452,10 +1517,10 @@ impl Bindgen for FunctionBindgen<'_> {
                     self.clear_resource_borrows = false;
                 }
 
-                // // For non-async calls, the current task can end immediately
-                // if !async_ {
-                //     self.end_current_task();
-                // }
+                // For non-async calls, the current task can end immediately
+                if !async_ {
+                    self.end_current_task();
+                }
             }
 
             Instruction::Return {
@@ -2162,16 +2227,21 @@ impl Bindgen for FunctionBindgen<'_> {
             // - '[task-return]some-func'
             //
             // At this point in code generation, the following things have already been set:
+            // - `parentTask`: A parent task, if one was executing before
+            // - `subtask`: A subtask, if the current task is a subtask of a parent task
+            // - `task`: the currently executing task
             // - `ret`: the original function return value, via (i.e. via `CallWasm`/`CallInterface`)
+            // - `hostProvided`: whether the original function was a host-provided (i.e. host provided import)
             //
             Instruction::AsyncTaskReturn { name, params } => {
                 let debug_log_fn = self.intrinsic(Intrinsic::DebugLog);
                 uwriteln!(
                     self.src,
-                    "{debug_log_fn}('{prefix} [Instruction::AsyncTaskReturn]', {{
+                    "{debug_log_fn}('{prefix}  [Instruction::AsyncTaskReturn]', {{
                          funcName: '{name}',
                          paramCount: {param_count},
-                         postReturn: {post_return_present}
+                         postReturn: {post_return_present},
+                         hostProvided,
                       }});",
                     param_count = params.len(),
                     post_return_present = self.post_return.is_some(),
@@ -2201,8 +2271,7 @@ impl Bindgen for FunctionBindgen<'_> {
                 let get_or_create_async_state_fn = self.intrinsic(Intrinsic::Component(
                     ComponentIntrinsic::GetOrCreateAsyncState,
                 ));
-                let current_task_get_fn =
-                    self.intrinsic(Intrinsic::AsyncTask(AsyncTaskIntrinsic::GetCurrentTask));
+
                 let component_instance_idx = self.canon_opts.instance.as_u32();
                 let is_async_js = self.requires_async_porcelain | self.is_async;
 
@@ -2212,17 +2281,31 @@ impl Bindgen for FunctionBindgen<'_> {
                 //
                 // e.g. right now a correctly formatted record would trigger the code below,
                 // and it should not.
+                //
+                // NOTE: if the import was host provided we *already* have the result via
+                // JSPI and simply calling the host provided JS function -- there is no need
+                // to drive the async loop as with an async import that came from a component.
+                //
+                //
+                // If a subtask is defined, then we're in the case of a lowered async import,
+                // which means that the first async call (to the callee fn) has occurred,
+                // and a subtask has been created, but has not been triggered as started.
+                //
+                // TODO: can we know that the latest subtask is the right one? Is it possible
+                // to have two subtasks prepped for this task on the same thread?
+                //
                 uwriteln!(
                     self.src,
                     r#"
-                      const componentState = {get_or_create_async_state_fn}({component_instance_idx});
-                      if (!componentState) {{ throw new Error('failed to lookup current component state'); }}
+                      if (hostProvided) {{ return ret; }}
 
-                      const taskMeta = {current_task_get_fn}({component_instance_idx});
-                      if (!taskMeta) {{ throw new Error('failed to find current task metadata'); }}
+                      const currentSubtask = task.getLatestSubtask();
+                      if (currentSubtask) {{
+                          currentSubtask.onStart();
+                      }}
 
-                      const task = taskMeta.task;
-                      if (!task) {{ throw new Error('missing/invalid task in current task metadata'); }}
+                      const componentState = {get_or_create_async_state_fn}({component_instance_idx});
+                      if (!componentState) {{ throw new Error('failed to lookup current component state'); }}
 
                       new Promise(async (resolve, reject) => {{
                           try {{