From 747ab5b90b8476684162e22b6859ca2e41e155b7 Mon Sep 17 00:00:00 2001
From: Yosuke Ota <otameshiyo23@gmail.com>
Date: Tue, 6 Jul 2021 00:01:54 +0900
Subject: [PATCH] Update `vue/no-watch-after-await` rule to support `<script
 setup>` (#1541)

---
 lib/rules/no-watch-after-await.js       | 83 +++++++++++++++++--------
 tests/lib/rules/no-watch-after-await.js | 57 +++++++++++++++++
 2 files changed, 115 insertions(+), 25 deletions(-)

diff --git a/lib/rules/no-watch-after-await.js b/lib/rules/no-watch-after-await.js
index 5f037fb30..8de30668e 100644
--- a/lib/rules/no-watch-after-await.js
+++ b/lib/rules/no-watch-after-await.js
@@ -57,20 +57,31 @@ module.exports = {
   /** @param {RuleContext} context */
   create(context) {
     const watchCallNodes = new Set()
-    /** @type {Map<FunctionExpression | ArrowFunctionExpression | FunctionDeclaration, { setupProperty: Property, afterAwait: boolean }>} */
-    const setupFunctions = new Map()
+    /**
+     * @typedef {object} SetupScopeData
+     * @property {boolean} afterAwait
+     * @property {[number,number]} range
+     */
+    /** @type {Map<FunctionExpression | ArrowFunctionExpression | FunctionDeclaration | Program, SetupScopeData>} */
+    const setupScopes = new Map()
 
     /**
      * @typedef {object} ScopeStack
      * @property {ScopeStack | null} upper
-     * @property {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration} functionNode
+     * @property {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration | Program} scopeNode
      */
     /** @type {ScopeStack | null} */
     let scopeStack = null
 
-    return Object.assign(
+    return utils.compositingVisitors(
       {
-        Program() {
+        /** @param {Program} node */
+        Program(node) {
+          scopeStack = {
+            upper: scopeStack,
+            scopeNode: node
+          }
+
           const tracker = new ReferenceTracker(context.getScope())
           const traceMap = {
             vue: {
@@ -87,39 +98,39 @@ module.exports = {
           for (const { node } of tracker.iterateEsmReferences(traceMap)) {
             watchCallNodes.add(node)
           }
-        }
-      },
-      utils.defineVueVisitor(context, {
+        },
         /** @param {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration} node */
         ':function'(node) {
           scopeStack = {
             upper: scopeStack,
-            functionNode: node
+            scopeNode: node
           }
         },
-        onSetupFunctionEnter(node) {
-          setupFunctions.set(node, {
-            setupProperty: node.parent,
-            afterAwait: false
-          })
+        ':function:exit'() {
+          scopeStack = scopeStack && scopeStack.upper
         },
-        AwaitExpression() {
+        /** @param {AwaitExpression} node */
+        AwaitExpression(node) {
           if (!scopeStack) {
             return
           }
-          const setupFunctionData = setupFunctions.get(scopeStack.functionNode)
-          if (!setupFunctionData) {
+          const setupScope = setupScopes.get(scopeStack.scopeNode)
+          if (!setupScope || !utils.inRange(setupScope.range, node)) {
             return
           }
-          setupFunctionData.afterAwait = true
+          setupScope.afterAwait = true
         },
         /** @param {CallExpression} node */
         CallExpression(node) {
           if (!scopeStack) {
             return
           }
-          const setupFunctionData = setupFunctions.get(scopeStack.functionNode)
-          if (!setupFunctionData || !setupFunctionData.afterAwait) {
+          const setupScope = setupScopes.get(scopeStack.scopeNode)
+          if (
+            !setupScope ||
+            !setupScope.afterAwait ||
+            !utils.inRange(setupScope.range, node)
+          ) {
             return
           }
 
@@ -129,12 +140,34 @@ 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
+          })
         },
-        /** @param {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration} node */
-        ':function:exit'(node) {
-          scopeStack = scopeStack && scopeStack.upper
-
-          setupFunctions.delete(node)
+        onSetupFunctionExit(node) {
+          setupScopes.delete(node)
         }
       })
     )
diff --git a/tests/lib/rules/no-watch-after-await.js b/tests/lib/rules/no-watch-after-await.js
index c7cf73150..ae5718672 100644
--- a/tests/lib/rules/no-watch-after-await.js
+++ b/tests/lib/rules/no-watch-after-await.js
@@ -118,6 +118,43 @@ tester.run('no-watch-after-await', rule, {
       }
       </script>
       `
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      import {watchEffect} from 'vue'
+      watchEffect(() => { /* ... */ })
+      await doSomething()
+      </script>
+      `,
+      parserOptions: { ecmaVersion: 2022 }
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      await doSomething()
+      </script>
+      <script>
+      import {watchEffect} from 'vue'
+      watchEffect(() => { /* ... */ }) // not error
+      </script>
+      `,
+      parserOptions: { ecmaVersion: 2022 }
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      </script>
+      <script>
+      import {watchEffect} from 'vue'
+      await doSomething()
+      watchEffect(() => { /* ... */ }) // not error
+      </script>
+      `,
+      parserOptions: { ecmaVersion: 2022 }
     }
   ],
   invalid: [
@@ -199,6 +236,26 @@ 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
+        }
+      ]
     }
   ]
 })