From acc5e547a197ea7ff11bdb0550973d5676f7f034 Mon Sep 17 00:00:00 2001
From: Michel Weststrate <mweststrate@fb.com>
Date: Tue, 4 Feb 2020 21:22:41 +0000
Subject: [PATCH] fix: Fixed issue where drafts from nested producers were
 frozen. Fixes #522, #524

---
 __tests__/base.js | 94 +++++++++++++++++++++++++++++++++++++++++++++++
 src/finalize.ts   |  6 +--
 2 files changed, 97 insertions(+), 3 deletions(-)

diff --git a/__tests__/base.js b/__tests__/base.js
index 3d4151ca..20e9a989 100644
--- a/__tests__/base.js
+++ b/__tests__/base.js
@@ -1773,6 +1773,100 @@ function runBaseTest(name, useProxies, autoFreeze, useListener) {
 		}
 	})
 
+	it(`works with spread #524 - ${name}`, () => {
+		const state = {
+			level1: {
+				level2: {
+					level3: "data"
+				}
+			}
+		}
+
+		const nextState = produce(state, draft => {
+			return {...draft}
+		})
+		const nextState1 = produce(state, draft => {
+			const newLevel1 = produce(draft.level1, level1Draft => {
+				return {...level1Draft}
+			})
+			draft.level1 = newLevel1
+		})
+
+		const nextState2 = produce(state, draft => {
+			const newLevel1 = produce(draft.level1, level1Draft => {
+				return {...level1Draft}
+			})
+			return {
+				level1: newLevel1
+			}
+		})
+
+		const nextState3 = produce(state, draft => {
+			const newLevel1 = produce(draft.level1, level1Draft => {
+				return Object.assign({}, level1Draft)
+			})
+			return Object.assign(draft, {
+				level1: newLevel1
+			})
+		})
+
+		const expected = {level1: {level2: {level3: "data"}}}
+		expect(nextState1).toEqual(expected)
+		expect(nextState2).toEqual(expected)
+		expect(nextState3).toEqual(expected)
+	})
+
+	it(`Something with nested producers #522 ${name}`, () => {
+		function initialStateFactory() {
+			return {
+				substate: {
+					array: [
+						{id: "id1", value: 0},
+						{id: "id2", value: 0}
+					]
+				},
+				array: [
+					{id: "id1", value: 0},
+					{id: "id2", value: 0}
+				]
+			}
+		}
+
+		const globalProducer = produce(draft => {
+			draft.substate = subProducer(draft.substate)
+			draft.array = arrayProducer(draft.array)
+		}, initialStateFactory())
+
+		const subProducer = produce(draftSubState => {
+			draftSubState.array = arrayProducer(draftSubState.array)
+		})
+
+		const arrayProducer = produce(draftArray => {
+			draftArray[0].value += 5
+		})
+
+		{
+			const state = globalProducer(undefined)
+			expect(state.array[0].value).toBe(5)
+			expect(state.array.length).toBe(2)
+			expect(state.array[1]).toMatchObject({
+				id: "id2",
+				value: 0
+			})
+		}
+
+		{
+			const state = globalProducer(undefined)
+			expect(state.substate.array[0].value).toBe(5)
+			expect(state.substate.array.length).toBe(2)
+			expect(state.substate.array[1]).toMatchObject({
+				id: "id2",
+				value: 0
+			})
+			expect(state.substate.array).toMatchObject(state.array)
+		}
+	})
+
 	describe(`isDraft - ${name}`, () => {
 		it("returns true for object drafts", () => {
 			produce({}, state => {
diff --git a/src/finalize.ts b/src/finalize.ts
index e57b6235..898b05e0 100644
--- a/src/finalize.ts
+++ b/src/finalize.ts
@@ -33,7 +33,7 @@ export function processResult(immer: Immer, result: any, scope: ImmerScope) {
 		if (isDraftable(result)) {
 			// Finalize the result in case it contains (or is) a subset of the draft.
 			result = finalize(immer, result, scope)
-			maybeFreeze(immer, result)
+			if (!scope.parent) maybeFreeze(immer, result)
 		}
 		if (scope.patches) {
 			scope.patches.push({
@@ -179,7 +179,7 @@ function finalizeProperty(
 	// Search new objects for unfinalized drafts. Frozen objects should never contain drafts.
 	// TODO: the recursion over here looks weird, shouldn't non-draft stuff have it's own recursion?
 	// especially the passing on of root and rootState doesn't make sense...
-	else if (isDraftable(childValue) && !Object.isFrozen(childValue)) {
+	else if (isDraftable(childValue)) {
 		each(childValue, (key, grandChild) =>
 			finalizeProperty(
 				immer,
@@ -192,7 +192,7 @@ function finalizeProperty(
 				rootPath
 			)
 		)
-		maybeFreeze(immer, childValue)
+		if (!scope.parent) maybeFreeze(immer, childValue)
 	}
 
 	if (isDraftProp && immer.onAssign && !isSetMember) {