implement a more robust solution

lib/import-export-visitor.js

@@ -6,6 +6,7 @@ const utils = require("./utils.js");
 
 const MagicString = require("magic-string");
 const Visitor = require("./visitor.js");
+const { isTopLevelAwait } = require('./is-top-level-await');
 
 const codeOfCR = "\r".charCodeAt(0);
 const codeOfDoubleQuote = '"'.charCodeAt(0);
@@ -390,19 +391,13 @@ class ImportExportVisitor extends Visitor {
     if (!this.hasTopLevelAwait) {
       const parent = path.getParentNode(1);
 
-      if (
-        parent.type === 'VariableDeclarator' &&
-        (path.getParentNode(3).type === 'Program' ||
-        path.getParentNode(4).type === 'Program')
-      ) {
-        this.hasTopLevelAwait = true;
-      }
-
       if (
         parent.type === 'ExpressionStatement' &&
         path.getParentNode(2).type === 'Program'
       ) {
         this.hasTopLevelAwait = true;
+      } else {
+        this.hasTopLevelAwait = isTopLevelAwait(path.stack);
       }
     }
 

lib/is-top-level-await.js

@@ -0,0 +1,19 @@
+const FUNCTION_TYPES = [
+  'FunctionDeclaration',
+  'FunctionExpression',
+  'ArrowFunctionExpression',
+  'ClassMethod',
+  'ObjectMethod',
+  'ClassPrivateMethod',
+]
+
+function isTopLevelAwait(ancestors) {
+  for (let ancestor of ancestors) {
+    if (FUNCTION_TYPES.includes(ancestor.type)) {
+      return false;
+    }
+  }
+  return true;
+}
+
+module.exports = { isTopLevelAwait }

test/tla/await-inside-function.js

@@ -0,0 +1,54 @@
+// FunctionDeclaration
+async function functionDeclaration() {
+  const result = await Promise.resolve("This is a function declaration");
+  return result;
+}
+
+// FunctionExpression
+var functionExpression = async function() {
+  const result = await Promise.resolve("This is a function expression");
+  return result;
+};
+
+// ArrowFunctionExpression
+var arrowFunctionExpression = async () => {
+  const result = await Promise.resolve("This is an arrow function expression");
+  return result;
+};
+
+// ClassMethod
+class MyClass {
+  async classMethod() {
+    const result = await Promise.resolve("This is a class method");
+    return result;
+  }
+}
+
+// ObjectMethod
+var myObject = {
+  objectMethod: async function() {
+    const result = await Promise.resolve("This is an object method");
+    return result;
+  }
+};
+
+// ClassPrivateMethod
+class MyClassWithPrivateMethod {
+  async #privateMethod() {
+    const result = await Promise.resolve("This is a private class method");
+    return result;
+  }
+
+  async callPrivateMethod() {
+    return this.#privateMethod();
+  }
+}
+
+module.exports = {
+  functionDeclaration,
+  functionExpression,
+  arrowFunctionExpression,
+  MyClass,
+  myObject,
+  MyClassWithPrivateMethod
+};

test/top-level-await-tests.js

@@ -85,16 +85,21 @@ import { importSync, importAsync, importAsyncEvaluated } from './tla/nested/pare
     assert(exports.a === 1);
   });
 
-  it('should detect exported declarations', async () => {
+  it('should detect tla on exported declarations', async () => {
     const exports = await require('./tla/exported-declaration.js');
     assert(exports.test === 'value');
   })
 
-  it('should detect non exported declarations', async () => {
+  it('should detect tla on non exported declarations', async () => {
     const exports = await require('./tla/non-exported-declaration.js');
     assert(exports.redirected === 'value');
   })
 
+  it('should not detect tla if they are inside any function type', async () => {
+    const exports = require('./tla/await-inside-function.js');
+    assert(!(exports instanceof Promise))
+  })
+
   describe('errors', () => {
     it('should synchronously throw error for sync module', () => {
       try {