From ee5ea4be9195d83df24f194310c0e4c226181e35 Mon Sep 17 00:00:00 2001
From: Yosuke Ota <otameshiyo23@gmail.com>
Date: Sat, 3 Jul 2021 14:44:14 +0900
Subject: [PATCH] Update `vue/no-lifecycle-after-await` rule to support
 `<script setup>` (#1539)

---
 lib/rules/no-lifecycle-after-await.js         | 87 +++++++++++++------
 lib/utils/index.js                            | 27 +++++-
 tests/lib/rules/no-lifecycle-after-await.js   | 55 ++++++++++++
 typings/eslint-plugin-vue/util-types/utils.ts |  4 +
 4 files changed, 146 insertions(+), 27 deletions(-)

diff --git a/lib/rules/no-lifecycle-after-await.js b/lib/rules/no-lifecycle-after-await.js
index 7688d5b70..aba61794d 100644
--- a/lib/rules/no-lifecycle-after-await.js
+++ b/lib/rules/no-lifecycle-after-await.js
@@ -41,26 +41,34 @@ module.exports = {
   /** @param {RuleContext} context */
   create(context) {
     /**
-     * @typedef {object} SetupFunctionData
-     * @property {Property} setupProperty
+     * @typedef {object} SetupScopeData
      * @property {boolean} afterAwait
+     * @property {[number,number]} range
      */
     /**
      * @typedef {object} ScopeStack
      * @property {ScopeStack | null} upper
-     * @property {FunctionDeclaration | FunctionExpression | ArrowFunctionExpression} functionNode
+     * @property {FunctionDeclaration | FunctionExpression | ArrowFunctionExpression | Program} scopeNode
      */
     /** @type {Set<ESNode>} */
     const lifecycleHookCallNodes = new Set()
-    /** @type {Map<FunctionDeclaration | FunctionExpression | ArrowFunctionExpression, SetupFunctionData>} */
-    const setupFunctions = new Map()
+    /** @type {Map<FunctionDeclaration | FunctionExpression | ArrowFunctionExpression | Program, SetupScopeData>} */
+    const setupScopes = new Map()
 
     /** @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 = {
             /** @type {TraceMap} */
@@ -77,37 +85,41 @@ module.exports = {
           for (const { node } of tracker.iterateEsmReferences(traceMap)) {
             lifecycleHookCallNodes.add(node)
           }
-        }
-      },
-      utils.defineVueVisitor(context, {
+        },
+        /**
+         * @param {FunctionExpression | FunctionDeclaration | ArrowFunctionExpression} 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
           }
 
@@ -121,11 +133,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
+          })
         },
-        ':function:exit'(node) {
-          scopeStack = scopeStack && scopeStack.upper
-
-          setupFunctions.delete(node)
+        onSetupFunctionExit(node) {
+          setupScopes.delete(node)
         }
       })
     )
diff --git a/lib/utils/index.js b/lib/utils/index.js
index d259b486f..a75a212d8 100644
--- a/lib/utils/index.js
+++ b/lib/utils/index.js
@@ -1003,7 +1003,12 @@ module.exports = {
         vueStack = vueStack.parent
       }
     }
-    if (visitor.onSetupFunctionEnter || visitor.onRenderFunctionEnter) {
+    if (
+      visitor.onSetupFunctionEnter ||
+      visitor.onSetupFunctionExit ||
+      visitor.onRenderFunctionEnter
+    ) {
+      const setups = new Set()
       /** @param { (FunctionExpression | ArrowFunctionExpression) & { parent: Property } } node */
       vueVisitor[
         'Property[value.type=/^(Arrow)?FunctionExpression$/] > :function'
@@ -1014,6 +1019,7 @@ module.exports = {
           const name = getStaticPropertyName(prop)
           if (name === 'setup') {
             callVisitor('onSetupFunctionEnter', node)
+            setups.add(node)
           } else if (name === 'render') {
             callVisitor('onRenderFunctionEnter', node)
           }
@@ -1023,6 +1029,17 @@ module.exports = {
           node
         )
       }
+      if (visitor.onSetupFunctionExit) {
+        /** @param { (FunctionExpression | ArrowFunctionExpression) & { parent: Property } } node */
+        vueVisitor[
+          'Property[value.type=/^(Arrow)?FunctionExpression$/] > :function:exit'
+        ] = (node) => {
+          if (setups.has(node)) {
+            callVisitor('onSetupFunctionExit', node)
+            setups.delete(node)
+          }
+        }
+      }
     }
 
     return vueVisitor
@@ -1511,6 +1528,14 @@ module.exports = {
     }
     return dp[alen][blen]
   },
+  /**
+   * Checks whether the target node is within the given range.
+   * @param { [number, number] } range
+   * @param {ASTNode} target
+   */
+  inRange(range, target) {
+    return range[0] <= target.range[0] && target.range[1] <= range[1]
+  },
   /**
    * Checks whether the given node is Property.
    */
diff --git a/tests/lib/rules/no-lifecycle-after-await.js b/tests/lib/rules/no-lifecycle-after-await.js
index 17dfada47..9fc17f3ab 100644
--- a/tests/lib/rules/no-lifecycle-after-await.js
+++ b/tests/lib/rules/no-lifecycle-after-await.js
@@ -106,6 +106,43 @@ tester.run('no-lifecycle-after-await', rule, {
       }
       </script>
       `
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      import {onMounted} from 'vue'
+      onMounted(() => { /* ... */ })
+      await doSomething()
+      </script>
+      `,
+      parserOptions: { ecmaVersion: 2022 }
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      await doSomething()
+      </script>
+      <script>
+      import {onMounted} from 'vue'
+      onMounted(() => { /* ... */ }) // not error
+      </script>
+      `,
+      parserOptions: { ecmaVersion: 2022 }
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      </script>
+      <script>
+      import {onMounted} from 'vue'
+      await doSomething()
+      onMounted(() => { /* ... */ }) // not error
+      </script>
+      `,
+      parserOptions: { ecmaVersion: 2022 }
     }
   ],
   invalid: [
@@ -224,6 +261,24 @@ 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/typings/eslint-plugin-vue/util-types/utils.ts b/typings/eslint-plugin-vue/util-types/utils.ts
index 86e2648b3..adece39a9 100644
--- a/typings/eslint-plugin-vue/util-types/utils.ts
+++ b/typings/eslint-plugin-vue/util-types/utils.ts
@@ -19,6 +19,10 @@ export interface VueVisitor extends VueVisitorBase {
     node: (FunctionExpression | ArrowFunctionExpression) & { parent: Property },
     obj: VueObjectData
   ): void
+  onSetupFunctionExit?(
+    node: (FunctionExpression | ArrowFunctionExpression) & { parent: Property },
+    obj: VueObjectData
+  ): void
   onRenderFunctionEnter?(
     node: (FunctionExpression | ArrowFunctionExpression) & { parent: Property },
     obj: VueObjectData