feat: better async transform (yield replace await)

src/lib/task.ts

@@ -116,7 +116,7 @@ function _task<TArgs = undefined, TReturn = unknown>(
 									if (abort_controller.signal.aborted) {
 										break;
 									}
-									next_val = await gen_or_value.next();
+									next_val = await gen_or_value.next(next_val.value);
 								}
 								if (next_val.done) {
 									const last_result = next_val.value;

src/lib/tests/expected-transforms/add-yield-after-await-assignment/transform.js

@@ -7,87 +7,31 @@ function fn(val) {
 function str(quasi, strings) {}
 
 task(async function* () {
-	const assign = await Promise.resolve();
-
-	yield;
-
-	const array = [await Promise.resolve()];
-
-	yield;
-
-	const sum = await Promise.resolve(1) + 2;
-
-	yield;
-
-	const sum2 = 2 + await Promise.resolve(1);
-
-	yield;
-
-	const function_call = fn(await Promise.resolve(2));
-
-	yield;
-
-	const conditional1 = await Promise.resolve(true) ? 1 : 2;
-
-	yield;
-
-	const conditional2 = true ? await Promise.resolve(1) : 2;
-
-	yield;
-
-	const conditional3 = false ? 1 : await Promise.resolve(2);
-
-	yield;
-
-	const logical1 = await Promise.resolve(null) ?? 3;
-
-	yield;
-
-	const logical2 = null ?? await Promise.resolve(null);
-
-	yield;
-
-	const logical3 = null || await Promise.resolve(null);
-
-	yield;
-
-	const logical4 = await Promise.resolve(null) || null;
-
-	yield;
-
-	const logical5 = await Promise.resolve(null) && null;
-
-	yield;
-
-	const logical6 = true && await Promise.resolve(null);
-
-	yield;
-
-	const object_expression1 = { awaited: await Promise.resolve(2) };
-
-	yield;
+	const assign = yield Promise.resolve();
+	const array = [yield Promise.resolve()];
+	const sum = (yield Promise.resolve(1)) + 2;
+	const sum2 = 2 + (yield Promise.resolve(1));
+	const function_call = fn(yield Promise.resolve(2));
+	const conditional1 = (yield Promise.resolve(true)) ? 1 : 2;
+	const conditional2 = true ? yield Promise.resolve(1) : 2;
+	const conditional3 = false ? 1 : yield Promise.resolve(2);
+	const logical1 = (yield Promise.resolve(null)) ?? 3;
+	const logical2 = null ?? (yield Promise.resolve(null));
+	const logical3 = null || (yield Promise.resolve(null));
+	const logical4 = (yield Promise.resolve(null)) || null;
+	const logical5 = (yield Promise.resolve(null)) && null;
+	const logical6 = true && (yield Promise.resolve(null));
+	const object_expression1 = { awaited: yield Promise.resolve(2) };
 
 	const object_expression2 = {
-		awaited: { nested: await Promise.resolve(2) }
+		awaited: { nested: yield Promise.resolve(2) }
 	};
 
-	yield;
-
 	const object_expression3 = {
-		awaited: { nested: [await Promise.resolve(2)] }
+		awaited: { nested: [yield Promise.resolve(2)] }
 	};
 
-	yield;
-
-	const tagged_template = str`something ${await Promise.resolve("")}`;
-
-	yield;
-
-	const template_literal = `something ${await Promise.resolve("")}`;
-
-	yield;
-
-	const unary = !await Promise.resolve(true);
-
-	yield;
+	const tagged_template = str`something ${yield Promise.resolve("")}`;
+	const template_literal = `something ${yield Promise.resolve("")}`;
+	const unary = !(yield Promise.resolve(true));
 });

src/lib/tests/expected-transforms/add-yield-after-await/transform.js

@@ -1,6 +1,5 @@
 import { task } from "svelte-concurrency";
 
 task(async function* () {
-	await Promise.resolve();
-	yield;
+	yield Promise.resolve();
 });

src/lib/tests/expected-transforms/add-yield-after-multiple-await/transform.js

@@ -1,8 +1,6 @@
 import { task } from "svelte-concurrency";
 
 task(async function* () {
-	await Promise.resolve();
-	yield;
-	await Promise.resolve();
-	yield;
+	yield Promise.resolve();
+	yield Promise.resolve();
 });

src/lib/tests/expected-transforms/task-aliased-plus-task-from-another-library/transform.js

@@ -7,8 +7,6 @@ task(async () => {
 });
 
 other_name(async function* () {
-	await Promise.resolve();
-	yield;
-	await Promise.resolve();
-	yield;
+	yield Promise.resolve();
+	yield Promise.resolve();
 });

src/lib/tests/expected-transforms/task-aliased-plus-task-function/transform.js

@@ -10,8 +10,6 @@ task(async () => {
 });
 
 other_name(async function* () {
-	await Promise.resolve();
-	yield;
-	await Promise.resolve();
-	yield;
+	yield Promise.resolve();
+	yield Promise.resolve();
 });

src/lib/tests/expected-transforms/task-aliased/transform.js

@@ -1,8 +1,6 @@
 import { task as other_name } from "svelte-concurrency";
 
 other_name(async function* () {
-	await Promise.resolve();
-	yield;
-	await Promise.resolve();
-	yield;
+	yield Promise.resolve();
+	yield Promise.resolve();
 });

src/lib/vite.ts

@@ -9,71 +9,70 @@ import type {
 	CallExpression,
 	Statement,
 	Expression,
+	YieldExpression,
 } from 'acorn';
 
 type Nodes = ImportDeclaration | FunctionExpression | ArrowFunctionExpression | CallExpression;
 
 /**
  * This should handle most situation where you can top level await in an async function
  */
-function has_expression_await(expression: Expression): boolean {
+function get_expressions_await(expression: Expression): Expression[] {
 	switch (expression.type) {
 		case 'ArrayExpression':
 			// const x = [await promise];
-			return expression.elements.some(
-				(element) =>
-					element?.type !== 'SpreadElement' && element !== null && has_expression_await(element),
+			return expression.elements.flatMap((element) =>
+				element?.type !== 'SpreadElement' && element !== null ? get_expressions_await(element) : [],
 			);
 		case 'AssignmentExpression':
 			// x = await promise; x = { value: await promise }; x = [await promise]
 			// TODO: [x = await promise] = [undefined]
-			return has_expression_await(expression.right);
+			return get_expressions_await(expression.right);
 		case 'AwaitExpression':
 			// await promise;
-			return true;
+			return [expression];
 		case 'BinaryExpression':
 			// await promise + something;
-			return (
-				(expression.left.type !== 'PrivateIdentifier' && has_expression_await(expression.left)) ||
-				has_expression_await(expression.right)
-			);
+			if (expression.left.type !== 'PrivateIdentifier') {
+				return get_expressions_await(expression.left).concat(
+					get_expressions_await(expression.right),
+				);
+			}
+			return get_expressions_await(expression.right);
 		case 'CallExpression':
 			// fn_call(await stuff);
 			// TODO: fns[await name]();
-			return expression.arguments.some(
-				(argument) => argument.type !== 'SpreadElement' && has_expression_await(argument),
+			return expression.arguments.flatMap((argument) =>
+				argument.type !== 'SpreadElement' ? get_expressions_await(argument) : [],
 			);
 		case 'ConditionalExpression':
 			// await test ? await consequent : await alternate;
-			return (
-				has_expression_await(expression.alternate) ||
-				has_expression_await(expression.consequent) ||
-				has_expression_await(expression.test)
-			);
+			return get_expressions_await(expression.alternate)
+				.concat(get_expressions_await(expression.consequent))
+				.concat(get_expressions_await(expression.test));
 		case 'LogicalExpression':
 			// await promise || await another;
-			return has_expression_await(expression.left) || has_expression_await(expression.right);
+			return get_expressions_await(expression.left).concat(get_expressions_await(expression.right));
 		case 'MemberExpression':
-			return (
-				expression.property.type !== 'PrivateIdentifier' &&
-				has_expression_await(expression.property)
-			);
+			return expression.property.type !== 'PrivateIdentifier'
+				? get_expressions_await(expression.property)
+				: [];
 		case 'ObjectExpression':
-			return expression.properties.some(
-				(property) =>
-					property.type !== 'SpreadElement' &&
-					(has_expression_await(property.key) || has_expression_await(property.value)),
+			return expression.properties.flatMap((property) =>
+				property.type !== 'SpreadElement'
+					? get_expressions_await(property.key).concat(get_expressions_await(property.value))
+					: [],
 			);
 		case 'TaggedTemplateExpression':
-			return has_expression_await(expression.quasi);
+			return get_expressions_await(expression.quasi);
 		case 'TemplateLiteral':
-			return expression.expressions.some((template_expression) =>
-				has_expression_await(template_expression),
+			return expression.expressions.flatMap((template_expression) =>
+				get_expressions_await(template_expression),
 			);
 		case 'UnaryExpression':
-			return has_expression_await(expression.argument);
+			return get_expressions_await(expression.argument);
 		default:
-			return false;
+			return [];
 	}
 }
 
@@ -82,26 +81,22 @@ function update_body(task: FunctionExpression) {
 
 	for (const statement of task.body.body) {
 		body.push(statement);
-		if (
-			(statement.type === 'ExpressionStatement' && has_expression_await(statement.expression)) ||
-			(statement.type === 'VariableDeclaration' &&
-				statement.declarations.some((declarator) => {
-					return declarator.init && has_expression_await(declarator.init);
-				})) ||
-			(statement.type === 'ForInStatement' && has_expression_await(statement.right)) ||
-			(statement.type === 'ForOfStatement' && has_expression_await(statement.right))
-		) {
-			body.push({
-				type: 'ExpressionStatement',
-				expression: {
-					type: 'YieldExpression',
-					delegate: false,
-					start: 0,
-					end: 0,
-				},
-				start: 0,
-				end: 0,
-			});
+		let expressions: Expression[] = [];
+		if (statement.type === 'ExpressionStatement') {
+			expressions = get_expressions_await(statement.expression);
+		} else if (statement.type === 'VariableDeclaration') {
+			for (const declaration of statement.declarations) {
+				if (declaration.init) {
+					expressions = expressions.concat(get_expressions_await(declaration.init));
+				}
+			}
+		} else if (statement.type === 'ForInStatement' || statement.type === 'ForOfStatement') {
+			expressions = get_expressions_await(statement.right);
+		}
+		for (const expression of expressions) {
+			if (expression.type === 'AwaitExpression') {
+				(expression as unknown as YieldExpression).type = 'YieldExpression';
+			}
 		}
 	}
 	task.body.body = body;