Skip to content

Commit e1996ce

Browse files
authored
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`
1 parent 40aa814 commit e1996ce

File tree

6 files changed

+119
-30
lines changed

6 files changed

+119
-30
lines changed

perf-testing/immutability-benchmarks.mjs

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -59,8 +59,8 @@ const MAX = 1
5959

6060
const BENCHMARK_CONFIG = {
6161
iterations: 1,
62-
arraySize: 10000,
63-
nestedArraySize: 100,
62+
arraySize: 100,
63+
nestedArraySize: 10,
6464
multiUpdateCount: 5,
6565
reuseStateIterations: 10
6666
}

src/core/current.ts

Lines changed: 9 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -21,18 +21,24 @@ function currentImpl(value: any): any {
2121
if (!isDraftable(value) || isFrozen(value)) return value
2222
const state: ImmerState | undefined = value[DRAFT_STATE]
2323
let copy: any
24+
let strict = true // Default to strict for compatibility
2425
if (state) {
2526
if (!state.modified_) return state.base_
2627
// Optimization: avoid generating new drafts during copying
2728
state.finalized_ = true
2829
copy = shallowCopy(value, state.scope_.immer_.useStrictShallowCopy_)
30+
strict = state.scope_.immer_.shouldUseStrictIteration()
2931
} else {
3032
copy = shallowCopy(value, true)
3133
}
3234
// recurse
33-
each(copy, (key, childValue) => {
34-
set(copy, key, currentImpl(childValue))
35-
})
35+
each(
36+
copy,
37+
(key, childValue) => {
38+
set(copy, key, currentImpl(childValue))
39+
},
40+
strict
41+
)
3642
if (state) {
3743
state.finalized_ = false
3844
}

src/core/finalize.ts

Lines changed: 42 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -56,11 +56,16 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) {
5656
// Don't recurse in tho recursive data structures
5757
if (isFrozen(value)) return value
5858

59+
const useStrictIteration = rootScope.immer_.shouldUseStrictIteration()
60+
5961
const state: ImmerState = value[DRAFT_STATE]
6062
// A plain object, might need freezing, might contain drafts
6163
if (!state) {
62-
each(value, (key, childValue) =>
63-
finalizeProperty(rootScope, state, value, key, childValue, path)
64+
each(
65+
value,
66+
(key, childValue) =>
67+
finalizeProperty(rootScope, state, value, key, childValue, path),
68+
useStrictIteration
6469
)
6570
return value
6671
}
@@ -87,8 +92,19 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) {
8792
result.clear()
8893
isSet = true
8994
}
90-
each(resultEach, (key, childValue) =>
91-
finalizeProperty(rootScope, state, result, key, childValue, path, isSet)
95+
each(
96+
resultEach,
97+
(key, childValue) =>
98+
finalizeProperty(
99+
rootScope,
100+
state,
101+
result,
102+
key,
103+
childValue,
104+
path,
105+
isSet
106+
),
107+
useStrictIteration
92108
)
93109
// everything inside is frozen, we can freeze here
94110
maybeFreeze(rootScope, result, false)
@@ -114,6 +130,18 @@ function finalizeProperty(
114130
rootPath?: PatchPath,
115131
targetIsSet?: boolean
116132
) {
133+
if (childValue == null) {
134+
return
135+
}
136+
137+
if (typeof childValue !== "object" && !targetIsSet) {
138+
return
139+
}
140+
const childIsFrozen = isFrozen(childValue)
141+
if (childIsFrozen && !targetIsSet) {
142+
return
143+
}
144+
117145
if (process.env.NODE_ENV !== "production" && childValue === targetObject)
118146
die(5)
119147
if (isDraft(childValue)) {
@@ -136,7 +164,7 @@ function finalizeProperty(
136164
targetObject.add(childValue)
137165
}
138166
// Search new objects for unfinalized drafts. Frozen objects should never contain drafts.
139-
if (isDraftable(childValue) && !isFrozen(childValue)) {
167+
if (isDraftable(childValue) && !childIsFrozen) {
140168
if (!rootScope.immer_.autoFreeze_ && rootScope.unfinalizedDrafts_ < 1) {
141169
// optimization: if an object is not a draft, and we don't have to
142170
// deepfreeze everything, and we are sure that no drafts are left in the remaining object
@@ -145,6 +173,15 @@ function finalizeProperty(
145173
// See add-data.js perf test
146174
return
147175
}
176+
if (
177+
parentState &&
178+
parentState.base_ &&
179+
parentState.base_[prop] === childValue &&
180+
childIsFrozen
181+
) {
182+
// Object is unchanged from base - no need to process further
183+
return
184+
}
148185
finalize(rootScope, childValue)
149186
// Immer deep freezes plain objects, so if there is no parent state, we freeze as well
150187
// Per #590, we never freeze symbolic properties. Just to make sure don't accidentally interfere

src/core/immerClass.ts

Lines changed: 19 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -31,20 +31,24 @@ interface ProducersFns {
3131
produceWithPatches: IProduceWithPatches
3232
}
3333

34-
export type StrictMode = boolean | "class_only";
34+
export type StrictMode = boolean | "class_only"
3535

3636
export class Immer implements ProducersFns {
3737
autoFreeze_: boolean = true
3838
useStrictShallowCopy_: StrictMode = false
39+
useStrictIteration_: boolean = true
3940

4041
constructor(config?: {
4142
autoFreeze?: boolean
4243
useStrictShallowCopy?: StrictMode
44+
useStrictIteration?: boolean
4345
}) {
4446
if (typeof config?.autoFreeze === "boolean")
4547
this.setAutoFreeze(config!.autoFreeze)
4648
if (typeof config?.useStrictShallowCopy === "boolean")
4749
this.setUseStrictShallowCopy(config!.useStrictShallowCopy)
50+
if (typeof config?.useStrictIteration === "boolean")
51+
this.setUseStrictIteration(config!.useStrictIteration)
4852
}
4953

5054
/**
@@ -172,6 +176,20 @@ export class Immer implements ProducersFns {
172176
this.useStrictShallowCopy_ = value
173177
}
174178

179+
/**
180+
* Pass false to use faster iteration that skips non-enumerable properties
181+
* but still handles symbols for compatibility.
182+
*
183+
* By default, strict iteration is enabled (includes all own properties).
184+
*/
185+
setUseStrictIteration(value: boolean) {
186+
this.useStrictIteration_ = value
187+
}
188+
189+
shouldUseStrictIteration(): boolean {
190+
return this.useStrictIteration_
191+
}
192+
175193
applyPatches<T extends Objectish>(base: T, patches: readonly Patch[]): T {
176194
// If a patch replaces the entire state, take that replacement as base
177195
// before applying patches

src/immer.ts

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -71,6 +71,16 @@ export const setUseStrictShallowCopy = /* @__PURE__ */ immer.setUseStrictShallow
7171
immer
7272
)
7373

74+
/**
75+
* Pass false to use loose iteration that only processes enumerable string properties.
76+
* This skips symbols and non-enumerable properties for maximum performance.
77+
*
78+
* By default, strict iteration is enabled (includes all own properties).
79+
*/
80+
export const setUseStrictIteration = /* @__PURE__ */ immer.setUseStrictIteration.bind(
81+
immer
82+
)
83+
7484
/**
7585
* Apply an array of Immer patches to the first argument.
7686
*

src/utils/common.ts

Lines changed: 37 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -35,22 +35,26 @@ export function isDraftable(value: any): boolean {
3535
}
3636

3737
const objectCtorString = Object.prototype.constructor.toString()
38+
const cachedCtorStrings = new WeakMap()
3839
/*#__PURE__*/
3940
export function isPlainObject(value: any): boolean {
4041
if (!value || typeof value !== "object") return false
41-
const proto = getPrototypeOf(value)
42-
if (proto === null) {
43-
return true
44-
}
42+
const proto = Object.getPrototypeOf(value)
43+
if (proto === null || proto === Object.prototype) return true
44+
4545
const Ctor =
4646
Object.hasOwnProperty.call(proto, "constructor") && proto.constructor
47-
4847
if (Ctor === Object) return true
4948

50-
return (
51-
typeof Ctor == "function" &&
52-
Function.toString.call(Ctor) === objectCtorString
53-
)
49+
if (typeof Ctor !== "function") return false
50+
51+
let ctorString = cachedCtorStrings.get(Ctor)
52+
if (ctorString === undefined) {
53+
ctorString = Function.toString.call(Ctor)
54+
cachedCtorStrings.set(Ctor, ctorString)
55+
}
56+
57+
return ctorString === objectCtorString
5458
}
5559

5660
/** Get the underlying object that is represented by the given draft */
@@ -64,15 +68,23 @@ export function original(value: Drafted<any>): any {
6468
/**
6569
* Each iterates a map, set or array.
6670
* Or, if any other kind of object, all of its own properties.
67-
* Regardless whether they are enumerable or symbols
71+
*
72+
* @param obj The object to iterate over
73+
* @param iter The iterator function
74+
* @param strict When true (default), includes symbols and non-enumerable properties.
75+
* When false, uses looseiteration over only enumerable string properties.
6876
*/
6977
export function each<T extends Objectish>(
7078
obj: T,
71-
iter: (key: string | number, value: any, source: T) => void
79+
iter: (key: string | number, value: any, source: T) => void,
80+
strict?: boolean
7281
): void
73-
export function each(obj: any, iter: any) {
82+
export function each(obj: any, iter: any, strict: boolean = true) {
7483
if (getArchtype(obj) === ArchType.Object) {
75-
Reflect.ownKeys(obj).forEach(key => {
84+
// If strict, we do a full iteration including symbols and non-enumerable properties
85+
// Otherwise, we only iterate enumerable string properties for performance
86+
const keys = strict ? Reflect.ownKeys(obj) : Object.keys(obj)
87+
keys.forEach(key => {
7688
iter(key, obj[key], obj)
7789
})
7890
} else {
@@ -198,12 +210,12 @@ export function freeze<T>(obj: T, deep?: boolean): T
198210
export function freeze<T>(obj: any, deep: boolean = false): T {
199211
if (isFrozen(obj) || isDraft(obj) || !isDraftable(obj)) return obj
200212
if (getArchtype(obj) > 1 /* Map or Set */) {
201-
Object.defineProperties(obj, {
202-
set: {value: dontMutateFrozenCollections as any},
203-
add: {value: dontMutateFrozenCollections as any},
204-
clear: {value: dontMutateFrozenCollections as any},
205-
delete: {value: dontMutateFrozenCollections as any}
206-
})
213+
Object.defineProperties(obj, {
214+
set: dontMutateMethodOverride,
215+
add: dontMutateMethodOverride,
216+
clear: dontMutateMethodOverride,
217+
delete: dontMutateMethodOverride
218+
})
207219
}
208220
Object.freeze(obj)
209221
if (deep)
@@ -217,6 +229,12 @@ function dontMutateFrozenCollections() {
217229
die(2)
218230
}
219231

232+
const dontMutateMethodOverride = {
233+
value: dontMutateFrozenCollections
234+
}
235+
220236
export function isFrozen(obj: any): boolean {
237+
// Fast path: primitives and null/undefined are always "frozen"
238+
if (obj === null || typeof obj !== "object") return true
221239
return Object.isFrozen(obj)
222240
}

0 commit comments

Comments
 (0)