From af9490c9dd80dddaef01dda18c839e10e408626f Mon Sep 17 00:00:00 2001
From: Leonardo Venturini <leovenbag@gmail.com>
Date: Mon, 27 May 2024 09:25:16 -0400
Subject: [PATCH] implement a more robust solution

---
 lib/import-export-visitor.js      | 11 ++-----
 lib/is-top-level-await.js         | 19 +++++++++++
 test/tla/await-inside-function.js | 54 +++++++++++++++++++++++++++++++
 test/top-level-await-tests.js     |  9 ++++--
 4 files changed, 83 insertions(+), 10 deletions(-)
 create mode 100644 lib/is-top-level-await.js
 create mode 100644 test/tla/await-inside-function.js

diff --git a/lib/import-export-visitor.js b/lib/import-export-visitor.js
index 6622ced..04589bc 100644
--- a/lib/import-export-visitor.js
+++ b/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);
       }
     }
 
diff --git a/lib/is-top-level-await.js b/lib/is-top-level-await.js
new file mode 100644
index 0000000..f86b2f7
--- /dev/null
+++ b/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 }
\ No newline at end of file
diff --git a/test/tla/await-inside-function.js b/test/tla/await-inside-function.js
new file mode 100644
index 0000000..bac2bf4
--- /dev/null
+++ b/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
+};
\ No newline at end of file
diff --git a/test/top-level-await-tests.js b/test/top-level-await-tests.js
index 558fcf3..e566766 100644
--- a/test/top-level-await-tests.js
+++ b/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 {