fix: Set finalization can get stuck in a loop, fixes #628

Author: mweststrate

__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())
+		})
 	})
 }

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)