fix(jco): async return writing for caller memory

This commit fixes a bug where the memory for writing to was being
retrieved too eagerly (@ trampoline fn binding time), and was not
available properly when performing async `task.return`s.

Author: vados-cosmonic

crates/js-component-bindgen/src/intrinsics/p3/async_task.rs

@@ -303,12 +303,14 @@ impl AsyncTaskIntrinsic {
                 let current_task_get_fn = Self::GetCurrentTask.name();
 
                 output.push_str(&format!("
-                    function {task_return_fn}(componentIdx, useDirectParams, memory, callbackFnIdx, liftFns) {{
-                        const params = [...arguments].slice(5);
+                    function {task_return_fn}(ctx) {{
+                        const {{ componentIdx, useDirectParams, getMemoryFn, callbackFnIdx, liftFns }} = ctx;
+                        const params = [...arguments].slice(1);
+                        const memory = getMemoryFn();
                         {debug_log_fn}('[{task_return_fn}()] args', {{
                             componentIdx,
-                            memory,
                             callbackFnIdx,
+                            memory,
                             liftFns,
                             params,
                         }});
@@ -326,18 +328,11 @@ impl AsyncTaskIntrinsic {
                         }}
 
                         let liftCtx = {{ memory, useDirectParams, params }};
-                        if (useDirectParams) {{
-                            liftCtx.memory = memory;
-                            liftCtx.params = params;
-                        }} else {{
+                        if (!useDirectParams) {{
                             liftCtx.storagePtr = params[0];
                             liftCtx.storageLen = params[1];
                         }}
 
-                        if (memory && liftCtx.storageLen === undefined) {{
-                            throw new Error('storageLen must be provided if a memory is present');
-                        }}
-
                         const results = [];
                         {debug_log_fn}('[{task_return_fn}()] lifting results out of memory', {{ liftCtx }});
                         for (const liftFn of liftFns) {{

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

@@ -1996,9 +1996,9 @@ impl<'a> Instantiator<'a, '_> {
                 }
                 let lift_fns_js = format!("[{}]", lift_fns.join(","));
 
-                let memory_js = memory
-                    .map(|idx| format!("memory{}", idx.as_u32()))
-                    .unwrap_or_else(|| "null".into());
+                let get_memory_fn_js = memory
+                    .map(|idx| format!("() => memory{}", idx.as_u32()))
+                    .unwrap_or_else(|| "() => null".into());
                 let component_idx = instance.as_u32();
                 let task_return_fn = self
                     .bindgen
@@ -2011,11 +2011,13 @@ impl<'a> Instantiator<'a, '_> {
                     self.src.js,
                     "const trampoline{i} = {task_return_fn}.bind(
                          null,
-                         {component_idx},
-                         {use_direct_params},
-                         {memory_js},
-                         {callback_fn_idx},
-                         {lift_fns_js},
+                         {{
+                             componentIdx: {component_idx},
+                             useDirectParams: {use_direct_params},
+                             getMemoryFn: {get_memory_fn_js},
+                             callbackFnIdx: {callback_fn_idx},
+                             liftFns: {lift_fns_js},
+                         }},
                      );",
                 );
             }

packages/jco/test/fixtures/components/p3/async-simple-string-return.wasm

@@ -10,24 +10,16 @@ import { P3_COMPONENT_FIXTURES_DIR } from '../common.js';
 suite('Async (WASI P3)', () => {
     // see: https://github.com/bytecodealliance/jco/issues/1076
     test('incorrect task return params offloading', async () => {
-        const componentPath = join(
-            P3_COMPONENT_FIXTURES_DIR,
-            'async-flat-param-adder.wasm'
-        );
-
+        const name = 'async-flat-param-adder';
         const { instance, cleanup } = await setupAsyncTest({
             component: {
-                name: 'async-flat-param-adder',
-                path: componentPath,
+                name,
+                path: join(
+                    P3_COMPONENT_FIXTURES_DIR,
+                    `${name}.wasm`,
+                ),
                 imports: new WASIShim().getImportObject(),
             },
-            jco: {
-                transpile: {
-                    extraArgs: {
-                        minify: false,
-                    }
-                }
-            }
         });
 
         assert(instance.test3);
@@ -38,4 +30,26 @@ suite('Async (WASI P3)', () => {
 
         await cleanup();
     });
+
+    // https://bytecodealliance.zulipchat.com/#narrow/channel/206238-general/topic/Should.20StringLift.20be.20emitted.20for.20async.20return.20values.3F/with/561133720
+    test('simple string return', async () => {
+        const name = 'async-simple-string-return';
+        const { instance, cleanup } = await setupAsyncTest({
+            component: {
+                name,
+                path: join(
+                    P3_COMPONENT_FIXTURES_DIR,
+                    `${name}.wasm`,
+                ),
+                imports: new WASIShim().getImportObject(),
+            },
+        });
+
+        assert.typeOf(instance.asyncGetLiteral, 'function');
+
+        const result = await instance.asyncGetLiteral();
+        assert.strictEqual(result, "literal");
+
+        await cleanup();
+    });
 });