From 4c43d4e5b9df8732b601a269bf4030f9721d466f Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Fri, 10 Jul 2020 18:00:13 -0400
Subject: [PATCH] feat(compiler-sfc): `<script setup>` support (experimental)

This is the last commit for the feature which adds async/await detection.
---
 .../src/transforms/transformExpression.ts     |  5 +-
 .../__tests__/compileScript.spec.ts           | 43 +++++++++++++++--
 packages/compiler-sfc/src/compileScript.ts    | 46 +++++++++++++++----
 3 files changed, 80 insertions(+), 14 deletions(-)

diff --git a/packages/compiler-core/src/transforms/transformExpression.ts b/packages/compiler-core/src/transforms/transformExpression.ts
index 1e223067..6256569f 100644
--- a/packages/compiler-core/src/transforms/transformExpression.ts
+++ b/packages/compiler-core/src/transforms/transformExpression.ts
@@ -263,8 +263,9 @@ export function processExpression(
   return ret
 }
 
-const isFunction = (node: Node): node is Function =>
-  /Function(Expression|Declaration)$/.test(node.type)
+const isFunction = (node: Node): node is Function => {
+  return /Function(?:Expression|Declaration)$|Method$/.test(node.type)
+}
 
 const isStaticProperty = (node: Node): node is ObjectProperty =>
   node &&
diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index 4be8f5e1..4015683f 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -49,10 +49,6 @@ describe('SFC compile <script setup>', () => {
     )
   })
 
-  test('async/await detection', () => {
-    // TODO
-  })
-
   describe('exports', () => {
     test('export const x = ...', () => {
       const { content, bindings } = compile(
@@ -333,6 +329,45 @@ describe('SFC compile <script setup>', () => {
     })
   })
 
+  describe('async/await detection', () => {
+    function assertAwaitDetection(code: string, shouldAsync = true) {
+      const { content } = compile(`<script setup>${code}</script>`)
+      expect(content).toMatch(
+        `export ${shouldAsync ? `async ` : ``}function setup`
+      )
+    }
+
+    test('expression statement', () => {
+      assertAwaitDetection(`await foo`)
+    })
+
+    test('variable', () => {
+      assertAwaitDetection(`const a = 1 + (await foo)`)
+    })
+
+    test('export', () => {
+      assertAwaitDetection(`export const a = 1 + (await foo)`)
+    })
+
+    test('nested statements', () => {
+      assertAwaitDetection(`if (ok) { await foo } else { await bar }`)
+    })
+
+    test('should ignore await inside functions', () => {
+      // function declaration
+      assertAwaitDetection(`export async function foo() { await bar }`, false)
+      // function expression
+      assertAwaitDetection(`const foo = async () => { await bar }`, false)
+      // object method
+      assertAwaitDetection(`const obj = { async method() { await bar }}`, false)
+      // class method
+      assertAwaitDetection(
+        `const cls = class Foo { async method() { await bar }}`,
+        false
+      )
+    })
+  })
+
   describe('errors', () => {
     test('<script> and <script setup> must have same lang', () => {
       expect(
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index 33d3c40f..3bcd2b2f 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -11,6 +11,7 @@ import {
   ExpressionStatement,
   ArrowFunctionExpression,
   ExportSpecifier,
+  Function as FunctionNode,
   TSType,
   TSTypeLiteral,
   TSFunctionType,
@@ -87,7 +88,8 @@ export function compileScript(
   const setupExports: Record<string, boolean> = {}
   let exportAllIndex = 0
   let defaultExport: Node | undefined
-  let needDefaultExportRefCheck: boolean = false
+  let needDefaultExportRefCheck = false
+  let hasAwait = false
 
   const checkDuplicateDefaultExport = (node: Node) => {
     if (defaultExport) {
@@ -230,7 +232,11 @@ export function compileScript(
 
   // 3. parse <script setup> and  walk over top level statements
   for (const node of parse(scriptSetup.content, {
-    plugins,
+    plugins: [
+      ...plugins,
+      // allow top level await but only inside <script setup>
+      'topLevelAwait'
+    ],
     sourceType: 'module'
   }).program.body) {
     const start = node.start! + startOffset
@@ -439,6 +445,27 @@ export function compileScript(
       recordType(node, declaredTypes)
       s.move(start, end, 0)
     }
+
+    // walk statements & named exports / variable declarations for top level
+    // await
+    if (
+      node.type === 'VariableDeclaration' ||
+      (node.type === 'ExportNamedDeclaration' &&
+        node.declaration &&
+        node.declaration.type === 'VariableDeclaration') ||
+      node.type.endsWith('Statement')
+    ) {
+      ;(walk as any)(node, {
+        enter(node: Node) {
+          if (isFunction(node)) {
+            this.skip()
+          }
+          if (node.type === 'AwaitExpression') {
+            hasAwait = true
+          }
+        }
+      })
+    }
   }
 
   // 4. check default export to make sure it doesn't reference setup scope
@@ -503,7 +530,10 @@ export function compileScript(
   // 6. wrap setup code with function.
   // export the content of <script setup> as a named export, `setup`.
   // this allows `import { setup } from '*.vue'` for testing purposes.
-  s.prependLeft(startOffset, `\nexport function setup(${args}) {\n`)
+  s.prependLeft(
+    startOffset,
+    `\nexport ${hasAwait ? `async ` : ``}function setup(${args}) {\n`
+  )
 
   // generate return statement
   let returned = `{ ${Object.keys(setupExports).join(', ')} }`
@@ -867,11 +897,7 @@ function checkDefaultExport(
               )
           )
         }
-      } else if (
-        node.type === 'FunctionDeclaration' ||
-        node.type === 'FunctionExpression' ||
-        node.type === 'ArrowFunctionExpression'
-      ) {
+      } else if (isFunction(node)) {
         // walk function expressions and add its arguments to known identifiers
         // so that we don't prefix them
         node.params.forEach(p =>
@@ -927,6 +953,10 @@ function isStaticPropertyKey(node: Node, parent: Node): boolean {
   )
 }
 
+function isFunction(node: Node): node is FunctionNode {
+  return /Function(?:Expression|Declaration)$|Method$/.test(node.type)
+}
+
 /**
  * Analyze bindings in normal `<script>`
  * Note that `compileScriptSetup` already analyzes bindings as part of its