From d297c3943d916285af3e02835c21e66d01370ee1 Mon Sep 17 00:00:00 2001
From: daiwei <daiwei521@126.com>
Date: Fri, 27 Aug 2021 17:47:54 +0800
Subject: [PATCH 1/4] fix(compiler-sfc): support nested await statements

---
 .../compiler-sfc/__tests__/compileScript.spec.ts  | 15 +++++++++++++++
 packages/compiler-sfc/src/compileScript.ts        |  5 ++++-
 2 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index 35ad4876bad..3f467ca17db 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -1097,6 +1097,21 @@ const emit = defineEmits(['a', 'b'])
       )
     })
 
+    test('nested await', () => {
+      assertAwaitDetection(
+        `await (await foo)`,
+        `;(([__temp,__restore]=_withAsyncContext(()=>(((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))))),__temp=await __temp,__restore())`
+      )
+      assertAwaitDetection(
+        `await ((await foo))`,
+        `;(([__temp,__restore]=_withAsyncContext(()=>((((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp)))))),__temp=await __temp,__restore())`
+      )
+      assertAwaitDetection(
+        `await (await (await foo))`,
+        `;(([__temp,__restore]=_withAsyncContext(()=>(((([__temp,__restore]=_withAsyncContext(()=>(((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))))),__temp=await __temp,__restore(),__temp))))),__temp=await __temp,__restore())`
+      )
+    })
+
     test('nested statements', () => {
       assertAwaitDetection(`if (ok) { await foo } else { await bar }`, code => {
         return (
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index cd189be5c8f..4a46dfbc63c 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -512,9 +512,12 @@ export function compileScript(
    * (([__temp, __restore] = withAsyncContext(() => foo())),__temp=await __temp,__restore(),__temp)
    */
   function processAwait(node: AwaitExpression, isStatement: boolean) {
+    const end = node.argument.extra
+      ? (node.argument.extra.parenStart as number)
+      : node.argument.start!
     s.overwrite(
       node.start! + startOffset,
-      node.argument.start! + startOffset,
+      end + startOffset,
       `${isStatement ? `;` : ``}(([__temp,__restore]=${helper(
         `withAsyncContext`
       )}(()=>(`

From 15de5a2eecaf3d5929098e2510cba4ce7ecadc03 Mon Sep 17 00:00:00 2001
From: daiwei <daiwei521@126.com>
Date: Fri, 27 Aug 2021 18:19:48 +0800
Subject: [PATCH 2/4] chore: improve test case

---
 .../__tests__/compileScript.spec.ts             | 17 ++++++++++++++---
 1 file changed, 14 insertions(+), 3 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index 3f467ca17db..3645ccb92da 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -1100,15 +1100,26 @@ const emit = defineEmits(['a', 'b'])
     test('nested await', () => {
       assertAwaitDetection(
         `await (await foo)`,
-        `;(([__temp,__restore]=_withAsyncContext(()=>(((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))))),__temp=await __temp,__restore())`
+        `;(([__temp,__restore]=_withAsyncContext(` +
+          `()=>(((([__temp,__restore]=_withAsyncContext(` +
+          `()=>(foo))),__temp=await __temp,__restore(),__temp))))` +
+          `),__temp=await __temp,__restore())`
       )
       assertAwaitDetection(
         `await ((await foo))`,
-        `;(([__temp,__restore]=_withAsyncContext(()=>((((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp)))))),__temp=await __temp,__restore())`
+        `;(([__temp,__restore]=_withAsyncContext(` +
+          `()=>((((([__temp,__restore]=_withAsyncContext(` +
+          `()=>(foo))),__temp=await __temp,__restore(),__temp)))))` +
+          `),__temp=await __temp,__restore())`
       )
       assertAwaitDetection(
         `await (await (await foo))`,
-        `;(([__temp,__restore]=_withAsyncContext(()=>(((([__temp,__restore]=_withAsyncContext(()=>(((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))))),__temp=await __temp,__restore(),__temp))))),__temp=await __temp,__restore())`
+        `;(([__temp,__restore]=_withAsyncContext(` +
+          `()=>(((([__temp,__restore]=_withAsyncContext(` +
+          `()=>(((([__temp,__restore]=_withAsyncContext(` +
+          `()=>(foo))),__temp=await __temp,__restore(),__temp))))` +
+          `),__temp=await __temp,__restore(),__temp))))` +
+          `),__temp=await __temp,__restore())`
       )
     })
 

From f4da5d71b3a38392bcfbbd1755213d2b92ebe648 Mon Sep 17 00:00:00 2001
From: daiwei <daiwei521@126.com>
Date: Fri, 27 Aug 2021 22:18:00 +0800
Subject: [PATCH 3/4] chore: improve code

---
 packages/compiler-sfc/__tests__/compileScript.spec.ts | 8 ++++----
 packages/compiler-sfc/src/compileScript.ts            | 2 +-
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index 3645ccb92da..3a066178632 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -1101,22 +1101,22 @@ const emit = defineEmits(['a', 'b'])
       assertAwaitDetection(
         `await (await foo)`,
         `;(([__temp,__restore]=_withAsyncContext(` +
-          `()=>(((([__temp,__restore]=_withAsyncContext(` +
+          `async ()=>(((([__temp,__restore]=_withAsyncContext(` +
           `()=>(foo))),__temp=await __temp,__restore(),__temp))))` +
           `),__temp=await __temp,__restore())`
       )
       assertAwaitDetection(
         `await ((await foo))`,
         `;(([__temp,__restore]=_withAsyncContext(` +
-          `()=>((((([__temp,__restore]=_withAsyncContext(` +
+          `async ()=>((((([__temp,__restore]=_withAsyncContext(` +
           `()=>(foo))),__temp=await __temp,__restore(),__temp)))))` +
           `),__temp=await __temp,__restore())`
       )
       assertAwaitDetection(
         `await (await (await foo))`,
         `;(([__temp,__restore]=_withAsyncContext(` +
-          `()=>(((([__temp,__restore]=_withAsyncContext(` +
-          `()=>(((([__temp,__restore]=_withAsyncContext(` +
+          `async ()=>(((([__temp,__restore]=_withAsyncContext(` +
+          `async ()=>(((([__temp,__restore]=_withAsyncContext(` +
           `()=>(foo))),__temp=await __temp,__restore(),__temp))))` +
           `),__temp=await __temp,__restore(),__temp))))` +
           `),__temp=await __temp,__restore())`
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index 4a46dfbc63c..7fb02cea464 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -520,7 +520,7 @@ export function compileScript(
       end + startOffset,
       `${isStatement ? `;` : ``}(([__temp,__restore]=${helper(
         `withAsyncContext`
-      )}(()=>(`
+      )}(${node.argument.type === 'AwaitExpression' ? `async ` : ``}()=>(`
     )
     s.appendLeft(
       node.end! + startOffset,

From 2b7bf4259ef5c38cae525502017c507f0d54e8ee Mon Sep 17 00:00:00 2001
From: Evan You <yyx990803@gmail.com>
Date: Thu, 16 Sep 2021 16:12:24 -0400
Subject: [PATCH 4/4] refactor: improve codegen and test

---
 .../__snapshots__/compileScript.spec.ts.snap  | 227 ++++++++++++++++++
 .../__tests__/compileScript.spec.ts           |  83 +------
 packages/compiler-sfc/src/compileScript.ts    |  27 ++-
 3 files changed, 258 insertions(+), 79 deletions(-)

diff --git a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
index c613cf50217..855ddd4a610 100644
--- a/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
+++ b/packages/compiler-sfc/__tests__/__snapshots__/compileScript.spec.ts.snap
@@ -96,6 +96,233 @@ export default /*#__PURE__*/ Object.assign(__default__, {
 })"
 `;
 
+exports[`SFC compile <script setup> async/await detection expression statement 1`] = `
+"import { withAsyncContext as _withAsyncContext } from 'vue'
+
+export default {
+  async setup(__props, { expose }) {
+  expose()
+
+let __temp, __restore
+;(
+  ([__temp,__restore] = _withAsyncContext(() => {
+    return foo
+  })),
+  __temp = await __temp,
+  __restore()
+)
+return {  }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection nested await 1`] = `
+"import { withAsyncContext as _withAsyncContext } from 'vue'
+
+export default {
+  async setup(__props, { expose }) {
+  expose()
+
+let __temp, __restore
+;(
+  ([__temp,__restore] = _withAsyncContext(async () => {
+    return ((
+  ([__temp,__restore] = _withAsyncContext(() => {
+    return foo
+  })),
+  __temp = await __temp,
+  __restore(),
+  __temp
+))
+  })),
+  __temp = await __temp,
+  __restore()
+)
+return {  }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection nested await 2`] = `
+"import { withAsyncContext as _withAsyncContext } from 'vue'
+
+export default {
+  async setup(__props, { expose }) {
+  expose()
+
+let __temp, __restore
+;(
+  ([__temp,__restore] = _withAsyncContext(async () => {
+    return (((
+  ([__temp,__restore] = _withAsyncContext(() => {
+    return foo
+  })),
+  __temp = await __temp,
+  __restore(),
+  __temp
+)))
+  })),
+  __temp = await __temp,
+  __restore()
+)
+return {  }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection nested await 3`] = `
+"import { withAsyncContext as _withAsyncContext } from 'vue'
+
+export default {
+  async setup(__props, { expose }) {
+  expose()
+
+let __temp, __restore
+;(
+  ([__temp,__restore] = _withAsyncContext(async () => {
+    return ((
+  ([__temp,__restore] = _withAsyncContext(async () => {
+    return ((
+  ([__temp,__restore] = _withAsyncContext(() => {
+    return foo
+  })),
+  __temp = await __temp,
+  __restore(),
+  __temp
+))
+  })),
+  __temp = await __temp,
+  __restore(),
+  __temp
+))
+  })),
+  __temp = await __temp,
+  __restore()
+)
+return {  }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection nested statements 1`] = `
+"import { withAsyncContext as _withAsyncContext } from 'vue'
+
+export default {
+  async setup(__props, { expose }) {
+  expose()
+
+let __temp, __restore
+if (ok) { ;(
+  ([__temp,__restore] = _withAsyncContext(() => {
+    return foo
+  })),
+  __temp = await __temp,
+  __restore()
+) } else { ;(
+  ([__temp,__restore] = _withAsyncContext(() => {
+    return bar
+  })),
+  __temp = await __temp,
+  __restore()
+) }
+return {  }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection ref 1`] = `
+"import { withAsyncContext as _withAsyncContext, ref as _ref } from 'vue'
+
+export default {
+  async setup(__props, { expose }) {
+  expose()
+
+let __temp, __restore
+let a = _ref(1 + ((
+  ([__temp,__restore] = _withAsyncContext(() => {
+    return foo
+  })),
+  __temp = await __temp,
+  __restore(),
+  __temp
+)))
+return { a }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection should ignore await inside functions 1`] = `
+"export default {
+  setup(__props, { expose }) {
+  expose()
+async function foo() { await bar }
+return { foo }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection should ignore await inside functions 2`] = `
+"export default {
+  setup(__props, { expose }) {
+  expose()
+const foo = async () => { await bar }
+return { foo }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection should ignore await inside functions 3`] = `
+"export default {
+  setup(__props, { expose }) {
+  expose()
+const obj = { async method() { await bar }}
+return { obj }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection should ignore await inside functions 4`] = `
+"export default {
+  setup(__props, { expose }) {
+  expose()
+const cls = class Foo { async method() { await bar }}
+return { cls }
+}
+
+}"
+`;
+
+exports[`SFC compile <script setup> async/await detection variable 1`] = `
+"import { withAsyncContext as _withAsyncContext } from 'vue'
+
+export default {
+  async setup(__props, { expose }) {
+  expose()
+
+let __temp, __restore
+const a = 1 + ((
+  ([__temp,__restore] = _withAsyncContext(() => {
+    return foo
+  })),
+  __temp = await __temp,
+  __restore(),
+  __temp
+))
+return { a }
+}
+
+}"
+`;
+
 exports[`SFC compile <script setup> binding analysis for destructur 1`] = `
 "export default {
   setup(__props, { expose }) {
diff --git a/packages/compiler-sfc/__tests__/compileScript.spec.ts b/packages/compiler-sfc/__tests__/compileScript.spec.ts
index 3a066178632..64b4fcfcf1b 100644
--- a/packages/compiler-sfc/__tests__/compileScript.spec.ts
+++ b/packages/compiler-sfc/__tests__/compileScript.spec.ts
@@ -1057,11 +1057,7 @@ const emit = defineEmits(['a', 'b'])
   })
 
   describe('async/await detection', () => {
-    function assertAwaitDetection(
-      code: string,
-      expected: string | ((content: string) => boolean),
-      shouldAsync = true
-    ) {
+    function assertAwaitDetection(code: string, shouldAsync = true) {
       const { content } = compile(`<script setup>${code}</script>`, {
         refSugar: true
       })
@@ -1069,96 +1065,41 @@ const emit = defineEmits(['a', 'b'])
         expect(content).toMatch(`let __temp, __restore`)
       }
       expect(content).toMatch(`${shouldAsync ? `async ` : ``}setup(`)
-      if (typeof expected === 'string') {
-        expect(content).toMatch(expected)
-      } else {
-        expect(expected(content)).toBe(true)
-      }
+      assertCode(content)
     }
 
     test('expression statement', () => {
-      assertAwaitDetection(
-        `await foo`,
-        `;(([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore())`
-      )
+      assertAwaitDetection(`await foo`)
     })
 
     test('variable', () => {
-      assertAwaitDetection(
-        `const a = 1 + (await foo)`,
-        `1 + ((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))`
-      )
+      assertAwaitDetection(`const a = 1 + (await foo)`)
     })
 
     test('ref', () => {
-      assertAwaitDetection(
-        `let a = $ref(1 + (await foo))`,
-        `1 + ((([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore(),__temp))`
-      )
+      assertAwaitDetection(`let a = $ref(1 + (await foo))`)
     })
 
     test('nested await', () => {
-      assertAwaitDetection(
-        `await (await foo)`,
-        `;(([__temp,__restore]=_withAsyncContext(` +
-          `async ()=>(((([__temp,__restore]=_withAsyncContext(` +
-          `()=>(foo))),__temp=await __temp,__restore(),__temp))))` +
-          `),__temp=await __temp,__restore())`
-      )
-      assertAwaitDetection(
-        `await ((await foo))`,
-        `;(([__temp,__restore]=_withAsyncContext(` +
-          `async ()=>((((([__temp,__restore]=_withAsyncContext(` +
-          `()=>(foo))),__temp=await __temp,__restore(),__temp)))))` +
-          `),__temp=await __temp,__restore())`
-      )
-      assertAwaitDetection(
-        `await (await (await foo))`,
-        `;(([__temp,__restore]=_withAsyncContext(` +
-          `async ()=>(((([__temp,__restore]=_withAsyncContext(` +
-          `async ()=>(((([__temp,__restore]=_withAsyncContext(` +
-          `()=>(foo))),__temp=await __temp,__restore(),__temp))))` +
-          `),__temp=await __temp,__restore(),__temp))))` +
-          `),__temp=await __temp,__restore())`
-      )
+      assertAwaitDetection(`await (await foo)`)
+      assertAwaitDetection(`await ((await foo))`)
+      assertAwaitDetection(`await (await (await foo))`)
     })
 
     test('nested statements', () => {
-      assertAwaitDetection(`if (ok) { await foo } else { await bar }`, code => {
-        return (
-          code.includes(
-            `;(([__temp,__restore]=_withAsyncContext(()=>(foo))),__temp=await __temp,__restore())`
-          ) &&
-          code.includes(
-            `;(([__temp,__restore]=_withAsyncContext(()=>(bar))),__temp=await __temp,__restore())`
-          )
-        )
-      })
+      assertAwaitDetection(`if (ok) { await foo } else { await bar }`)
     })
 
     test('should ignore await inside functions', () => {
       // function declaration
-      assertAwaitDetection(
-        `async function foo() { await bar }`,
-        `await bar`,
-        false
-      )
+      assertAwaitDetection(`async function foo() { await bar }`, false)
       // function expression
-      assertAwaitDetection(
-        `const foo = async () => { await bar }`,
-        `await bar`,
-        false
-      )
+      assertAwaitDetection(`const foo = async () => { await bar }`, false)
       // object method
-      assertAwaitDetection(
-        `const obj = { async method() { await bar }}`,
-        `await bar`,
-        false
-      )
+      assertAwaitDetection(`const obj = { async method() { await bar }}`, false)
       // class method
       assertAwaitDetection(
         `const cls = class Foo { async method() { await bar }}`,
-        `await bar`,
         false
       )
     })
diff --git a/packages/compiler-sfc/src/compileScript.ts b/packages/compiler-sfc/src/compileScript.ts
index 7fb02cea464..848bb322251 100644
--- a/packages/compiler-sfc/src/compileScript.ts
+++ b/packages/compiler-sfc/src/compileScript.ts
@@ -509,22 +509,33 @@ export function compileScript(
   /**
    * await foo()
    * -->
-   * (([__temp, __restore] = withAsyncContext(() => foo())),__temp=await __temp,__restore(),__temp)
+   * (([__temp, __restore] = withAsyncContext(async () => foo())),__temp=await __temp,__restore(),__temp)
    */
   function processAwait(node: AwaitExpression, isStatement: boolean) {
-    const end = node.argument.extra
-      ? (node.argument.extra.parenStart as number)
-      : node.argument.start!
+    const argumentStart =
+      node.argument.extra && node.argument.extra.parenthesized
+        ? (node.argument.extra.parenStart as number)
+        : node.argument.start!
+
+    const argumentStr = source.slice(
+      argumentStart + startOffset,
+      node.argument.end! + startOffset
+    )
+
+    const containsNestedAwait = /\bawait\b/.test(argumentStr)
+
     s.overwrite(
       node.start! + startOffset,
-      end + startOffset,
-      `${isStatement ? `;` : ``}(([__temp,__restore]=${helper(
+      argumentStart + startOffset,
+      `${isStatement ? `;` : ``}(\n  ([__temp,__restore] = ${helper(
         `withAsyncContext`
-      )}(${node.argument.type === 'AwaitExpression' ? `async ` : ``}()=>(`
+      )}(${containsNestedAwait ? `async ` : ``}() => {\n    return `
     )
     s.appendLeft(
       node.end! + startOffset,
-      `))),__temp=await __temp,__restore()${isStatement ? `` : `,__temp`})`
+      `\n  })),\n  __temp = await __temp,\n  __restore()${
+        isStatement ? `` : `,\n  __temp`
+      }\n)`
     )
   }