From 9dc78d005c7965ce8f5f9e7f4c143a504262bc34 Mon Sep 17 00:00:00 2001
From: Yosuke Ota <otameshiyo23@gmail.com>
Date: Sat, 17 Jul 2021 16:03:14 +0900
Subject: [PATCH] Fix false positives for `<script setup>` in
 `vue/no-lifecycle-after-await`, `vue/no-watch-after-await`,
 `vue/no-restricted-call-after-await` rules. (#1569)

---
 lib/rules/no-lifecycle-after-await.js         | 47 +++++--------------
 lib/rules/no-restricted-call-after-await.js   | 42 ++++-------------
 lib/rules/no-watch-after-await.js             | 45 +++++-------------
 tests/lib/rules/no-lifecycle-after-await.js   | 30 +++++-------
 .../rules/no-restricted-call-after-await.js   | 31 +++++-------
 tests/lib/rules/no-watch-after-await.js       | 34 ++++++--------
 6 files changed, 70 insertions(+), 159 deletions(-)

diff --git a/lib/rules/no-lifecycle-after-await.js b/lib/rules/no-lifecycle-after-await.js
index aba61794d..d38a97a59 100644
--- a/lib/rules/no-lifecycle-after-await.js
+++ b/lib/rules/no-lifecycle-after-await.js
@@ -48,11 +48,11 @@ module.exports = {
     /**
      * @typedef {object} ScopeStack
      * @property {ScopeStack | null} upper
-     * @property {FunctionDeclaration | FunctionExpression | ArrowFunctionExpression | Program} scopeNode
+     * @property {FunctionDeclaration | FunctionExpression | ArrowFunctionExpression} scopeNode
      */
     /** @type {Set<ESNode>} */
     const lifecycleHookCallNodes = new Set()
-    /** @type {Map<FunctionDeclaration | FunctionExpression | ArrowFunctionExpression | Program, SetupScopeData>} */
+    /** @type {Map<FunctionDeclaration | FunctionExpression | ArrowFunctionExpression, SetupScopeData>} */
     const setupScopes = new Map()
 
     /** @type {ScopeStack | null} */
@@ -60,15 +60,7 @@ module.exports = {
 
     return utils.compositingVisitors(
       {
-        /**
-         * @param {Program} node
-         */
-        Program(node) {
-          scopeStack = {
-            upper: scopeStack,
-            scopeNode: node
-          }
-
+        Program() {
           const tracker = new ReferenceTracker(context.getScope())
           const traceMap = {
             /** @type {TraceMap} */
@@ -85,6 +77,14 @@ module.exports = {
           for (const { node } of tracker.iterateEsmReferences(traceMap)) {
             lifecycleHookCallNodes.add(node)
           }
+        }
+      },
+      utils.defineVueVisitor(context, {
+        onSetupFunctionEnter(node) {
+          setupScopes.set(node, {
+            afterAwait: false,
+            range: node.range
+          })
         },
         /**
          * @param {FunctionExpression | FunctionDeclaration | ArrowFunctionExpression} node
@@ -133,31 +133,6 @@ module.exports = {
               messageId: 'forbidden'
             })
           }
-        }
-      },
-      (() => {
-        const scriptSetup = utils.getScriptSetupElement(context)
-        if (!scriptSetup) {
-          return {}
-        }
-        return {
-          /**
-           * @param {Program} node
-           */
-          Program(node) {
-            setupScopes.set(node, {
-              afterAwait: false,
-              range: scriptSetup.range
-            })
-          }
-        }
-      })(),
-      utils.defineVueVisitor(context, {
-        onSetupFunctionEnter(node) {
-          setupScopes.set(node, {
-            afterAwait: false,
-            range: node.range
-          })
         },
         onSetupFunctionExit(node) {
           setupScopes.delete(node)
diff --git a/lib/rules/no-restricted-call-after-await.js b/lib/rules/no-restricted-call-after-await.js
index 0c1989418..523d6af12 100644
--- a/lib/rules/no-restricted-call-after-await.js
+++ b/lib/rules/no-restricted-call-after-await.js
@@ -62,13 +62,13 @@ module.exports = {
 
     /** @type {Map<ESNode, string>} */
     const restrictedCallNodes = new Map()
-    /** @type {Map<FunctionExpression | ArrowFunctionExpression | FunctionDeclaration | Program, SetupScopeData>} */
+    /** @type {Map<FunctionExpression | ArrowFunctionExpression | FunctionDeclaration, SetupScopeData>} */
     const setupScopes = new Map()
 
     /**x
      * @typedef {object} ScopeStack
      * @property {ScopeStack | null} upper
-     * @property {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration | Program} scopeNode
+     * @property {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration} scopeNode
      */
     /** @type {ScopeStack | null} */
     let scopeStack = null
@@ -142,11 +142,6 @@ module.exports = {
       {
         /** @param {Program} node */
         Program(node) {
-          scopeStack = {
-            upper: scopeStack,
-            scopeNode: node
-          }
-
           const tracker = new ReferenceTracker(context.getScope())
 
           for (const option of context.options) {
@@ -181,6 +176,14 @@ module.exports = {
               }
             }
           }
+        }
+      },
+      utils.defineVueVisitor(context, {
+        onSetupFunctionEnter(node) {
+          setupScopes.set(node, {
+            afterAwait: false,
+            range: node.range
+          })
         },
         /** @param {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration} node */
         ':function'(node) {
@@ -225,31 +228,6 @@ module.exports = {
               data: { message }
             })
           }
-        }
-      },
-      (() => {
-        const scriptSetup = utils.getScriptSetupElement(context)
-        if (!scriptSetup) {
-          return {}
-        }
-        return {
-          /**
-           * @param {Program} node
-           */
-          Program(node) {
-            setupScopes.set(node, {
-              afterAwait: false,
-              range: scriptSetup.range
-            })
-          }
-        }
-      })(),
-      utils.defineVueVisitor(context, {
-        onSetupFunctionEnter(node) {
-          setupScopes.set(node, {
-            afterAwait: false,
-            range: node.range
-          })
         },
         onSetupFunctionExit(node) {
           setupScopes.delete(node)
diff --git a/lib/rules/no-watch-after-await.js b/lib/rules/no-watch-after-await.js
index 8de30668e..6ecc4dd75 100644
--- a/lib/rules/no-watch-after-await.js
+++ b/lib/rules/no-watch-after-await.js
@@ -62,26 +62,20 @@ module.exports = {
      * @property {boolean} afterAwait
      * @property {[number,number]} range
      */
-    /** @type {Map<FunctionExpression | ArrowFunctionExpression | FunctionDeclaration | Program, SetupScopeData>} */
+    /** @type {Map<FunctionExpression | ArrowFunctionExpression | FunctionDeclaration, SetupScopeData>} */
     const setupScopes = new Map()
 
     /**
      * @typedef {object} ScopeStack
      * @property {ScopeStack | null} upper
-     * @property {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration | Program} scopeNode
+     * @property {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration} scopeNode
      */
     /** @type {ScopeStack | null} */
     let scopeStack = null
 
     return utils.compositingVisitors(
       {
-        /** @param {Program} node */
-        Program(node) {
-          scopeStack = {
-            upper: scopeStack,
-            scopeNode: node
-          }
-
+        Program() {
           const tracker = new ReferenceTracker(context.getScope())
           const traceMap = {
             vue: {
@@ -98,6 +92,14 @@ module.exports = {
           for (const { node } of tracker.iterateEsmReferences(traceMap)) {
             watchCallNodes.add(node)
           }
+        }
+      },
+      utils.defineVueVisitor(context, {
+        onSetupFunctionEnter(node) {
+          setupScopes.set(node, {
+            afterAwait: false,
+            range: node.range
+          })
         },
         /** @param {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration} node */
         ':function'(node) {
@@ -140,31 +142,6 @@ module.exports = {
               messageId: 'forbidden'
             })
           }
-        }
-      },
-      (() => {
-        const scriptSetup = utils.getScriptSetupElement(context)
-        if (!scriptSetup) {
-          return {}
-        }
-        return {
-          /**
-           * @param {Program} node
-           */
-          Program(node) {
-            setupScopes.set(node, {
-              afterAwait: false,
-              range: scriptSetup.range
-            })
-          }
-        }
-      })(),
-      utils.defineVueVisitor(context, {
-        onSetupFunctionEnter(node) {
-          setupScopes.set(node, {
-            afterAwait: false,
-            range: node.range
-          })
         },
         onSetupFunctionExit(node) {
           setupScopes.delete(node)
diff --git a/tests/lib/rules/no-lifecycle-after-await.js b/tests/lib/rules/no-lifecycle-after-await.js
index 9fc17f3ab..bb7414682 100644
--- a/tests/lib/rules/no-lifecycle-after-await.js
+++ b/tests/lib/rules/no-lifecycle-after-await.js
@@ -139,6 +139,18 @@ tester.run('no-lifecycle-after-await', rule, {
       <script>
       import {onMounted} from 'vue'
       await doSomething()
+      onMounted(() => { /* ... */ }) // not error
+      </script>
+      `,
+      parserOptions: { ecmaVersion: 2022 }
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      import {onMounted} from 'vue'
+      await doSomething()
+
       onMounted(() => { /* ... */ }) // not error
       </script>
       `,
@@ -261,24 +273,6 @@ tester.run('no-lifecycle-after-await', rule, {
           messageId: 'forbidden'
         }
       ]
-    },
-    {
-      filename: 'test.vue',
-      code: `
-      <script setup>
-      import {onMounted} from 'vue'
-      await doSomething()
-
-      onMounted(() => { /* ... */ }) // error
-      </script>
-      `,
-      parserOptions: { ecmaVersion: 2022 },
-      errors: [
-        {
-          messageId: 'forbidden',
-          line: 6
-        }
-      ]
     }
   ]
 })
diff --git a/tests/lib/rules/no-restricted-call-after-await.js b/tests/lib/rules/no-restricted-call-after-await.js
index 92353eaf0..0eef6713b 100644
--- a/tests/lib/rules/no-restricted-call-after-await.js
+++ b/tests/lib/rules/no-restricted-call-after-await.js
@@ -137,6 +137,18 @@ tester.run('no-restricted-call-after-await', rule, {
       `,
       options: [{ module: 'vue-i18n', path: 'useI18n' }],
       parserOptions: { ecmaVersion: 2022 }
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      import {useI18n} from 'vue-i18n'
+      await doSomething()
+      useI18n()
+      </script>
+      `,
+      parserOptions: { ecmaVersion: 2022 },
+      options: [{ module: 'vue-i18n', path: 'useI18n' }]
     }
   ],
   invalid: [
@@ -389,25 +401,6 @@ tester.run('no-restricted-call-after-await', rule, {
           line: 7
         }
       ]
-    },
-    {
-      filename: 'test.vue',
-      code: `
-      <script setup>
-      import {useI18n} from 'vue-i18n'
-      await doSomething()
-      useI18n()
-      </script>
-      `,
-      parserOptions: { ecmaVersion: 2022 },
-      options: [{ module: 'vue-i18n', path: 'useI18n' }],
-      errors: [
-        {
-          message:
-            'The `import("vue-i18n").useI18n` after `await` expression are forbidden.',
-          line: 5
-        }
-      ]
     }
   ]
 })
diff --git a/tests/lib/rules/no-watch-after-await.js b/tests/lib/rules/no-watch-after-await.js
index ae5718672..e97c78c63 100644
--- a/tests/lib/rules/no-watch-after-await.js
+++ b/tests/lib/rules/no-watch-after-await.js
@@ -155,6 +155,20 @@ tester.run('no-watch-after-await', rule, {
       </script>
       `,
       parserOptions: { ecmaVersion: 2022 }
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      import {watch} from 'vue'
+      watch(foo, () => { /* ... */ })
+
+      await doSomething()
+
+      watch(foo, () => { /* ... */ })
+      </script>
+      `,
+      parserOptions: { ecmaVersion: 2022 }
     }
   ],
   invalid: [
@@ -236,26 +250,6 @@ tester.run('no-watch-after-await', rule, {
           line: 12
         }
       ]
-    },
-    {
-      filename: 'test.vue',
-      code: `
-      <script setup>
-      import {watch} from 'vue'
-      watch(foo, () => { /* ... */ })
-
-      await doSomething()
-
-      watch(foo, () => { /* ... */ })
-      </script>
-      `,
-      parserOptions: { ecmaVersion: 2022 },
-      errors: [
-        {
-          message: 'The `watch` after `await` expression are forbidden.',
-          line: 8
-        }
-      ]
     }
   ]
 })