Wait for SystemJS to reparent module before processing top level awaits

async-to-promises.test.js

@@ -40,6 +40,7 @@ const environments = {
 		},
 		pluginMapping: {
 			"transform-modules-commonjs": "@babel/plugin-transform-modules-commonjs",
+			"transform-modules-systemjs": "@babel/plugin-transform-modules-systemjs",
 			"transform-parameters": "@babel/plugin-transform-parameters",
 			"transform-classes": "@babel/plugin-transform-classes",
 			"external-helpers": "@babel/plugin-external-helpers",

async-to-promises.ts

@@ -290,7 +290,7 @@ type UsedHelpers = { [key in HelperName]: true };
 
 interface PluginState {
 	readonly opts: Partial<Readonly<AsyncToPromisesConfiguration>>;
-	processedTopLevelAwait?: true;
+	hasTopLevelAwait?: boolean;
 	usedHelpers?: UsedHelpers;
 }
 
@@ -349,6 +349,10 @@ const numberNames = ["zero", "one", "two", "three", "four", "five", "six", "seve
 
 type CompatibleSubset<New, Old> = Partial<New> & Pick<New, keyof New & keyof Old>;
 
+interface FindAwaitExpressionState {
+	awaitPath?: NodePath;
+}
+
 // function generate(node?: Node): string {
 // 	if (node === undefined) {
 // 		return "";
@@ -4791,38 +4795,59 @@ export default function ({
 		},
 	};
 
+	const findAwaitExpressionVisitor: Visitor<FindAwaitExpressionState> = {
+		AwaitExpression(path) {
+			this.awaitPath = path;
+			path.stop();
+		},
+	};
+
 	// Main babel plugin implementation and top level visitor
 	return {
 		manipulateOptions(_options: any, parserOptions: { plugins: string[] }) {
 			parserOptions.plugins.push("asyncGenerators");
 		},
 		visitor: {
 			AwaitExpression(path) {
-				if (!path.getFunctionParent() && !this.processedTopLevelAwait) {
-					this.processedTopLevelAwait = true;
-					switch (readConfigKey(this.opts, "topLevelAwait")) {
-						case "simple": {
-							const programPath = path.scope.getProgramParent().path;
-							rewriteAsyncBlock({ state: this }, programPath, [], undefined, false, true);
-							break;
-						}
-						case "return": {
-							const programPath = path.scope.getProgramParent().path;
-							rewriteAsyncBlock({ state: this }, programPath, [], undefined, false, false);
-							break;
-						}
-						case "ignore":
-							break;
-						default:
-							throw path.buildCodeFrameError(
-								`Top level await is not supported unless experimental topLevelAwait: "simple" or topLevelAwait: "return" options are specified!`,
-								TypeError
-							);
-					}
+				if (!path.getFunctionParent() && !this.hasTopLevelAwait) {
+					this.hasTopLevelAwait = true;
 				}
 			},
 			Program: {
 				exit(path) {
+					if (this.hasTopLevelAwait) {
+						// rediscover the top level await path since it may have been reorganized by a module plugin
+						let rediscoverState = {} as FindAwaitExpressionState;
+						path.traverse(findAwaitExpressionVisitor, rediscoverState);
+						if (rediscoverState.awaitPath !== undefined) {
+							const functionParent = rediscoverState.awaitPath.getFunctionParent();
+							const topLevelAwaitParent = functionParent ? functionParent.get("body") : path;
+							switch (readConfigKey(this.opts, "topLevelAwait")) {
+								case "simple": {
+									rewriteAsyncBlock({ state: this }, topLevelAwaitParent, [], undefined, false, true);
+									break;
+								}
+								case "return": {
+									rewriteAsyncBlock(
+										{ state: this },
+										topLevelAwaitParent,
+										[],
+										undefined,
+										false,
+										false
+									);
+									break;
+								}
+								case "ignore":
+									break;
+								default:
+									throw rediscoverState.awaitPath.buildCodeFrameError(
+										`Top level await is not supported unless experimental topLevelAwait: "simple" or topLevelAwait: "return" options are specified!`,
+										TypeError
+									);
+							}
+						}
+					}
 					const usedHelpers = this.usedHelpers;
 					if (usedHelpers !== undefined) {
 						const file = getFile(path);

package.json

@@ -34,6 +34,7 @@
     "@babel/plugin-proposal-class-properties": "^7.16.5",
     "@babel/plugin-transform-classes": "^7.16.5",
     "@babel/plugin-transform-modules-commonjs": "^7.16.5",
+    "@babel/plugin-transform-modules-systemjs": "^7.16.5",
     "@babel/plugin-transform-parameters": "^7.16.5",
     "@babel/preset-env": "^7.16.5",
     "@babel/traverse": "^7.16.5",

tests/top level await export/input.js

@@ -0,0 +1,5 @@
+export const head = 1;
+const response = await fetch("https://www.example.com/");
+export const body = 2;
+const json = await response.json();
+export const tail = json;

tests/top level await export/options.json

@@ -0,0 +1,9 @@
+{
+	"supportedBabels": ["babel 7"],
+	"checkSyntax": false,
+	"module": true,
+	"plugins": ["transform-modules-systemjs"],
+	"options": {
+		"topLevelAwait": "return"
+	}
+}