From 39cebf5f7a8f72338030844fca4a75ffc913c518 Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Thu, 16 Sep 2021 16:49:59 -0400
Subject: [PATCH] fix(compiler-sfc): fix TLA codegen semicolon insertion

fix #4596
---
 .../__snapshots__/compileScript.spec.ts.snap  | 74 +++++++++++++++++--
 .../__tests__/compileScript.spec.ts           | 14 ++++
 packages/compiler-sfc/src/compileScript.ts    | 21 +++---
 3 files changed, 92 insertions(+), 17 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
index 855ddd4a610..fe8d095f4c5 100644
--- a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
+++ b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
@@ -96,6 +96,36 @@ export default /*#__PURE__*/ Object.assign(__default__, {
 })"
 `;
 
+exports[`SFC compile <script setup> async/await detection await in expression statement 1`] = `
+"import { withAsyncContext as _withAsyncContext } from 'vue'
+
+export default {
+  async setup(__props, { expose }) {
+  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
+)
+return {  }
+}
+
+}"
+`;
+
 exports[`SFC compile <script setup> async/await detection expression statement 1`] = `
 "import { withAsyncContext as _withAsyncContext } from 'vue'
 
@@ -109,7 +139,8 @@ let __temp, __restore
     return foo
   })),
   __temp = await __temp,
-  __restore()
+  __restore(),
+  __temp
 )
 return {  }
 }
@@ -137,7 +168,8 @@ let __temp, __restore
 ))
   })),
   __temp = await __temp,
-  __restore()
+  __restore(),
+  __temp
 )
 return {  }
 }
@@ -165,7 +197,8 @@ let __temp, __restore
 )))
   })),
   __temp = await __temp,
-  __restore()
+  __restore(),
+  __temp
 )
 return {  }
 }
@@ -200,7 +233,8 @@ let __temp, __restore
 ))
   })),
   __temp = await __temp,
-  __restore()
+  __restore(),
+  __temp
 )
 return {  }
 }
@@ -216,18 +250,20 @@ export default {
   expose()
 
 let __temp, __restore
-if (ok) { ;(
+if (ok) { (
   ([__temp,__restore] = _withAsyncContext(() => {
     return foo
   })),
   __temp = await __temp,
-  __restore()
-) } else { ;(
+  __restore(),
+  __temp
+) } else { (
   ([__temp,__restore] = _withAsyncContext(() => {
     return bar
   })),
   __temp = await __temp,
-  __restore()
+  __restore(),
+  __temp
 ) }
 return {  }
 }
@@ -301,6 +337,28 @@ return { cls }
 }"
 `;
 
+exports[`SFC compile <script setup> async/await detection single line conditions 1`] = `
+"import { withAsyncContext as _withAsyncContext } from 'vue'
+
+export default {
+  async setup(__props, { expose }) {
+  expose()
+
+let __temp, __restore
+if (false) (
+  ([__temp,__restore] = _withAsyncContext(() => {
+    return foo()
+  })),
+  __temp = await __temp,
+  __restore(),
+  __temp
+)
+return {  }
+}
+
+}"
+`;
+
 exports[`SFC compile <script setup> async/await detection variable 1`] = `
 "import { withAsyncContext as _withAsyncContext } from 'vue'
 
diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index 64b4fcfcf1b..9aec5707e01 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -1066,6 +1066,7 @@ const emit = defineEmits(['a', 'b'])
       }
       expect(content).toMatch(`${shouldAsync ? `async ` : ``}setup(`)
       assertCode(content)
+      return content
     }
 
     test('expression statement', () => {
@@ -1080,12 +1081,25 @@ const emit = defineEmits(['a', 'b'])
       assertAwaitDetection(`let a = $ref(1 + (await foo))`)
     })
 
+    // #4448
     test('nested await', () => {
       assertAwaitDetection(`await (await foo)`)
       assertAwaitDetection(`await ((await foo))`)
       assertAwaitDetection(`await (await (await foo))`)
     })
 
+    // should prepend semicolon
+    test('await in expression statement', () => {
+      const code = assertAwaitDetection(`foo()\nawait 1 + await 2`)
+      expect(code).toMatch(`foo()\n;(`)
+    })
+
+    // #4596 should NOT prepend semicolon
+    test('single line conditions', () => {
+      const code = assertAwaitDetection(`if (false) await foo()`)
+      expect(code).not.toMatch(`if (false) ;(`)
+    })
+
     test('nested statements', () => {
       assertAwaitDetection(`if (ok) { await foo } else { await bar }`)
     })
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index 848bb322251..fc0dd171b9f 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -507,11 +507,13 @@ export function compileScript(
   }
 
   /**
-   * await foo()
-   * -->
-   * (([__temp, __restore] = withAsyncContext(async () => foo())),__temp=await __temp,__restore(),__temp)
+   * await foo() -->
+   *
+   * (([__temp, __restore] = withAsyncContext(async () => {
+   *   return foo()
+   * })),__temp=await __temp,__restore(),__temp)
    */
-  function processAwait(node: AwaitExpression, isStatement: boolean) {
+  function processAwait(node: AwaitExpression, needSemi: boolean) {
     const argumentStart =
       node.argument.extra && node.argument.extra.parenthesized
         ? (node.argument.extra.parenStart as number)
@@ -527,15 +529,13 @@ export function compileScript(
     s.overwrite(
       node.start! + startOffset,
       argumentStart + startOffset,
-      `${isStatement ? `;` : ``}(\n  ([__temp,__restore] = ${helper(
+      `${needSemi ? `;` : ``}(\n  ([__temp,__restore] = ${helper(
         `withAsyncContext`
       )}(${containsNestedAwait ? `async ` : ``}() => {\n    return `
     )
     s.appendLeft(
       node.end! + startOffset,
-      `\n  })),\n  __temp = await __temp,\n  __restore()${
-        isStatement ? `` : `,\n  __temp`
-      }\n)`
+      `\n  })),\n  __temp = await __temp,\n  __restore(),\n  __temp\n)`
     )
   }
 
@@ -934,7 +934,10 @@ export function compileScript(
           }
           if (child.type === 'AwaitExpression') {
             hasAwait = true
-            processAwait(child, parent.type === 'ExpressionStatement')
+            const needsSemi = scriptSetupAst.body.some(n => {
+              return n.type === 'ExpressionStatement' && n.start === child.start
+            })
+            processAwait(child, needsSemi)
           }
         }
       })