Faster iteration experiment

__tests__/not-strict-copy.ts

@@ -1,37 +1,75 @@
-import {produce, setUseStrictShallowCopy} from "../src/immer"
+import {
+	immerable,
+	produce,
+	setUseStrictShallowCopy,
+	setAutoFreeze,
+	StrictMode
+} from "../src/immer"
 
-describe("setUseStrictShallowCopy(true)", () => {
-	test("keep descriptors", () => {
-		setUseStrictShallowCopy(true)
+describe.each([true, false, "class_only" as const])(
+	"setUseStrictShallowCopy(true)",
+	(strictMode: StrictMode) => {
+		test("keep descriptors, mode: " + strictMode, () => {
+			setUseStrictShallowCopy(strictMode)
 
-		const base: Record<string, unknown> = {}
-		Object.defineProperty(base, "foo", {
-			value: "foo",
-			writable: false,
-			configurable: false
+			const base: Record<string, unknown> = {}
+			Object.defineProperty(base, "foo", {
+				value: "foo",
+				writable: false,
+				configurable: false
+			})
+			const copy = produce(base, (draft: any) => {
+				draft.bar = "bar"
+			})
+			if (strictMode === true) {
+				expect(Object.getOwnPropertyDescriptor(copy, "foo")).toStrictEqual(
+					Object.getOwnPropertyDescriptor(base, "foo")
+				)
+			} else {
+				expect(Object.getOwnPropertyDescriptor(copy, "foo")).toBeUndefined()
+			}
 		})
-		const copy = produce(base, (draft: any) => {
-			draft.bar = "bar"
-		})
-		expect(Object.getOwnPropertyDescriptor(copy, "foo")).toStrictEqual(
-			Object.getOwnPropertyDescriptor(base, "foo")
-		)
-	})
-})
-
-describe("setUseStrictShallowCopy(false)", () => {
-	test("ignore descriptors", () => {
-		setUseStrictShallowCopy(false)
-
-		const base: Record<string, unknown> = {}
-		Object.defineProperty(base, "foo", {
-			value: "foo",
-			writable: false,
-			configurable: false
-		})
-		const copy = produce(base, (draft: any) => {
-			draft.bar = "bar"
+
+		test("keep non-enumerable class descriptors, mode: " + strictMode, () => {
+			setUseStrictShallowCopy(strictMode)
+			setAutoFreeze(false)
+
+			class X {
+				[immerable] = true
+				foo = "foo"
+				bar!: string
+				constructor() {
+					Object.defineProperty(this, "bar", {
+						get() {
+							return this.foo + "bar"
+						},
+						configurable: false,
+						enumerable: false
+					})
+				}
+
+				get baz() {
+					return this.foo + "baz"
+				}
+			}
+
+			const copy = produce(new X(), (draft: any) => {
+				draft.foo = "FOO"
+			})
+
+			const strict = strictMode === true || strictMode === "class_only"
+
+			// descriptors on the prototype are unaffected, so this is still a getter
+			expect(copy.baz).toBe("FOObaz")
+			// descriptors on the instance are found, even when non-enumerable, and read during copy
+			// so new values won't be reflected
+			expect(copy.bar).toBe(strict ? "foobar" : undefined)
+
+			copy.foo = "fluff"
+			// not updated, the own prop became a value
+			expect(copy.bar).toBe(strict ? "foobar" : undefined)
+			// updated, it is still a getter
+			expect(copy.baz).toBe("fluffbaz")
 		})
-		expect(Object.getOwnPropertyDescriptor(copy, "foo")).toBeUndefined()
-	})
-})
+	}
+)

src/core/current.ts

@@ -21,18 +21,22 @@ function currentImpl(value: any): any {
 	if (!isDraftable(value) || isFrozen(value)) return value
 	const state: ImmerState | undefined = value[DRAFT_STATE]
 	let copy: any
+	let strict = true // we might not know, so true just to be safe
 	if (state) {
 		if (!state.modified_) return state.base_
 		// Optimization: avoid generating new drafts during copying
 		state.finalized_ = true
-		copy = shallowCopy(value, state.scope_.immer_.useStrictShallowCopy_)
-	} else {
-		copy = shallowCopy(value, true)
+		strict = state.scope_.immer_.useStrict_(value)
 	}
+	copy = shallowCopy(value, strict)
 	// recurse
-	each(copy, (key, childValue) => {
-		set(copy, key, currentImpl(childValue))
-	})
+	each(
+		copy,
+		(key, childValue) => {
+			set(copy, key, currentImpl(childValue))
+		},
+		strict
+	)
 	if (state) {
 		state.finalized_ = false
 	}

src/core/finalize.ts

@@ -58,8 +58,11 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) {
 	const state: ImmerState = value[DRAFT_STATE]
 	// A plain object, might need freezing, might contain drafts
 	if (!state) {
-		each(value, (key, childValue) =>
-			finalizeProperty(rootScope, state, value, key, childValue, path)
+		each(
+			value,
+			(key, childValue) =>
+				finalizeProperty(rootScope, state, value, key, childValue, path),
+			rootScope.immer_.useStrict_(value)
 		)
 		return value
 	}
@@ -86,8 +89,19 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) {
 			result.clear()
 			isSet = true
 		}
-		each(resultEach, (key, childValue) =>
-			finalizeProperty(rootScope, state, result, key, childValue, path, isSet)
+		each(
+			resultEach,
+			(key, childValue) =>
+				finalizeProperty(
+					rootScope,
+					state,
+					result,
+					key,
+					childValue,
+					path,
+					isSet
+				),
+			state.scope_.immer_.useStrict_(resultEach)
 		)
 		// everything inside is frozen, we can freeze here
 		maybeFreeze(rootScope, result, false)

src/core/immerClass.ts

@@ -23,19 +23,25 @@ import {
 	getCurrentScope,
 	NOTHING,
 	freeze,
-	current
+	current,
+	isPlainObject
 } from "../internal"
 
 interface ProducersFns {
 	produce: IProduce
 	produceWithPatches: IProduceWithPatches
 }
 
+export type StrictMode = boolean | "class_only"
+
 export class Immer implements ProducersFns {
 	autoFreeze_: boolean = true
-	useStrictShallowCopy_: boolean = false
+	useStrictShallowCopy_: StrictMode = false
 
-	constructor(config?: {autoFreeze?: boolean; useStrictShallowCopy?: boolean}) {
+	constructor(config?: {
+		autoFreeze?: boolean
+		useStrictShallowCopy?: StrictMode
+	}) {
 		if (typeof config?.autoFreeze === "boolean")
 			this.setAutoFreeze(config!.autoFreeze)
 		if (typeof config?.useStrictShallowCopy === "boolean")
@@ -163,7 +169,7 @@ export class Immer implements ProducersFns {
 	 *
 	 * By default, immer does not copy the object descriptors such as getter, setter and non-enumrable properties.
 	 */
-	setUseStrictShallowCopy(value: boolean) {
+	setUseStrictShallowCopy(value: StrictMode) {
 		this.useStrictShallowCopy_ = value
 	}
 
@@ -194,6 +200,13 @@ export class Immer implements ProducersFns {
 			applyPatchesImpl(draft, patches)
 		)
 	}
+
+	useStrict_(obj: any): boolean {
+		return (
+			this.useStrictShallowCopy_ === true ||
+			(this.useStrictShallowCopy_ === "class_only" && !isPlainObject(obj))
+		)
+	}
 }
 
 export function createProxy<T extends Objectish>(

src/immer.ts

@@ -18,7 +18,8 @@ export {
 	NOTHING as nothing,
 	DRAFTABLE as immerable,
 	freeze,
-	Objectish
+	Objectish,
+	StrictMode
 } from "./internal"
 
 const immer = new Immer()

src/plugins/patches.ts

@@ -132,21 +132,25 @@ export function enablePatches() {
 		inversePatches: Patch[]
 	) {
 		const {base_, copy_} = state
-		each(state.assigned_!, (key, assignedValue) => {
-			const origValue = get(base_, key)
-			const value = get(copy_!, key)
-			const op = !assignedValue ? REMOVE : has(base_, key) ? REPLACE : ADD
-			if (origValue === value && op === REPLACE) return
-			const path = basePath.concat(key as any)
-			patches.push(op === REMOVE ? {op, path} : {op, path, value})
-			inversePatches.push(
-				op === ADD
-					? {op: REMOVE, path}
-					: op === REMOVE
-					? {op: ADD, path, value: clonePatchValueIfNeeded(origValue)}
-					: {op: REPLACE, path, value: clonePatchValueIfNeeded(origValue)}
-			)
-		})
+		each(
+			state.assigned_!,
+			(key, assignedValue) => {
+				const origValue = get(base_, key)
+				const value = get(copy_!, key)
+				const op = !assignedValue ? REMOVE : has(base_, key) ? REPLACE : ADD
+				if (origValue === value && op === REPLACE) return
+				const path = basePath.concat(key as any)
+				patches.push(op === REMOVE ? {op, path} : {op, path, value})
+				inversePatches.push(
+					op === ADD
+						? {op: REMOVE, path}
+						: op === REMOVE
+						? {op: ADD, path, value: clonePatchValueIfNeeded(origValue)}
+						: {op: REPLACE, path, value: clonePatchValueIfNeeded(origValue)}
+				)
+			},
+			state.scope_.immer_.useStrict_(base_)
+		)
 	}
 
 	function generateSetPatches(

src/utils/common.ts

@@ -8,7 +8,8 @@ import {
 	AnySet,
 	ImmerState,
 	ArchType,
-	die
+	die,
+	StrictMode
 } from "../internal"
 
 export const getPrototypeOf = Object.getPrototypeOf
@@ -67,11 +68,13 @@ export function original(value: Drafted<any>): any {
  */
 export function each<T extends Objectish>(
 	obj: T,
-	iter: (key: string | number, value: any, source: T) => void
+	iter: (key: string | number, value: any, source: T) => void,
+	strict?: boolean
 ): void
-export function each(obj: any, iter: any) {
+export function each(obj: any, iter: any, strict: boolean = true) {
 	if (getArchtype(obj) === ArchType.Object) {
-		Reflect.ownKeys(obj).forEach(key => {
+		const keys = strict ? Reflect.ownKeys(obj) : Object.keys(obj)
+		keys.forEach(key => {
 			iter(key, obj[key], obj)
 		})
 	} else {
@@ -140,7 +143,7 @@ export function latest(state: ImmerState): any {
 }
 
 /*#__PURE__*/
-export function shallowCopy(base: any, strict: boolean) {
+export function shallowCopy(base: any, strict: StrictMode) {
 	if (isMap(base)) {
 		return new Map(base)
 	}
@@ -149,36 +152,41 @@ export function shallowCopy(base: any, strict: boolean) {
 	}
 	if (Array.isArray(base)) return Array.prototype.slice.call(base)
 
-	if (!strict && isPlainObject(base)) {
-		if (!getPrototypeOf(base)) {
-			const obj = Object.create(null)
-			return Object.assign(obj, base)
+	const isPlain = isPlainObject(base)
+
+	if (strict === true || (strict === "class_only" && !isPlain)) {
+		// Perform a strict copy
+		const descriptors = Object.getOwnPropertyDescriptors(base)
+		delete descriptors[DRAFT_STATE as any]
+		let keys = Reflect.ownKeys(descriptors)
+		for (let i = 0; i < keys.length; i++) {
+			const key: any = keys[i]
+			const desc = descriptors[key]
+			if (desc.writable === false) {
+				desc.writable = true
+				desc.configurable = true
+			}
+			// like object.assign, we will read any _own_, get/set accessors. This helps in dealing
+			// with libraries that trap values, like mobx or vue
+			// unlike object.assign, non-enumerables will be copied as well
+			if (desc.get || desc.set)
+				descriptors[key] = {
+					configurable: true,
+					writable: true, // could live with !!desc.set as well here...
+					enumerable: desc.enumerable,
+					value: base[key]
+				}
 		}
-		return {...base}
-	}
-
-	const descriptors = Object.getOwnPropertyDescriptors(base)
-	delete descriptors[DRAFT_STATE as any]
-	let keys = Reflect.ownKeys(descriptors)
-	for (let i = 0; i < keys.length; i++) {
-		const key: any = keys[i]
-		const desc = descriptors[key]
-		if (desc.writable === false) {
-			desc.writable = true
-			desc.configurable = true
+		return Object.create(getPrototypeOf(base), descriptors)
+	} else {
+		// perform a sloppy copy
+		const proto = getPrototypeOf(base)
+		if (proto !== null && isPlain) {
+			return {...base} // assumption: better inner class optimization than the assign below
 		}
-		// like object.assign, we will read any _own_, get/set accessors. This helps in dealing
-		// with libraries that trap values, like mobx or vue
-		// unlike object.assign, non-enumerables will be copied as well
-		if (desc.get || desc.set)
-			descriptors[key] = {
-				configurable: true,
-				writable: true, // could live with !!desc.set as well here...
-				enumerable: desc.enumerable,
-				value: base[key]
-			}
+		const obj = Object.create(proto)
+		return Object.assign(obj, base)
 	}
-	return Object.create(getPrototypeOf(base), descriptors)
 }
 
 /**