From cab9541ffd20b38cfdc27ea2a29a2653cf2b42a0 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Thu, 16 Sep 2021 17:07:29 -0400
Subject: [PATCH] refactor(compiler-sfc): improve TLA codegen

---
 .../__snapshots__/compileScript.spec.ts.snap  | 117 ++++++------------
 .../__tests__/compileScript.spec.ts           |   2 +-
 packages/compiler-sfc/src/compileScript.ts    |  37 ++++--
 3 files changed, 70 insertions(+), 86 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
index fe8d095f4c5..74c0ad4670e 100644
--- a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
+++ b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
@@ -96,7 +96,7 @@ export default /*#__PURE__*/ Object.assign(__default__, {
 })"
 `;
 
-exports[`SFC compile <script setup> async/await detection await in expression statement 1`] = `
+exports[`SFC compile <script setup> async/await detection expression statement 1`] = `
 "import { withAsyncContext as _withAsyncContext } from 'vue'
 
 export default {
@@ -104,21 +104,10 @@ export default {
   expose()
 
 let __temp, __restore
-foo()
 ;(
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return 1
-  })),
-  __temp = await __temp,
-  __restore(),
-  __temp
-) + (
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return 2
-  })),
-  __temp = await __temp,
-  __restore(),
-  __temp
+  ([__temp,__restore] = _withAsyncContext(() => foo)),
+  await __temp,
+  __restore()
 )
 return {  }
 }
@@ -126,7 +115,7 @@ return {  }
 }"
 `;
 
-exports[`SFC compile <script setup> async/await detection expression statement 1`] = `
+exports[`SFC compile <script setup> async/await detection nested await 1`] = `
 "import { withAsyncContext as _withAsyncContext } from 'vue'
 
 export default {
@@ -135,12 +124,14 @@ export default {
 
 let __temp, __restore
 ;(
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return foo
-  })),
+  ([__temp,__restore] = _withAsyncContext(async () => ((
+  ([__temp,__restore] = _withAsyncContext(() => foo)),
   __temp = await __temp,
   __restore(),
   __temp
+)))),
+  await __temp,
+  __restore()
 )
 return {  }
 }
@@ -148,7 +139,7 @@ return {  }
 }"
 `;
 
-exports[`SFC compile <script setup> async/await detection nested await 1`] = `
+exports[`SFC compile <script setup> async/await detection nested await 2`] = `
 "import { withAsyncContext as _withAsyncContext } from 'vue'
 
 export default {
@@ -157,19 +148,14 @@ export default {
 
 let __temp, __restore
 ;(
-  ([__temp,__restore] = _withAsyncContext(async () => {
-    return ((
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return foo
-  })),
-  __temp = await __temp,
-  __restore(),
-  __temp
-))
-  })),
+  ([__temp,__restore] = _withAsyncContext(async () => (((
+  ([__temp,__restore] = _withAsyncContext(() => foo)),
   __temp = await __temp,
   __restore(),
   __temp
+))))),
+  await __temp,
+  __restore()
 )
 return {  }
 }
@@ -177,7 +163,7 @@ return {  }
 }"
 `;
 
-exports[`SFC compile <script setup> async/await detection nested await 2`] = `
+exports[`SFC compile <script setup> async/await detection nested await 3`] = `
 "import { withAsyncContext as _withAsyncContext } from 'vue'
 
 export default {
@@ -186,19 +172,19 @@ export default {
 
 let __temp, __restore
 ;(
-  ([__temp,__restore] = _withAsyncContext(async () => {
-    return (((
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return foo
-  })),
+  ([__temp,__restore] = _withAsyncContext(async () => ((
+  ([__temp,__restore] = _withAsyncContext(async () => ((
+  ([__temp,__restore] = _withAsyncContext(() => foo)),
   __temp = await __temp,
   __restore(),
   __temp
-)))
-  })),
+)))),
   __temp = await __temp,
   __restore(),
   __temp
+)))),
+  await __temp,
+  __restore()
 )
 return {  }
 }
@@ -206,7 +192,7 @@ return {  }
 }"
 `;
 
-exports[`SFC compile <script setup> async/await detection nested await 3`] = `
+exports[`SFC compile <script setup> async/await detection nested leading await in expression statement 1`] = `
 "import { withAsyncContext as _withAsyncContext } from 'vue'
 
 export default {
@@ -214,24 +200,14 @@ export default {
   expose()
 
 let __temp, __restore
+foo()
 ;(
-  ([__temp,__restore] = _withAsyncContext(async () => {
-    return ((
-  ([__temp,__restore] = _withAsyncContext(async () => {
-    return ((
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return foo
-  })),
+  ([__temp,__restore] = _withAsyncContext(() => 1)),
   __temp = await __temp,
   __restore(),
   __temp
-))
-  })),
-  __temp = await __temp,
-  __restore(),
-  __temp
-))
-  })),
+) + (
+  ([__temp,__restore] = _withAsyncContext(() => 2)),
   __temp = await __temp,
   __restore(),
   __temp
@@ -251,19 +227,13 @@ export default {
 
 let __temp, __restore
 if (ok) { (
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return foo
-  })),
-  __temp = await __temp,
-  __restore(),
-  __temp
+  ([__temp,__restore] = _withAsyncContext(() => foo)),
+  await __temp,
+  __restore()
 ) } else { (
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return bar
-  })),
-  __temp = await __temp,
-  __restore(),
-  __temp
+  ([__temp,__restore] = _withAsyncContext(() => bar)),
+  await __temp,
+  __restore()
 ) }
 return {  }
 }
@@ -280,9 +250,7 @@ export default {
 
 let __temp, __restore
 let a = _ref(1 + ((
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return foo
-  })),
+  ([__temp,__restore] = _withAsyncContext(() => foo)),
   __temp = await __temp,
   __restore(),
   __temp
@@ -346,12 +314,9 @@ export default {
 
 let __temp, __restore
 if (false) (
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return foo()
-  })),
-  __temp = await __temp,
-  __restore(),
-  __temp
+  ([__temp,__restore] = _withAsyncContext(() => foo())),
+  await __temp,
+  __restore()
 )
 return {  }
 }
@@ -368,9 +333,7 @@ export default {
 
 let __temp, __restore
 const a = 1 + ((
-  ([__temp,__restore] = _withAsyncContext(() => {
-    return foo
-  })),
+  ([__temp,__restore] = _withAsyncContext(() => foo)),
   __temp = await __temp,
   __restore(),
   __temp
diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index 9aec5707e01..135f6be1268 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -1089,7 +1089,7 @@ const emit = defineEmits(['a', 'b'])
     })
 
     // should prepend semicolon
-    test('await in expression statement', () => {
+    test('nested leading await in expression statement', () => {
       const code = assertAwaitDetection(`foo()\nawait 1 + await 2`)
       expect(code).toMatch(`foo()\n;(`)
     })
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index fc0dd171b9f..766fe98e70a 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -507,13 +507,28 @@ export function compileScript(
   }
 
   /**
-   * await foo() -->
+   * await foo()
+   * -->
+   * ;(
+   *   ([__temp,__restore] = withAsyncContext(() => foo())),
+   *   await __temp,
+   *   __restore()
+   * )
    *
-   * (([__temp, __restore] = withAsyncContext(async () => {
-   *   return foo()
-   * })),__temp=await __temp,__restore(),__temp)
+   * const a = await foo()
+   * -->
+   * const a = (
+   *   ([__temp, __restore] = withAsyncContext(() => foo())),
+   *   __temp = await __temp,
+   *   __restore(),
+   *   __temp
+   * )
    */
-  function processAwait(node: AwaitExpression, needSemi: boolean) {
+  function processAwait(
+    node: AwaitExpression,
+    needSemi: boolean,
+    isStatement: boolean
+  ) {
     const argumentStart =
       node.argument.extra && node.argument.extra.parenthesized
         ? (node.argument.extra.parenStart as number)
@@ -531,11 +546,13 @@ export function compileScript(
       argumentStart + startOffset,
       `${needSemi ? `;` : ``}(\n  ([__temp,__restore] = ${helper(
         `withAsyncContext`
-      )}(${containsNestedAwait ? `async ` : ``}() => {\n    return `
+      )}(${containsNestedAwait ? `async ` : ``}() => `
     )
     s.appendLeft(
       node.end! + startOffset,
-      `\n  })),\n  __temp = await __temp,\n  __restore(),\n  __temp\n)`
+      `)),\n  ${isStatement ? `` : `__temp = `}await __temp,\n  __restore()${
+        isStatement ? `` : `,\n  __temp`
+      }\n)`
     )
   }
 
@@ -937,7 +954,11 @@ export function compileScript(
             const needsSemi = scriptSetupAst.body.some(n => {
               return n.type === 'ExpressionStatement' && n.start === child.start
             })
-            processAwait(child, needsSemi)
+            processAwait(
+              child,
+              needsSemi,
+              parent.type === 'ExpressionStatement'
+            )
           }
         }
       })