From b1321910b354165df345f4c20c89e0d492eaf36c Mon Sep 17 00:00:00 2001
From: Yosuke Ota <otameshiyo23@gmail.com>
Date: Tue, 6 Jul 2021 18:23:18 +0900
Subject: [PATCH] Update `vue/no-restricted-call-after-await` rule to support
 `<script setup>` (#1551)

---
 lib/rules/no-restricted-call-after-await.js   | 79 +++++++++++++------
 .../rules/no-restricted-call-after-await.js   | 31 ++++++++
 2 files changed, 87 insertions(+), 23 deletions(-)

diff --git a/lib/rules/no-restricted-call-after-await.js b/lib/rules/no-restricted-call-after-await.js
index 6723d91ce..0c1989418 100644
--- a/lib/rules/no-restricted-call-after-await.js
+++ b/lib/rules/no-restricted-call-after-await.js
@@ -54,15 +54,21 @@ module.exports = {
   },
   /** @param {RuleContext} context */
   create(context) {
+    /**
+     * @typedef {object} SetupScopeData
+     * @property {boolean} afterAwait
+     * @property {[number,number]} range
+     */
+
     /** @type {Map<ESNode, string>} */
     const restrictedCallNodes = new Map()
-    /** @type {Map<FunctionExpression | ArrowFunctionExpression | FunctionDeclaration, { setupProperty: Property, afterAwait: boolean }>} */
-    const setupFunctions = new Map()
+    /** @type {Map<FunctionExpression | ArrowFunctionExpression | FunctionDeclaration | Program, SetupScopeData>} */
+    const setupScopes = new Map()
 
     /**x
      * @typedef {object} ScopeStack
      * @property {ScopeStack | null} upper
-     * @property {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration} functionNode
+     * @property {FunctionExpression | ArrowFunctionExpression | FunctionDeclaration | Program} scopeNode
      */
     /** @type {ScopeStack | null} */
     let scopeStack = null
@@ -136,6 +142,11 @@ module.exports = {
       {
         /** @param {Program} node */
         Program(node) {
+          scopeStack = {
+            upper: scopeStack,
+            scopeNode: node
+          }
+
           const tracker = new ReferenceTracker(context.getScope())
 
           for (const option of context.options) {
@@ -170,39 +181,39 @@ module.exports = {
               }
             }
           }
-        }
-      },
-      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
           }
 
@@ -214,12 +225,34 @@ 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
+          })
         },
-        /** @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-restricted-call-after-await.js b/tests/lib/rules/no-restricted-call-after-await.js
index 530e7cd39..92353eaf0 100644
--- a/tests/lib/rules/no-restricted-call-after-await.js
+++ b/tests/lib/rules/no-restricted-call-after-await.js
@@ -125,6 +125,18 @@ tester.run('no-restricted-call-after-await', rule, {
         { module: './baz', path: 'qux' },
         { module: 'vue-i18n', path: 'useI18n' }
       ]
+    },
+    {
+      filename: 'test.vue',
+      code: `
+      <script setup>
+      import {useI18n} from 'vue-i18n'
+      useI18n()
+      await doSomething()
+      </script>
+      `,
+      options: [{ module: 'vue-i18n', path: 'useI18n' }],
+      parserOptions: { ecmaVersion: 2022 }
     }
   ],
   invalid: [
@@ -377,6 +389,25 @@ 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
+        }
+      ]
     }
   ]
 })