From b12e5c93d8b28142499e380997d2d17eceab82f4 Mon Sep 17 00:00:00 2001
From: Michel Weststrate <mweststrate@fb.com>
Date: Fri, 24 Jul 2020 20:34:57 +0100
Subject: [PATCH] fix: Set finalization can get stuck in a loop, fixes #628

---
 __tests__/regressions.js | 24 ++++++++++++++++++++++++
 src/core/finalize.ts     | 11 ++++++++---
 2 files changed, 32 insertions(+), 3 deletions(-)

diff --git a/__tests__/regressions.js b/__tests__/regressions.js
index e997f3e3..a8e65c71 100644
--- a/__tests__/regressions.js
+++ b/__tests__/regressions.js
@@ -116,5 +116,29 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
 			expect(state2.length).toBe(3)
 			expect(state2).toEqual([undefined, "v1", "v2"])
 		})
+
+		test("#628 set removal hangs", () => {
+			let arr = []
+			let set = new Set([arr])
+
+			let result = produce(set, draft1 => {
+				produce(draft1, draft2 => {
+					draft2.delete(arr)
+				})
+			})
+			expect(result).toEqual(new Set([[]])) // N.B. this outcome doesn't seem not correct, but then again,
+			// double produce without return looks iffy as well, so not sure what the expected outcome in the
+			// original report was
+		})
+
+		test("#628 - 2 set removal hangs", () => {
+			let arr = []
+			let set = new Set([arr])
+
+			let result = produce(set, draft2 => {
+				draft2.delete(arr)
+			})
+			expect(result).toEqual(new Set())
+		})
 	})
 }
diff --git a/src/core/finalize.ts b/src/core/finalize.ts
index d83c040a..273c6386 100644
--- a/src/core/finalize.ts
+++ b/src/core/finalize.ts
@@ -87,9 +87,14 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) {
 			state.type_ === ProxyTypeES5Object || state.type_ === ProxyTypeES5Array
 				? (state.copy_ = shallowCopy(state.draft_))
 				: state.copy_
-		// finalize all children of the copy
-		each(result as any, (key, childValue) =>
-			finalizeProperty(rootScope, state, result, key, childValue, path)
+		// Finalize all children of the copy
+		// For sets we clone before iterating, otherwise we can get in endless loop due to modifying during iteration, see #628
+		// Although the original test case doesn't seem valid anyway, so if this in the way we can turn the next line
+		// back to each(result, ....)
+		each(
+			state.type_ === ProxyTypeSet ? new Set(result) : result,
+			(key, childValue) =>
+				finalizeProperty(rootScope, state, result, key, childValue, path)
 		)
 		// everything inside is frozen, we can freeze here
 		maybeFreeze(rootScope, result, false)