From c0e6749e8df3e20d880d61b726b1395167ba2088 Mon Sep 17 00:00:00 2001
From: Michel Weststrate <mweststrate@gmail.com>
Date: Tue, 20 Oct 2020 21:34:40 +0100
Subject: [PATCH] fix: reconcile if the original value is assigned after
 creating a draft. Fixes #659

---
 __tests__/regressions.js | 81 ++++++++++++++++++++++++++++++++++++++++
 src/core/proxy.ts        | 22 ++++++++---
 2 files changed, 97 insertions(+), 6 deletions(-)

diff --git a/__tests__/regressions.js b/__tests__/regressions.js
index 9ae9138b..d3e63cac 100644
--- a/__tests__/regressions.js
+++ b/__tests__/regressions.js
@@ -158,5 +158,86 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
 			})
 			expect(newData.foo[0].isActive).toBe(true)
 		})
+
+		test("#659 no reconciliation after read", () => {
+			const bar = {}
+			const foo = {bar}
+
+			const next = produce(foo, draft => {
+				draft.bar
+				draft.bar = bar
+			})
+			expect(next).toBe(foo)
+		})
+
+		test("#659 no reconciliation after read - 2", () => {
+			const bar = {}
+			const foo = {bar}
+
+			const next = produce(foo, draft => {
+				const subDraft = draft.bar
+				draft.bar = bar
+				subDraft.x = 3 // this subDraft is not part of the end result, so ignore
+			})
+
+			expect(next).toEqual(foo)
+		})
+
+		test("#659 no reconciliation after read - 3", () => {
+			const bar = {}
+			const foo = {bar}
+
+			const next = produce(foo, draft => {
+				const subDraft = draft.bar
+				subDraft.x = 3 // this subDraft is not part of the end result, so ignore
+				draft.bar = bar
+			})
+			expect(next).toEqual(foo)
+		})
+
+		// Disabled: these are optimizations that would be nice if they
+		// could be detected, but don't change the correctness of the result
+		test.skip("#659 no reconciliation after read - 4", () => {
+			const bar = {}
+			const foo = {bar}
+
+			const next = produce(foo, draft => {
+				const subDraft = draft.bar
+				draft.bar = bar
+				subDraft.x = 3 // this subDraft is not part of the end result, so ignore
+			})
+
+			expect(next).toBe(foo)
+		})
+
+		// Disabled: these are optimizations that would be nice if they
+		// could be detected, but don't change the correctness of the result
+		test.skip("#659 no reconciliation after read - 5", () => {
+			const bar = {}
+			const foo = {bar}
+
+			const next = produce(foo, draft => {
+				const subDraft = draft.bar
+				subDraft.x = 3 // this subDraft is not part of the end result, so ignore
+				draft.bar = bar
+			})
+			expect(next).toBe(foo)
+		})
+
+		test("#659 no reconciliation after read - 6", () => {
+			const bar = {}
+			const foo = {bar}
+
+			const next = produce(foo, draft => {
+				const subDraft = draft.bar
+				subDraft.x = 3 // this subDraft is not part of the end result, so ignore
+				draft.bar = bar
+				draft.bar = subDraft
+			})
+			expect(next).not.toBe(foo)
+			expect(next).toEqual({
+				bar: {x: 3}
+			})
+		})
 	})
 }
diff --git a/src/core/proxy.ts b/src/core/proxy.ts
index 13140a5c..dbccbc20 100644
--- a/src/core/proxy.ts
+++ b/src/core/proxy.ts
@@ -18,6 +18,7 @@ import {
 	ProxyTypeProxyObject,
 	ProxyTypeProxyArray
 } from "../internal"
+import {isDraft} from "../utils/common"
 
 interface ProxyBaseState extends ImmerBaseState {
 	assigned_: {
@@ -129,7 +130,11 @@ export const objectTraps: ProxyHandler<ProxyState> = {
 	ownKeys(state) {
 		return Reflect.ownKeys(latest(state))
 	},
-	set(state, prop: string /* strictly not, but helps TS */, value) {
+	set(
+		state: ProxyObjectState,
+		prop: string /* strictly not, but helps TS */,
+		value
+	) {
 		const desc = getDescriptorFromProto(latest(state), prop)
 		if (desc?.set) {
 			// special case: if this write is captured by a setter, we have
@@ -137,20 +142,25 @@ export const objectTraps: ProxyHandler<ProxyState> = {
 			desc.set.call(state.draft_, value)
 			return true
 		}
-		state.assigned_[prop] = true
 		if (!state.modified_) {
 			// the last check is because we need to be able to distinguish setting a non-existig to undefined (which is a change)
 			// from setting an existing property with value undefined to undefined (which is not a change)
-			if (
-				is(value, peek(latest(state), prop)) &&
-				(value !== undefined || has(state.base_, prop))
-			)
+			const current = peek(latest(state), prop)
+			// special case, if we assigning the original value to a draft, we can ignore the assignment
+			const currentState: ProxyObjectState = current?.[DRAFT_STATE]
+			if (currentState && currentState.base_ === value) {
+				state.copy_![prop] = value
+				state.assigned_[prop] = false
+				return true
+			}
+			if (is(value, current) && (value !== undefined || has(state.base_, prop)))
 				return true
 			prepareCopy(state)
 			markChanged(state)
 		}
 		// @ts-ignore
 		state.copy_![prop] = value
+		state.assigned_[prop] = true
 		return true
 	},
 	deleteProperty(state, prop: string) {