Add Optional Support For Multiple References to an Object

Some state trees may need to reference an object more than once (such as the tree for my [fomod](https://www.npmjs.com/package/fomod) library. By using a combination of a WeakMap to store existing drafts and an off-by-default configuration option, this should be a painless solution. I've tested it within the scope of my project but there may always be issues I haven't foreseen.

Author: BellCubeDev

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<Objectish, ImmerState>) {
 	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)
+			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)
 	}
 	revokeScope(scope)
 	if (scope.patches_) {
@@ -51,17 +54,18 @@ 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<Objectish, ImmerState>): 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) {
 		each(
 			value,
 			(key, childValue) =>
-				finalizeProperty(rootScope, state, value, key, childValue, path),
+				finalizeProperty(rootScope, state, value, key, childValue, path, undefined, existingStateMap),
 			true // See #590, don't recurse into non-enumerable of non drafted objects
 		)
 		return value
@@ -70,7 +74,7 @@ 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
@@ -90,7 +94,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)
 		)
 		// everything inside is frozen, we can freeze here
 		maybeFreeze(rootScope, result, false)
@@ -104,6 +108,7 @@ function finalize(rootScope: ImmerScope, value: any, path?: PatchPath) {
 			)
 		}
 	}
+
 	return state.copy_
 }
 
@@ -114,10 +119,19 @@ function finalizeProperty(
 	prop: string | number,
 	childValue: any,
 	rootPath?: PatchPath,
-	targetIsSet?: boolean
+	targetIsSet?: boolean,
+	existingStateMap?: WeakMap<Objectish, ImmerState>,
 ) {
 	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_
+		}
+	}
+
 	if (isDraft(childValue)) {
 		const path =
 			rootPath &&
@@ -127,7 +141,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)
 		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 +151,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 +162,10 @@ function finalizeProperty(
 			// See add-data.js perf test
 			return
 		}
-		finalize(rootScope, childValue)
+		finalize(rootScope, childValue, undefined, existingStateMap)
 		// 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);
 	}
 }
 

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)
 		} 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<infer T> ? 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_) 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<T extends Objectish>(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<T extends Objectish>(
 	value: T,
-	parent?: ImmerState
+	parent?: ImmerState,
+	stateMap?: WeakMap<Objectish, ImmerState>
 ): Drafted<T, ImmerState> {
 	// 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
 }

src/core/proxy.ts

@@ -33,13 +33,15 @@ export interface ProxyObjectState extends ProxyBaseState {
 	base_: any
 	copy_: any
 	draft_: Drafted<AnyObject, ProxyObjectState>
+	stateMap_?: WeakMap<Objectish, ImmerState> | undefined
 }
 
 export interface ProxyArrayState extends ProxyBaseState {
 	type_: ArchType.Array
 	base_: AnyArray
 	copy_: AnyArray | null
 	draft_: Drafted<AnyArray, ProxyArrayState>
+	stateMap_?: WeakMap<Objectish, ImmerState> | undefined
 }
 
 type ProxyState = ProxyObjectState | ProxyArrayState
@@ -51,8 +53,10 @@ type ProxyState = ProxyObjectState | ProxyArrayState
  */
 export function createProxyProxy<T extends Objectish>(
 	base: T,
-	parent?: ImmerState
+	parent?: ImmerState,
+	stateMap?: WeakMap<Objectish, ImmerState>
 ): Drafted<T, ProxyState> {
+
 	const isArray = Array.isArray(base)
 	const state: ProxyState = {
 		type_: isArray ? ArchType.Array : (ArchType.Object as any),
@@ -74,7 +78,8 @@ export function createProxyProxy<T extends Objectish>(
 		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<ProxyState> = {
 		// 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)
 }

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)

src/utils/plugins.ts

@@ -60,6 +60,7 @@ export interface MapState extends ImmerBaseState {
 	base_: AnyMap
 	revoked_: boolean
 	draft_: Drafted<AnyMap, MapState>
+	stateMap_?: WeakMap<AnyMap, ImmerState> | undefined
 }
 
 export interface SetState extends ImmerBaseState {
@@ -69,6 +70,7 @@ export interface SetState extends ImmerBaseState {
 	drafts_: Map<any, Drafted> // maps the original value to the draft value in the new set
 	revoked_: boolean
 	draft_: Drafted<AnySet, SetState>
+	stateMap_?: WeakMap<AnySet, ImmerState> | undefined
 }
 
 /** Patches plugin */

website/docs/pitfalls.md

@@ -17,6 +17,8 @@ Never reassign the `draft` argument (example: `draft = myCoolNewState`). Instead
 
 ### Immer only supports unidirectional trees
 
+<!-- TODO: Discuss what to do in the docs regarding the multiple references PR -->
+
 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