diff --git a/src/core/finalize.ts b/src/core/finalize.ts index 7c446316..a4fab4cd 100644 --- a/src/core/finalize.ts +++ b/src/core/finalize.ts @@ -15,10 +15,13 @@ import { getPlugin, die, revokeScope, - isFrozen + isFrozen, + type Objectish, + type Drafted, + prepareCopy } from "../internal" -export function processResult(result: any, scope: ImmerScope) { +export function processResult(result: any, scope: ImmerScope, existingStateMap?: WeakMap, existingFinalizationMap?: WeakMap) { scope.unfinalizedDrafts_ = scope.drafts_.length const baseDraft = scope.drafts_![0] const isReplaced = result !== undefined && result !== baseDraft @@ -29,8 +32,8 @@ export function processResult(result: any, scope: ImmerScope) { } if (isDraftable(result)) { // Finalize the result in case it contains (or is) a subset of the draft. - result = finalize(scope, result) - if (!scope.parent_) maybeFreeze(scope, result) + result = finalize(scope, result, undefined, existingStateMap, existingFinalizationMap) + if (!scope.parent_) maybeFreeze(scope, result, false) } if (scope.patches_) { getPlugin("Patches").generateReplacementPatches_( @@ -42,7 +45,7 @@ export function processResult(result: any, scope: ImmerScope) { } } else { // Finalize the base draft. - result = finalize(scope, baseDraft, []) + result = finalize(scope, baseDraft, [], existingStateMap, existingFinalizationMap) } revokeScope(scope) if (scope.patches_) { @@ -51,17 +54,19 @@ export function processResult(result: any, scope: ImmerScope) { return result !== NOTHING ? result : undefined } -function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) { +function finalize(rootScope: ImmerScope, value: any, path?: PatchPath, existingStateMap?: WeakMap, existingFinalizationMap?: WeakMap): any { // Don't recurse in tho recursive data structures if (isFrozen(value)) return value - const state: ImmerState = value[DRAFT_STATE] + let state: ImmerState = value[DRAFT_STATE] + // A plain object, might need freezing, might contain drafts if (!state) { + existingFinalizationMap?.set(value, value) each( value, (key, childValue) => - finalizeProperty(rootScope, state, value, key, childValue, path), + finalizeProperty(rootScope, state, value, key, childValue, path, undefined, existingStateMap, existingFinalizationMap), true // See #590, don't recurse into non-enumerable of non drafted objects ) return value @@ -70,11 +75,13 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) { if (state.scope_ !== rootScope) return value // Unmodified draft, return the (frozen) original if (!state.modified_) { - maybeFreeze(rootScope, state.base_, true) + maybeFreeze(rootScope, state.copy_ ?? state.base_, true); return state.base_ } // Not finalized yet, let's do that now if (!state.finalized_) { + existingFinalizationMap?.set(state.base_, state.copy_) + state.finalized_ = true state.scope_.unfinalizedDrafts_-- const result = state.copy_ @@ -90,7 +97,7 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) { isSet = true } each(resultEach, (key, childValue) => - finalizeProperty(rootScope, state, result, key, childValue, path, isSet) + finalizeProperty(rootScope, state, result, key, childValue, path, isSet, existingStateMap, existingFinalizationMap) ) // everything inside is frozen, we can freeze here maybeFreeze(rootScope, result, false) @@ -104,6 +111,7 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) { ) } } + return state.copy_ } @@ -114,10 +122,23 @@ function finalizeProperty( prop: string | number, childValue: any, rootPath?: PatchPath, - targetIsSet?: boolean + targetIsSet?: boolean, + existingStateMap?: WeakMap, + existingFinalizationMap?: WeakMap, ) { if (process.env.NODE_ENV !== "production" && childValue === targetObject) die(5) + + if (!isDraft(childValue) && isDraftable(childValue)) { + const existingState = existingStateMap?.get(childValue) + if (existingState) { + childValue = existingState.draft_ + } else { + const existingFinalization = existingFinalizationMap?.get(childValue) + if (existingFinalization) return set(targetObject, prop, existingFinalization) + } + } + if (isDraft(childValue)) { const path = rootPath && @@ -127,7 +148,7 @@ function finalizeProperty( ? rootPath!.concat(prop) : undefined // Drafts owned by `scope` are finalized here. - const res = finalize(rootScope, childValue, path) + const res = finalize(rootScope, childValue, path, existingStateMap, existingFinalizationMap) set(targetObject, prop, res) // Drafts from another scope must prevented to be frozen // if we got a draft back from finalize, we're in a nested produce and shouldn't freeze @@ -137,6 +158,7 @@ function finalizeProperty( } else if (targetIsSet) { targetObject.add(childValue) } + // Search new objects for unfinalized drafts. Frozen objects should never contain drafts. if (isDraftable(childValue) && !isFrozen(childValue)) { if (!rootScope.immer_.autoFreeze_ && rootScope.unfinalizedDrafts_ < 1) { @@ -147,10 +169,10 @@ function finalizeProperty( // See add-data.js perf test return } - finalize(rootScope, childValue) + finalize(rootScope, childValue, undefined, existingStateMap, existingFinalizationMap) // immer deep freezes plain objects, so if there is no parent state, we freeze as well if (!parentState || !parentState.scope_.parent_) - maybeFreeze(rootScope, childValue) + maybeFreeze(rootScope, childValue, false); } } diff --git a/src/core/immerClass.ts b/src/core/immerClass.ts index 6c673e0a..c770e2f5 100644 --- a/src/core/immerClass.ts +++ b/src/core/immerClass.ts @@ -34,12 +34,19 @@ interface ProducersFns { export class Immer implements ProducersFns { autoFreeze_: boolean = true useStrictShallowCopy_: boolean = false + allowMultiRefs_: boolean = false - constructor(config?: {autoFreeze?: boolean; useStrictShallowCopy?: boolean}) { + constructor(config?: { + autoFreeze?: boolean + useStrictShallowCopy?: boolean + allowMultiRefs: boolean + }) { if (typeof config?.autoFreeze === "boolean") this.setAutoFreeze(config!.autoFreeze) if (typeof config?.useStrictShallowCopy === "boolean") this.setUseStrictShallowCopy(config!.useStrictShallowCopy) + if (typeof config?.allowMultiRefs === "boolean") + this.setAllowMultiRefs(config!.allowMultiRefs) } /** @@ -86,7 +93,8 @@ export class Immer implements ProducersFns { // Only plain objects, arrays, and "immerable classes" are drafted. if (isDraftable(base)) { const scope = enterScope(this) - const proxy = createProxy(base, undefined) + const stateMap = this.allowMultiRefs_ ? new Map() : undefined + const proxy = createProxy(base, undefined, stateMap) let hasError = true try { result = recipe(proxy) @@ -97,7 +105,7 @@ export class Immer implements ProducersFns { else leaveScope(scope) } usePatchesInScope(scope, patchListener) - return processResult(result, scope) + return processResult(result, scope, stateMap, this.allowMultiRefs_ ? new WeakMap() : undefined) } else if (!base || typeof base !== "object") { result = recipe(base) if (result === undefined) result = base @@ -132,7 +140,7 @@ export class Immer implements ProducersFns { if (!isDraftable(base)) die(8) if (isDraft(base)) base = current(base) const scope = enterScope(this) - const proxy = createProxy(base, undefined) + const proxy = createProxy(base, undefined, this.allowMultiRefs_ ? new WeakMap() : undefined) proxy[DRAFT_STATE].isManual_ = true leaveScope(scope) return proxy as any @@ -144,9 +152,11 @@ export class Immer implements ProducersFns { ): D extends Draft ? T : never { const state: ImmerState = draft && (draft as any)[DRAFT_STATE] if (!state || !state.isManual_) die(9) - const {scope_: scope} = state + + // @ts-ignore -- TODO: Remove this (add typings to state and map!) + const {scope_: scope, existingstateMap_} = state usePatchesInScope(scope, patchListener) - return processResult(undefined, scope) + return processResult(undefined, scope, existingstateMap_, this.allowMultiRefs_ ? new WeakMap() : undefined) as any } /** @@ -167,6 +177,11 @@ export class Immer implements ProducersFns { this.useStrictShallowCopy_ = value } + /** Pass true to allow multiple references to the same object in the same state tree. */ + setAllowMultiRefs(value: boolean) { + this.allowMultiRefs_ = value + } + applyPatches(base: T, patches: Patch[]): T { // If a patch replaces the entire state, take that replacement as base // before applying patches @@ -198,16 +213,19 @@ export class Immer implements ProducersFns { export function createProxy( value: T, - parent?: ImmerState + parent?: ImmerState, + stateMap?: WeakMap ): Drafted { // precondition: createProxy should be guarded by isDraftable, so we know we can safely draft const draft: Drafted = isMap(value) ? getPlugin("MapSet").proxyMap_(value, parent) : isSet(value) ? getPlugin("MapSet").proxySet_(value, parent) - : createProxyProxy(value, parent) + : createProxyProxy(value, parent, stateMap) const scope = parent ? parent.scope_ : getCurrentScope() + scope.drafts_.push(draft) + return draft } diff --git a/src/core/proxy.ts b/src/core/proxy.ts index 3ce06aa8..66e07f74 100644 --- a/src/core/proxy.ts +++ b/src/core/proxy.ts @@ -33,6 +33,7 @@ export interface ProxyObjectState extends ProxyBaseState { base_: any copy_: any draft_: Drafted + stateMap_?: WeakMap | undefined } export interface ProxyArrayState extends ProxyBaseState { @@ -40,6 +41,7 @@ export interface ProxyArrayState extends ProxyBaseState { base_: AnyArray copy_: AnyArray | null draft_: Drafted + stateMap_?: WeakMap | undefined } type ProxyState = ProxyObjectState | ProxyArrayState @@ -51,8 +53,10 @@ type ProxyState = ProxyObjectState | ProxyArrayState */ export function createProxyProxy( base: T, - parent?: ImmerState + parent?: ImmerState, + stateMap?: WeakMap ): Drafted { + const isArray = Array.isArray(base) const state: ProxyState = { type_: isArray ? ArchType.Array : (ArchType.Object as any), @@ -74,7 +78,8 @@ export function createProxyProxy( copy_: null, // Called by the `produce` function. revoke_: null as any, - isManual_: false + isManual_: false, + stateMap_: stateMap } // the traps must target something, a bit like the 'real' base. @@ -116,7 +121,11 @@ export const objectTraps: ProxyHandler = { // Assigned values are never drafted. This catches any drafts we created, too. if (value === peek(state.base_, prop)) { prepareCopy(state) - return (state.copy_![prop as any] = createProxy(value, state)) + return (state.copy_![prop as any] = createProxy( + value, + state, + state.stateMap_ + )) } return value }, @@ -278,15 +287,19 @@ export function markChanged(state: ImmerState) { } } -export function prepareCopy(state: { - base_: any - copy_: any - scope_: ImmerScope -}) { - if (!state.copy_) { - state.copy_ = shallowCopy( - state.base_, - state.scope_.immer_.useStrictShallowCopy_ - ) +export function prepareCopy(state: ImmerState) { + if (state.copy_) return + + const existing = state.stateMap_?.get(state.base_) + if (existing) { + Object.assign(state, existing) + return } + + state.copy_ = shallowCopy( + state.base_, + state.scope_.immer_.useStrictShallowCopy_ + ) + + state.stateMap_?.set(state.base_, state) } diff --git a/src/plugins/mapset.ts b/src/plugins/mapset.ts index edc628a7..faaa58ac 100644 --- a/src/plugins/mapset.ts +++ b/src/plugins/mapset.ts @@ -34,7 +34,8 @@ export function enableMapSet() { base_: target, draft_: this as any, isManual_: false, - revoked_: false + revoked_: false, + stateMap_: parent?.stateMap_ as any } } @@ -167,6 +168,7 @@ export function enableMapSet() { if (!state.copy_) { state.assigned_ = new Map() state.copy_ = new Map(state.base_) + state.stateMap_?.set(state.base_, state) } } @@ -185,7 +187,8 @@ export function enableMapSet() { draft_: this, drafts_: new Map(), revoked_: false, - isManual_: false + isManual_: false, + stateMap_: parent?.stateMap_ as any } } @@ -284,6 +287,7 @@ export function enableMapSet() { if (!state.copy_) { // create drafts for all entries to preserve insertion order state.copy_ = new Set() + state.stateMap_?.set(state.base_, state) state.base_.forEach(value => { if (isDraftable(value)) { const draft = createProxy(value, state) diff --git a/src/utils/plugins.ts b/src/utils/plugins.ts index 36cc1d70..6e84bb97 100644 --- a/src/utils/plugins.ts +++ b/src/utils/plugins.ts @@ -60,6 +60,7 @@ export interface MapState extends ImmerBaseState { base_: AnyMap revoked_: boolean draft_: Drafted + stateMap_?: WeakMap | undefined } export interface SetState extends ImmerBaseState { @@ -69,6 +70,7 @@ export interface SetState extends ImmerBaseState { drafts_: Map // maps the original value to the draft value in the new set revoked_: boolean draft_: Drafted + stateMap_?: WeakMap | undefined } /** Patches plugin */ diff --git a/website/docs/pitfalls.md b/website/docs/pitfalls.md index 55889881..c5a3da07 100644 --- a/website/docs/pitfalls.md +++ b/website/docs/pitfalls.md @@ -17,6 +17,8 @@ Never reassign the `draft` argument (example: `draft = myCoolNewState`). Instead ### Immer only supports unidirectional trees + + Immer assumes your state to be a unidirectional tree. That is, no object should appear twice in the tree, there should be no circular references. There should be exactly one path from the root to any node of the tree. ### Never explicitly return `undefined` from a producer