feat: Optimize Immer performance where possible, introduce setUseStrictIteration (#1164)

* Use WeakMap caching implementation of isPlainObject

* Add some early returns to `finalizeProperty`

* Add `strictIteration` option

* Add non-strict iteration handling

* Use strict iteration option

* Switch back to default strict iteration

* Fix strict iteration checks

* Shorten benchmark array sizes for faster results

* Dedupe Map/Set method overrides

* Removed old isPlainObject impl

* Add early bailout to `isFrozen`

Author: markerikson

perf-testing/immutability-benchmarks.mjs

@@ -59,8 +59,8 @@ const MAX = 1
 
 const BENCHMARK_CONFIG = {
 	iterations: 1,
-	arraySize: 10000,
-	nestedArraySize: 100,
+	arraySize: 100,
+	nestedArraySize: 10,
 	multiUpdateCount: 5,
 	reuseStateIterations: 10
 }

src/core/current.ts

@@ -21,18 +21,24 @@ 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 // Default to strict for compatibility
 	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_)
+		strict = state.scope_.immer_.shouldUseStrictIteration()
 	} else {
 		copy = shallowCopy(value, true)
 	}
 	// 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

@@ -56,11 +56,16 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) {
 	// Don't recurse in tho recursive data structures
 	if (isFrozen(value)) return value
 
+	const useStrictIteration = rootScope.immer_.shouldUseStrictIteration()
+
 	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),
+			useStrictIteration
 		)
 		return value
 	}
@@ -87,8 +92,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
+				),
+			useStrictIteration
 		)
 		// everything inside is frozen, we can freeze here
 		maybeFreeze(rootScope, result, false)
@@ -114,6 +130,18 @@ function finalizeProperty(
 	rootPath?: PatchPath,
 	targetIsSet?: boolean
 ) {
+	if (childValue == null) {
+		return
+	}
+
+	if (typeof childValue !== "object" && !targetIsSet) {
+		return
+	}
+	const childIsFrozen = isFrozen(childValue)
+	if (childIsFrozen && !targetIsSet) {
+		return
+	}
+
 	if (process.env.NODE_ENV !== "production" && childValue === targetObject)
 		die(5)
 	if (isDraft(childValue)) {
@@ -136,7 +164,7 @@ function finalizeProperty(
 		targetObject.add(childValue)
 	}
 	// Search new objects for unfinalized drafts. Frozen objects should never contain drafts.
-	if (isDraftable(childValue) && !isFrozen(childValue)) {
+	if (isDraftable(childValue) && !childIsFrozen) {
 		if (!rootScope.immer_.autoFreeze_ && rootScope.unfinalizedDrafts_ < 1) {
 			// optimization: if an object is not a draft, and we don't have to
 			// deepfreeze everything, and we are sure that no drafts are left in the remaining object
@@ -145,6 +173,15 @@ function finalizeProperty(
 			// See add-data.js perf test
 			return
 		}
+		if (
+			parentState &&
+			parentState.base_ &&
+			parentState.base_[prop] === childValue &&
+			childIsFrozen
+		) {
+			// Object is unchanged from base - no need to process further
+			return
+		}
 		finalize(rootScope, childValue)
 		// Immer deep freezes plain objects, so if there is no parent state, we freeze as well
 		// Per #590, we never freeze symbolic properties. Just to make sure don't accidentally interfere

src/core/immerClass.ts

@@ -31,20 +31,24 @@ interface ProducersFns {
 	produceWithPatches: IProduceWithPatches
 }
 
-export type StrictMode = boolean | "class_only";
+export type StrictMode = boolean | "class_only"
 
 export class Immer implements ProducersFns {
 	autoFreeze_: boolean = true
 	useStrictShallowCopy_: StrictMode = false
+	useStrictIteration_: boolean = true
 
 	constructor(config?: {
 		autoFreeze?: boolean
 		useStrictShallowCopy?: StrictMode
+		useStrictIteration?: boolean
 	}) {
 		if (typeof config?.autoFreeze === "boolean")
 			this.setAutoFreeze(config!.autoFreeze)
 		if (typeof config?.useStrictShallowCopy === "boolean")
 			this.setUseStrictShallowCopy(config!.useStrictShallowCopy)
+		if (typeof config?.useStrictIteration === "boolean")
+			this.setUseStrictIteration(config!.useStrictIteration)
 	}
 
 	/**
@@ -172,6 +176,20 @@ export class Immer implements ProducersFns {
 		this.useStrictShallowCopy_ = value
 	}
 
+	/**
+	 * Pass false to use faster iteration that skips non-enumerable properties
+	 * but still handles symbols for compatibility.
+	 *
+	 * By default, strict iteration is enabled (includes all own properties).
+	 */
+	setUseStrictIteration(value: boolean) {
+		this.useStrictIteration_ = value
+	}
+
+	shouldUseStrictIteration(): boolean {
+		return this.useStrictIteration_
+	}
+
 	applyPatches<T extends Objectish>(base: T, patches: readonly Patch[]): T {
 		// If a patch replaces the entire state, take that replacement as base
 		// before applying patches

src/immer.ts

@@ -71,6 +71,16 @@ export const setUseStrictShallowCopy = /* @__PURE__ */ immer.setUseStrictShallow
 	immer
 )
 
+/**
+ * Pass false to use loose iteration that only processes enumerable string properties.
+ * This skips symbols and non-enumerable properties for maximum performance.
+ *
+ * By default, strict iteration is enabled (includes all own properties).
+ */
+export const setUseStrictIteration = /* @__PURE__ */ immer.setUseStrictIteration.bind(
+	immer
+)
+
 /**
  * Apply an array of Immer patches to the first argument.
  *

src/utils/common.ts

@@ -35,22 +35,26 @@ export function isDraftable(value: any): boolean {
 }
 
 const objectCtorString = Object.prototype.constructor.toString()
+const cachedCtorStrings = new WeakMap()
 /*#__PURE__*/
 export function isPlainObject(value: any): boolean {
 	if (!value || typeof value !== "object") return false
-	const proto = getPrototypeOf(value)
-	if (proto === null) {
-		return true
-	}
+	const proto = Object.getPrototypeOf(value)
+	if (proto === null || proto === Object.prototype) return true
+
 	const Ctor =
 		Object.hasOwnProperty.call(proto, "constructor") && proto.constructor
-
 	if (Ctor === Object) return true
 
-	return (
-		typeof Ctor == "function" &&
-		Function.toString.call(Ctor) === objectCtorString
-	)
+	if (typeof Ctor !== "function") return false
+
+	let ctorString = cachedCtorStrings.get(Ctor)
+	if (ctorString === undefined) {
+		ctorString = Function.toString.call(Ctor)
+		cachedCtorStrings.set(Ctor, ctorString)
+	}
+
+	return ctorString === objectCtorString
 }
 
 /** Get the underlying object that is represented by the given draft */
@@ -64,15 +68,23 @@ export function original(value: Drafted<any>): any {
 /**
  * Each iterates a map, set or array.
  * Or, if any other kind of object, all of its own properties.
- * Regardless whether they are enumerable or symbols
+ *
+ * @param obj The object to iterate over
+ * @param iter The iterator function
+ * @param strict When true (default), includes symbols and non-enumerable properties.
+ *               When false, uses looseiteration over only enumerable string properties.
  */
 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 => {
+		// If strict, we do a full iteration including symbols and non-enumerable properties
+		// Otherwise, we only iterate enumerable string properties for performance
+		const keys = strict ? Reflect.ownKeys(obj) : Object.keys(obj)
+		keys.forEach(key => {
 			iter(key, obj[key], obj)
 		})
 	} else {
@@ -198,12 +210,12 @@ export function freeze<T>(obj: T, deep?: boolean): T
 export function freeze<T>(obj: any, deep: boolean = false): T {
 	if (isFrozen(obj) || isDraft(obj) || !isDraftable(obj)) return obj
 	if (getArchtype(obj) > 1 /* Map or Set */) {
-		 Object.defineProperties(obj, {
-                        set: {value: dontMutateFrozenCollections as any},
-                        add: {value: dontMutateFrozenCollections as any},
-                        clear: {value: dontMutateFrozenCollections as any},
-                        delete: {value: dontMutateFrozenCollections as any}
-                })
+		Object.defineProperties(obj, {
+			set: dontMutateMethodOverride,
+			add: dontMutateMethodOverride,
+			clear: dontMutateMethodOverride,
+			delete: dontMutateMethodOverride
+		})
 	}
 	Object.freeze(obj)
 	if (deep)
@@ -217,6 +229,12 @@ function dontMutateFrozenCollections() {
 	die(2)
 }
 
+const dontMutateMethodOverride = {
+	value: dontMutateFrozenCollections
+}
+
 export function isFrozen(obj: any): boolean {
+	// Fast path: primitives and null/undefined are always "frozen"
+	if (obj === null || typeof obj !== "object") return true
 	return Object.isFrozen(obj)
 }