From 2697430694dd4f68dcc48b1029d15624cd0b1c80 Mon Sep 17 00:00:00 2001
From: Michel Weststrate <mweststrate@fb.com>
Date: Fri, 24 Jul 2020 17:07:43 +0100
Subject: [PATCH] fix: Trigger setters with the correct context, fixes #604

---
 __tests__/regressions.js | 82 ++++++++++++++++++++++++++++++++++++++++
 src/core/proxy.ts        | 25 ++++++++++--
 2 files changed, 104 insertions(+), 3 deletions(-)
 create mode 100644 __tests__/regressions.js

diff --git a/__tests__/regressions.js b/__tests__/regressions.js
new file mode 100644
index 00000000..9c5bae33
--- /dev/null
+++ b/__tests__/regressions.js
@@ -0,0 +1,82 @@
+"use strict"
+import {
+	Immer,
+	nothing,
+	original,
+	isDraft,
+	immerable,
+	enableAllPlugins
+} from "../src/immer"
+
+enableAllPlugins()
+
+runBaseTest("proxy (no freeze)", true, false)
+runBaseTest("proxy (autofreeze)", true, true)
+runBaseTest("es5 (no freeze)", false, false)
+runBaseTest("es5 (autofreeze)", false, true)
+
+function runBaseTest(name, useProxies, autoFreeze, useListener) {
+	const listener = useListener ? function() {} : undefined
+	const {produce, produceWithPatches} = createPatchedImmer({
+		useProxies,
+		autoFreeze
+	})
+
+	// When `useListener` is true, append a function to the arguments of every
+	// uncurried `produce` call in every test. This makes tests easier to read.
+	function createPatchedImmer(options) {
+		const immer = new Immer(options)
+
+		const {produce} = immer
+		immer.produce = function(...args) {
+			return typeof args[1] === "function" && args.length < 3
+				? produce(...args, listener)
+				: produce(...args)
+		}
+
+		return immer
+	}
+
+	describe(`regressions ${name}`, () => {
+		test("#604 freeze inside class", () => {
+			class Thing {
+				[immerable] = true
+
+				constructor({x}) {
+					this._data = {x}
+				}
+
+				get x() {
+					return this._data.x
+				}
+
+				set x(x) {
+					this._data.x = x
+				}
+			}
+
+			let i = 1
+			let item = new Thing({x: i})
+			let item0 = item
+
+			const bump = () => {
+				item = produce(item, draft => {
+					// uncomment this to make things work
+					//draft._data
+					draft.x = ++i
+				})
+			}
+
+			bump()
+			bump()
+
+			expect(i).toBe(3)
+			expect(item._data).toEqual({
+				x: 3
+			})
+			expect(item0._data).toEqual({
+				x: 1
+			})
+		})
+	})
+}
diff --git a/src/core/proxy.ts b/src/core/proxy.ts
index 867c779c..94b89f85 100644
--- a/src/core/proxy.ts
+++ b/src/core/proxy.ts
@@ -130,6 +130,13 @@ export const objectTraps: ProxyHandler<ProxyState> = {
 		return Reflect.ownKeys(latest(state))
 	},
 	set(state, 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
+			// to trigger it with the correct context
+			desc.set.call(state.draft_, value)
+			return true
+		}
 		state.assigned_[prop] = true
 		if (!state.modified_) {
 			if (is(value, peek(latest(state), prop)) && value !== undefined)
@@ -208,14 +215,26 @@ function peek(draft: Drafted, prop: PropertyKey) {
 }
 
 function readPropFromProto(state: ImmerState, source: any, prop: PropertyKey) {
+	const desc = getDescriptorFromProto(source, prop)
+	return desc
+		? `value` in desc
+			? desc.value
+			: // This is a very special case, if the prop is a getter defined by the
+			  // prototype, we should invoke it with the draft as context!
+			  desc.get?.call(state.draft_)
+		: undefined
+}
+
+function getDescriptorFromProto(
+	source: any,
+	prop: PropertyKey
+): PropertyDescriptor | undefined {
 	// 'in' checks proto!
 	if (!(prop in source)) return undefined
 	let proto = Object.getPrototypeOf(source)
 	while (proto) {
 		const desc = Object.getOwnPropertyDescriptor(proto, prop)
-		// This is a very special case, if the prop is a getter defined by the
-		// prototype, we should invoke it with the draft as context!
-		if (desc) return `value` in desc ? desc.value : desc.get?.call(state.draft_)
+		if (desc) return desc
 		proto = Object.getPrototypeOf(proto)
 	}
 	return undefined