From b00910808433cbc59df3fc45b241fd9b7dbfd2a9 Mon Sep 17 00:00:00 2001
From: tanhauhau <lhtan93@gmail.com>
Date: Sat, 27 Nov 2021 23:30:05 +0800
Subject: [PATCH] opt-out loopGuard for async loops

---
 src/compiler/compile/Component.ts                | 16 +++++++++-------
 .../loop-protect-async-opt-out/_config.js        |  9 +++++++++
 .../loop-protect-async-opt-out/main.svelte       | 12 ++++++++++++
 .../loop-protect-generator-opt-out/_config.js    |  3 +++
 .../loop-protect-generator-opt-out/main.svelte   | 10 ++++++----
 5 files changed, 39 insertions(+), 11 deletions(-)
 create mode 100644 test/runtime/samples/loop-protect-async-opt-out/_config.js
 create mode 100644 test/runtime/samples/loop-protect-async-opt-out/main.svelte

diff --git a/src/compiler/compile/Component.ts b/src/compiler/compile/Component.ts
index 202e14fd98..0dd2fc6718 100644
--- a/src/compiler/compile/Component.ts
+++ b/src/compiler/compile/Component.ts
@@ -24,7 +24,7 @@ import TemplateScope from './nodes/shared/TemplateScope';
 import fuzzymatch from '../utils/fuzzymatch';
 import get_object from './utils/get_object';
 import Slot from './nodes/Slot';
-import { Node, ImportDeclaration, ExportNamedDeclaration, Identifier, ExpressionStatement, AssignmentExpression, Literal, Property, RestElement, ExportDefaultDeclaration, ExportAllDeclaration } from 'estree';
+import { Node, ImportDeclaration, ExportNamedDeclaration, Identifier, ExpressionStatement, AssignmentExpression, Literal, Property, RestElement, ExportDefaultDeclaration, ExportAllDeclaration, FunctionDeclaration, FunctionExpression } from 'estree';
 import add_to_set from './utils/add_to_set';
 import check_graph_for_cycles from './utils/check_graph_for_cycles';
 import { print, b } from 'code-red';
@@ -758,12 +758,13 @@ export default class Component {
 		};
 		let scope_updated = false;
 
-		let generator_count = 0;
+		const current_function_stack = [];
+		let current_function: FunctionDeclaration | FunctionExpression = null;
 
 		walk(content, {
 			enter(node: Node, parent: Node, prop, index) {
-				if ((node.type === 'FunctionDeclaration' || node.type === 'FunctionExpression') && node.generator === true) {
-					generator_count++;
+				if ((node.type === 'FunctionDeclaration' || node.type === 'FunctionExpression')) {
+					current_function_stack.push(current_function = node);
 				}
 
 				if (map.has(node)) {
@@ -792,12 +793,13 @@ export default class Component {
 			},
 
 			leave(node: Node) {
-				if ((node.type === 'FunctionDeclaration' || node.type === 'FunctionExpression') && node.generator === true) {
-					generator_count--;
+				if ((node.type === 'FunctionDeclaration' || node.type === 'FunctionExpression')) {
+					current_function_stack.pop();
+					current_function = current_function_stack[current_function_stack.length - 1];
 				}
 
 				// do it on leave, to prevent infinite loop
-				if (component.compile_options.dev && component.compile_options.loopGuardTimeout > 0 && generator_count <= 0) {
+				if (component.compile_options.dev && component.compile_options.loopGuardTimeout > 0 && (!current_function || (!current_function.generator && !current_function.async))) {
 					const to_replace_for_loop_protect = component.loop_protect(node, scope, component.compile_options.loopGuardTimeout);
 					if (to_replace_for_loop_protect) {
 						this.replace(to_replace_for_loop_protect);
diff --git a/test/runtime/samples/loop-protect-async-opt-out/_config.js b/test/runtime/samples/loop-protect-async-opt-out/_config.js
new file mode 100644
index 0000000000..b1855c83eb
--- /dev/null
+++ b/test/runtime/samples/loop-protect-async-opt-out/_config.js
@@ -0,0 +1,9 @@
+export default {
+	compileOptions: {
+		dev: true,
+		loopGuardTimeout: 1
+	},
+	async test({ component }) {
+		await component.run_async_function();
+	}
+};
diff --git a/test/runtime/samples/loop-protect-async-opt-out/main.svelte b/test/runtime/samples/loop-protect-async-opt-out/main.svelte
new file mode 100644
index 0000000000..92ebf1bf93
--- /dev/null
+++ b/test/runtime/samples/loop-protect-async-opt-out/main.svelte
@@ -0,0 +1,12 @@
+<script>
+	export async function run_async_function() {
+		await loop();
+	};
+
+	async function loop() {
+		let i = 0;
+		while (i ++ < 5) {
+			await new Promise(resolve => setTimeout(resolve, 20));
+		}
+	}
+</script>
\ No newline at end of file
diff --git a/test/runtime/samples/loop-protect-generator-opt-out/_config.js b/test/runtime/samples/loop-protect-generator-opt-out/_config.js
index 0fe83a36db..2680ecfdfb 100644
--- a/test/runtime/samples/loop-protect-generator-opt-out/_config.js
+++ b/test/runtime/samples/loop-protect-generator-opt-out/_config.js
@@ -2,5 +2,8 @@ export default {
 	compileOptions: {
 		dev: true,
 		loopGuardTimeout: 1
+	},
+	async test({ component }) {
+		await component.run_generator();
 	}
 };
diff --git a/test/runtime/samples/loop-protect-generator-opt-out/main.svelte b/test/runtime/samples/loop-protect-generator-opt-out/main.svelte
index a7aa790681..514089435d 100644
--- a/test/runtime/samples/loop-protect-generator-opt-out/main.svelte
+++ b/test/runtime/samples/loop-protect-generator-opt-out/main.svelte
@@ -1,10 +1,12 @@
 <script>
-	const it = gen();
-	it.next();
+	export async function run_generator() {
+		const it = gen();
+		it.next();
+
+		await new Promise(resolve => setTimeout(resolve, 20));
 
-	setTimeout(() => {
 		it.next();
-	}, 20)
+	}
 
 	function* gen() {
 		while (true) {