fix: Preserve insertion order of Sets, fixes #819 (#976)

* Proposed fix for #819 to preserve Set insertion order

* Proposed fix for #819 (B)

We are already cloning and iterating over the Set. If we go a step further and clear it before iteration, and re-add all children during iteration, we can preserve insertion order.

Author: jdip

__tests__/map-set.js

@@ -296,4 +296,38 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
 			expect(mapType1).toBe(mapType2)
 		})
 	})
+	describe("set issues " + name, () => {
+		test("#819.A - maintains order when adding", () => {
+			const objs = [
+				"a",
+				{
+					id: "b"
+				}
+			]
+
+			const set = new Set([objs[0]])
+			const newSet = produce(set, draft => {
+				draft.add(objs[1])
+			})
+
+			// passes
+			expect(Array.from(newSet)).toEqual([objs[0], objs[1]])
+		})
+
+		test("#819.B - maintains order when adding", () => {
+			const objs = [
+				{
+					id: "a"
+				},
+				"b"
+			]
+
+			const set = new Set([objs[0]])
+			const newSet = produce(set, draft => {
+				draft.add(objs[1])
+			})
+
+			expect(Array.from(newSet)).toEqual([objs[0], objs[1]])
+		})
+	})
 }

src/core/finalize.ts

@@ -87,12 +87,17 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) {
 				: state.copy_
 		// 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_ === ProxyType.Set ? new Set(result) : result,
-			(key, childValue) =>
-				finalizeProperty(rootScope, state, result, key, childValue, path)
+		// To preserve insertion order in all cases we then clear the set
+		// And we let finalizeProperty know it needs to re-add non-draft children back to the target
+		let resultEach = result
+		let isSet = false
+		if (state.type_ === ProxyType.Set) {
+			resultEach = new Set(result)
+			result.clear()
+			isSet = true
+		}
+		each(resultEach, (key, childValue) =>
+			finalizeProperty(rootScope, state, result, key, childValue, path, isSet)
 		)
 		// everything inside is frozen, we can freeze here
 		maybeFreeze(rootScope, result, false)
@@ -115,7 +120,8 @@ function finalizeProperty(
 	targetObject: any,
 	prop: string | number,
 	childValue: any,
-	rootPath?: PatchPath
+	rootPath?: PatchPath,
+	targetIsSet?: boolean
 ) {
 	if (__DEV__ && childValue === targetObject) die(5)
 	if (isDraft(childValue)) {
@@ -134,6 +140,8 @@ function finalizeProperty(
 		if (isDraft(res)) {
 			rootScope.canAutoFreeze_ = false
 		} else return
+	} else if (targetIsSet) {
+		targetObject.add(childValue)
 	}
 	// Search new objects for unfinalized drafts. Frozen objects should never contain drafts.
 	if (isDraftable(childValue) && !isFrozen(childValue)) {

src/utils/common.ts

@@ -132,7 +132,6 @@ export function set(thing: any, propOrOldValue: PropertyKey, value: any) {
 	const t = getArchtype(thing)
 	if (t === Archtype.Map) thing.set(propOrOldValue, value)
 	else if (t === Archtype.Set) {
-		thing.delete(propOrOldValue)
 		thing.add(value)
 	} else thing[propOrOldValue] = value
 }