fix: complete logic in getTopLevelAwait function

lib/compiler.test.ts

@@ -1,7 +1,7 @@
 import {
   getCompilerScriptTarget,
   getCompilerSourceMapOptions,
-  getTopLevelAwait,
+  getTopLevelAwaitLocation,
   SourceMapOptions,
 } from "./compiler.ts";
 import { ts } from "./mod.deps.ts";
@@ -35,18 +35,34 @@ Deno.test("script target should have expected outputs", () => {
 Deno.test("get has top level await", () => {
   runTest("const some = code;class SomeOtherCode {}", undefined);
   runTest("async function test() { await 5; }", undefined);
+  runTest(
+    "async function test() { for await (const item of items) {} }",
+    undefined,
+  );
   runTest("await test();", {
     line: 0,
     character: 0,
   });
+  runTest("for await (const item of items) {}", {
+    line: 0,
+    character: 0,
+  });
+  runTest("if (condition) { await test() }", {
+    line: 0,
+    character: 17,
+  });
+  runTest("const t = { prop: await test() };", {
+    line: 0,
+    character: 18,
+  });
 
   function runTest(code: string, expected: ts.LineAndCharacter | undefined) {
     const sourceFile = ts.createSourceFile(
       "file.ts",
       code,
       ts.ScriptTarget.Latest,
     );
-    assertEquals(getTopLevelAwait(sourceFile), expected);
+    assertEquals(getTopLevelAwaitLocation(sourceFile), expected);
   }
 });
 

lib/compiler.ts

@@ -53,18 +53,31 @@ export function getCompilerSourceMapOptions(
   }
 }
 
-export function getTopLevelAwait(sourceFile: ts.SourceFile) {
-  for (const statement of sourceFile.statements) {
+export function getTopLevelAwaitLocation(sourceFile: ts.SourceFile) {
+  const topLevelAwait = getTopLevelAwait(sourceFile);
+  if (topLevelAwait !== undefined) {
+    return sourceFile.getLineAndCharacterOfPosition(
+      topLevelAwait.getStart(sourceFile),
+    );
+  }
+  return undefined;
+}
+
+function getTopLevelAwait(node: ts.Node): ts.Node | undefined {
+  if (ts.isAwaitExpression(node)) {
+    return node;
+  }
+  if (ts.isForOfStatement(node) && node.awaitModifier !== undefined) {
+    return node;
+  }
+  return ts.forEachChild(node, (child) => {
     if (
-      ts.isExpressionStatement(statement) &&
-      ts.isAwaitExpression(statement.expression)
+      !ts.isFunctionDeclaration(child) && !ts.isFunctionExpression(child) &&
+      !ts.isArrowFunction(child) && !ts.isMethodDeclaration(child)
     ) {
-      return sourceFile.getLineAndCharacterOfPosition(
-        statement.expression.getStart(sourceFile),
-      );
+      return getTopLevelAwait(child);
     }
-  }
-  return undefined;
+  });
 }
 
 export function transformCodeToTarget(code: string, target: ts.ScriptTarget) {

mod.ts

@@ -3,7 +3,7 @@
 import {
   getCompilerScriptTarget,
   getCompilerSourceMapOptions,
-  getTopLevelAwait,
+  getTopLevelAwaitLocation,
   outputDiagnostics,
   SourceMapOptions,
   transformCodeToTarget,
@@ -246,7 +246,7 @@ export async function build(options: BuildOptions): Promise<void> {
 
     if (options.scriptModule) {
       // cjs does not support TLA so error fast if we find one
-      const tlaLocation = getTopLevelAwait(sourceFile);
+      const tlaLocation = getTopLevelAwaitLocation(sourceFile);
       if (tlaLocation) {
         warn(
           `Top level await cannot be used when distributing CommonJS/UMD ` +