From 9f785b39c06620cc2bdec20c17bb5aba783d645e Mon Sep 17 00:00:00 2001
From: "Lyu, Wei-Da" <36730922+jasonlyu123@users.noreply.github.com>
Date: Fri, 30 Apr 2021 22:18:07 +0800
Subject: [PATCH] Fix {#await} block scope when updating (#6219)

---
 .../compile/render_dom/wrappers/AwaitBlock.ts | 19 ++-----
 src/runtime/internal/await_block.ts           | 14 +++++
 .../await-with-update-catch-scope/_config.js  | 51 +++++++++++++++++++
 .../await-with-update-catch-scope/main.svelte | 13 +++++
 4 files changed, 81 insertions(+), 16 deletions(-)
 create mode 100644 test/runtime/samples/await-with-update-catch-scope/_config.js
 create mode 100644 test/runtime/samples/await-with-update-catch-scope/main.svelte

diff --git a/src/compiler/compile/render_dom/wrappers/AwaitBlock.ts b/src/compiler/compile/render_dom/wrappers/AwaitBlock.ts
index cc186d1c02..637d32676c 100644
--- a/src/compiler/compile/render_dom/wrappers/AwaitBlock.ts
+++ b/src/compiler/compile/render_dom/wrappers/AwaitBlock.ts
@@ -232,14 +232,7 @@ export default class AwaitBlockWrapper extends Wrapper {
 
 		const dependencies = this.node.expression.dynamic_dependencies();
 
-		let update_child_context;
-		if (this.then.value && this.catch.value) {
-			update_child_context = b`#child_ctx[${this.then.value_index}] = #child_ctx[${this.catch.value_index}] = ${info}.resolved;`;
-		} else if (this.then.value) {
-			update_child_context = b`#child_ctx[${this.then.value_index}] = ${info}.resolved;`;
-		} else if (this.catch.value) {
-			update_child_context = b`#child_ctx[${this.catch.value_index}] = ${info}.resolved;`;
-		}
+		const update_await_block_branch = b`@update_await_block_branch(${info}, #ctx, #dirty)`;
 
 		if (dependencies.length > 0) {
 			const condition = x`
@@ -256,9 +249,7 @@ export default class AwaitBlockWrapper extends Wrapper {
 					if (${condition}) {
 
 					} else {
-						const #child_ctx = #ctx.slice();
-						${update_child_context}
-						${info}.block.p(#child_ctx, #dirty);
+						${update_await_block_branch}
 					}
 				`);
 			} else {
@@ -269,11 +260,7 @@ export default class AwaitBlockWrapper extends Wrapper {
 		} else {
 			if (this.pending.block.has_update_method) {
 				block.chunks.update.push(b`
-					{
-						const #child_ctx = #ctx.slice();
-						${update_child_context}
-						${info}.block.p(#child_ctx, #dirty);
-					}
+					${update_await_block_branch}
 				`);
 			}
 		}
diff --git a/src/runtime/internal/await_block.ts b/src/runtime/internal/await_block.ts
index b93f216b32..ea6e8a187f 100644
--- a/src/runtime/internal/await_block.ts
+++ b/src/runtime/internal/await_block.ts
@@ -83,3 +83,17 @@ export function handle_promise(promise, info) {
 		info.resolved = promise;
 	}
 }
+
+export function update_await_block_branch(info, ctx, dirty) {
+	const child_ctx = ctx.slice();
+	const { resolved } = info;
+
+	if (info.current === info.then) {
+		child_ctx[info.value] = resolved;
+	}
+	if (info.current === info.catch) {
+		child_ctx[info.error] = resolved;
+	}
+
+	info.block.p(child_ctx, dirty);
+}
diff --git a/test/runtime/samples/await-with-update-catch-scope/_config.js b/test/runtime/samples/await-with-update-catch-scope/_config.js
new file mode 100644
index 0000000000..4c60ef2150
--- /dev/null
+++ b/test/runtime/samples/await-with-update-catch-scope/_config.js
@@ -0,0 +1,51 @@
+export default {
+	props: {
+		thePromise: new Promise((_) => {})
+	},
+
+	html: `
+		<div>error: undefined</div>
+	`,
+
+	async test({ assert, component, target }) {
+		await (component.thePromise = Promise.resolve('abc'));
+
+		assert.htmlEqual(
+			target.innerHTML,
+			`
+			<div>
+				error: undefined
+				After Resolve: undefined
+			</div>
+			`
+		);
+
+		component.error = 'external error occurred';
+
+		assert.htmlEqual(
+			target.innerHTML,
+			`
+			<div>
+				error: ${component.error}
+				After Resolve: ${component.error}
+			</div>
+			`
+		);
+
+		try {
+			await (component.thePromise = Promise.reject('failure'));
+		} catch (error) {
+			// ignore
+		}
+
+		assert.htmlEqual(
+			target.innerHTML,
+			`
+			<div>
+				error: ${component.error}
+				Rejected: failure
+			</div>
+			`
+		);
+	}
+};
diff --git a/test/runtime/samples/await-with-update-catch-scope/main.svelte b/test/runtime/samples/await-with-update-catch-scope/main.svelte
new file mode 100644
index 0000000000..4d4eb53f0b
--- /dev/null
+++ b/test/runtime/samples/await-with-update-catch-scope/main.svelte
@@ -0,0 +1,13 @@
+<script>
+	export let thePromise;
+	export let error;
+</script>
+
+<div>
+	error: {error}
+	{#await thePromise then _}
+		After Resolve: {error}
+	{:catch error}
+		Rejected: {error}
+	{/await}
+</div>