TypeScriptToLua
diff --git a/‎src/transformation/utils/typescript/index.ts‎
Lines changed: 16 additions & 0 deletions b/‎src/transformation/utils/typescript/index.ts‎
Lines changed: 16 additions & 0 deletions
diff --git a/‎src/transformation/visitors/function.ts‎
Lines changed: 4 additions & 1 deletion b/‎src/transformation/visitors/function.ts‎
Lines changed: 4 additions & 1 deletion
diff --git a/‎src/transformation/visitors/sourceFile.ts‎
Lines changed: 9 additions & 8 deletions b/‎src/transformation/visitors/sourceFile.ts‎
Lines changed: 9 additions & 8 deletions
diff --git a/‎test/unit/builtins/async-await.spec.ts‎
Lines changed: 14 additions & 8 deletions b/‎test/unit/builtins/async-await.spec.ts‎
Lines changed: 14 additions & 8 deletions
@@ -92,3 +92,19 @@ export function getFunctionTypeForCall(context: TransformationContext, node: ts.
     }
     return context.checker.getTypeFromTypeNode(typeDeclaration.type);
 }
+
+// https://github.com/microsoft/TypeScript/blob/663b19fe4a7c4d4ddaa61aedadd28da06acd27b6/src/services/documentHighlights.ts#L435
+// Do not cross function/class/interface/module/type boundaries.
+export function traverseWithoutCrossingFunction(node: ts.Node, cb: (node: ts.Node) => void) {
+    cb(node);
+    if (
+        !ts.isFunctionLike(node) &&
+        !ts.isClassLike(node) &&
+        !ts.isInterfaceDeclaration(node) &&
+        !ts.isModuleDeclaration(node) &&
+        !ts.isTypeAliasDeclaration(node) &&
+        !ts.isTypeNode(node)
+    ) {
+        ts.forEachChild(node, child => traverseWithoutCrossingFunction(child, cb));
+    }
+}
@@ -196,8 +196,11 @@ export function transformFunctionToExpression(
         spreadIdentifier,
         node
     );
+
+    const possiblyAsyncBody = isAsyncFunction(node) ? wrapInAsyncAwaiter(context, transformedBody) : transformedBody;
+
     const functionExpression = lua.createFunctionExpression(
-        lua.createBlock(isAsyncFunction(node) ? wrapInAsyncAwaiter(context, transformedBody) : transformedBody),
+        lua.createBlock(possiblyAsyncBody),
         paramNames,
         dotsLiteral,
         flags,
 
@@ -6,7 +6,7 @@ import { notAllowedTopLevelAwait } from "../utils/diagnostics";
 import { createExportsIdentifier } from "../utils/lua-ast";
 import { getUsedLuaLibFeatures } from "../utils/lualib";
 import { performHoisting, popScope, pushScope, ScopeType } from "../utils/scope";
-import { hasExportEquals } from "../utils/typescript";
+import { hasExportEquals, traverseWithoutCrossingFunction } from "../utils/typescript";
 
 export const transformSourceFileNode: FunctionVisitor<ts.SourceFile> = (node, context) => {
     let statements: lua.Statement[] = [];
@@ -52,11 +52,12 @@ export const transformSourceFileNode: FunctionVisitor<ts.SourceFile> = (node, co
 };
 
 function isTopLevelAwait(statement: ts.Statement) {
-    return (
-        (ts.isExpressionStatement(statement) && ts.isAwaitExpression(statement.expression)) ||
-        (ts.isVariableStatement(statement) &&
-            statement.declarationList.declarations.some(
-                declaration => declaration.initializer && ts.isAwaitExpression(declaration.initializer)
-            ))
-    );
+    // Check if expression contains an await child, without going into function declarations
+    let containsAwait = false;
+    traverseWithoutCrossingFunction(statement, node => {
+        if (ts.isAwaitExpression(node)) {
+            containsAwait = true;
+        }
+    });
+    return containsAwait;
 }
@@ -225,18 +225,24 @@ test("can call async function at top-level", () => {
         });
 });
 
-test.each(["await a();", "const b = await a();", "export const b = await a();"])(
-    "cannot await at top-level (%p)",
-    awaitUsage => {
-        util.testModule`
+test.each([
+    "await a();",
+    "const b = await a();",
+    "export const b = await a();",
+    "declare function foo(n: number): number; foo(await a());",
+    "declare function foo(n: number): number; const b = foo(await a());",
+    "const b = [await a()];",
+    "const b = [4, await a()];",
+    "const b = true ? 4 : await a();",
+])("cannot await at top-level (%p)", awaitUsage => {
+    util.testModule`
         async function a() {
             return 42;
         }
 
         ${awaitUsage}
         export {} // Required to make TS happy, cannot await without import/exports
     `
-            .setOptions({ module: ModuleKind.ESNext, target: ScriptTarget.ES2017 })
-            .expectToHaveDiagnostics([notAllowedTopLevelAwait.code]);
-    }
-);
+        .setOptions({ module: ModuleKind.ESNext, target: ScriptTarget.ES2017 })
+        .expectToHaveDiagnostics([notAllowedTopLevelAwait.code]);
+});