fix: track Map upserts and protect frozen Maps

__tests__/frozen.js

@@ -158,6 +158,28 @@ function runTests(name) {
 			expect(produce(res, d => void d.set("a", 2))).not.toBe(res)
 		})
 
+		it.runIf(
+			"getOrInsert" in Map.prototype && "getOrInsertComputed" in Map.prototype
+		)("will freeze map upsert methods", () => {
+			const res = produce(new Map([["a", 1]]), draft => {
+				draft.set("b", 2)
+			})
+			let calls = 0
+			expect(() => res.getOrInsert("a", 3)).toThrow()
+			expect(() => res.getOrInsert("c", 3)).toThrow()
+			expect(() =>
+				res.getOrInsertComputed("c", () => {
+					calls++
+					return 3
+				})
+			).toThrow()
+			expect(calls).toBe(0)
+			expect(Array.from(res)).toEqual([
+				["a", 1],
+				["b", 2]
+			])
+		})
+
 		it("will freeze sets", () => {
 			const base = new Set()
 			const res = produce(base, draft => {

__tests__/map-set.js

@@ -14,6 +14,9 @@ import {each, shallowCopy, isEnumerable, DRAFT_STATE} from "../src/utils/common"
 enableMapSet()
 enablePatches()
 
+const hasNativeMapUpsert =
+	"getOrInsert" in Map.prototype && "getOrInsertComputed" in Map.prototype
+
 vi.setConfig({
 	testTimeout: 1000
 })
@@ -362,6 +365,127 @@ function runBaseTest(name, autoFreeze, useListener) {
 			expect(result.map.keys().toArray()).toEqual([1, 3])
 		})
 	})
+	describe.runIf(hasNativeMapUpsert)("map upsert methods " + name, () => {
+		test("getOrInsert reads existing keys and inserts missing ones", () => {
+			const base = new Map([
+				["a", 1],
+				["u", undefined]
+			])
+			const [result, patches] = produceWithPatches(base, draft => {
+				expect(draft.getOrInsert("a", 9)).toBe(1)
+				expect(draft.getOrInsert("u", 9)).toBe(undefined)
+				expect(draft.getOrInsert("b", 2)).toBe(2)
+				expect(draft.getOrInsert("b", 3)).toBe(2)
+				expect(draft.has("b")).toBe(true)
+				expect(draft.get("b")).toBe(2)
+			})
+			expect(Array.from(result)).toEqual([
+				["a", 1],
+				["u", undefined],
+				["b", 2]
+			])
+			expect(patches).toEqual([{op: "add", path: ["b"], value: 2}])
+			expect(Array.from(base)).toEqual([
+				["a", 1],
+				["u", undefined]
+			])
+		})
+
+		test("getOrInsert returns a draft of an existing object value", () => {
+			const base = new Map([["a", {n: 1}]])
+			const fallback = {n: 9}
+			const result = produce(base, draft => {
+				const value = draft.getOrInsert("a", fallback)
+				expect(isDraft(value)).toBe(true)
+				value.n++
+			})
+			expect(result.get("a")).toEqual({n: 2})
+			expect(base.get("a")).toEqual({n: 1})
+			expect(fallback).toEqual({n: 9})
+		})
+
+		test("getOrInsertComputed only calls the callback for missing keys", () => {
+			const base = new Map([["a", {n: 1}]])
+			const calls = []
+			const [result, patches] = produceWithPatches(base, draft => {
+				const existing = draft.getOrInsertComputed("a", key => {
+					calls.push(key)
+					return {n: 9}
+				})
+				expect(isDraft(existing)).toBe(true)
+				existing.n++
+				const inserted = draft.getOrInsertComputed(-0, function (key) {
+					calls.push({
+						argc: arguments.length,
+						thisValue: this,
+						positiveZero: Object.is(key, 0)
+					})
+					return "zero"
+				})
+				expect(inserted).toBe("zero")
+				expect(draft.get(0)).toBe("zero")
+			})
+			expect(calls).toEqual([
+				{argc: 1, thisValue: undefined, positiveZero: true}
+			])
+			expect(Array.from(result)).toEqual([
+				["a", {n: 2}],
+				[0, "zero"]
+			])
+			expect(patches).toEqual([
+				{op: "replace", path: ["a", "n"], value: 2},
+				{op: "add", path: [0], value: "zero"}
+			])
+		})
+
+		test("getOrInsertComputed stores the computed value over a reentrant set", () => {
+			const result = produce(new Map(), draft => {
+				expect(
+					draft.getOrInsertComputed("c", key => {
+						draft.set(key, 7)
+						return 8
+					})
+				).toBe(8)
+				expect(draft.get("c")).toBe(8)
+			})
+			expect(Array.from(result)).toEqual([["c", 8]])
+		})
+
+		test("getOrInsertComputed inserts nothing when the callback throws", () => {
+			const base = new Map([["a", 1]])
+			const result = produce(base, draft => {
+				expect(() =>
+					draft.getOrInsertComputed("b", () => {
+						throw new Error("boom")
+					})
+				).toThrow("boom")
+				expect(draft.has("b")).toBe(false)
+			})
+			expect(result).toBe(base)
+		})
+
+		test("getOrInsertComputed validates the callback before lookup", () => {
+			produce(new Map([["a", 1]]), draft => {
+				expect(() => draft.getOrInsertComputed("a", null)).toThrow(TypeError)
+				expect(() => draft.getOrInsertComputed("b", 1)).toThrow(TypeError)
+				expect(draft.has("b")).toBe(false)
+			})
+		})
+
+		test("upsert methods throw on a revoked draft", () => {
+			let revoked
+			produce(new Map(), draft => {
+				revoked = draft
+			})
+			expect(() => revoked.getOrInsert("a", 1)).toThrow(
+				global.USES_BUILD ? "minified error nr: 3" : "revoked"
+			)
+			expect(() => revoked.getOrInsertComputed("a", () => 1)).toThrow(
+				global.USES_BUILD ? "minified error nr: 3" : "revoked"
+			)
+		})
+	})
+
 	describe("set issues " + name, () => {
 		test("#819.A - maintains order when adding", () => {
 			const objs = [
@@ -533,3 +657,35 @@ describe("RTK-5159: DraftSet leakage in Map", () => {
 		expect(() => JSON.stringify(result)).not.toThrow()
 	})
 })
+
+describe("Map upsert methods on runtimes without them", () => {
+	const upsertMethods = ["getOrInsert", "getOrInsertComputed"]
+
+	test("are not added to drafts or frozen Maps", async () => {
+		const descriptors = upsertMethods.map(method =>
+			Object.getOwnPropertyDescriptor(Map.prototype, method)
+		)
+		upsertMethods.forEach(method => delete Map.prototype[method])
+		vi.resetModules()
+		try {
+			const immer = await import("../src/immer")
+			immer.enableMapSet()
+			const result = immer.produce(new Map([["a", 1]]), draft => {
+				upsertMethods.forEach(method => {
+					expect(method in draft).toBe(false)
+				})
+				draft.set("b", 2)
+			})
+			expect(Object.isFrozen(result)).toBe(true)
+			upsertMethods.forEach(method => {
+				expect(method in result).toBe(false)
+			})
+		} finally {
+			upsertMethods.forEach((method, i) => {
+				if (descriptors[i])
+					Object.defineProperty(Map.prototype, method, descriptors[i])
+			})
+			vi.resetModules()
+		}
+	})
+})

src/plugins/mapset.ts

@@ -17,7 +17,8 @@ import {
 	each,
 	getValue,
 	PluginMapSet,
-	handleCrossReference
+	handleCrossReference,
+	nativeMapUpsertMethods
 } from "../internal"
 
 // Feature-detect `Iterator.from` (added in ES2025) without augmenting the
@@ -181,6 +182,42 @@ export function enableMapSet() {
 		}
 	}
 
+	const draftMapUpsertMethods = {
+		getOrInsert(this: DraftMap, key: unknown, value: unknown) {
+			if (this.has(key)) {
+				return this.get(key)
+			}
+			this.set(key, value)
+			return value
+		},
+
+		getOrInsertComputed(
+			this: DraftMap,
+			key: unknown,
+			callback: (key: unknown) => unknown
+		) {
+			assertUnrevoked(this[DRAFT_STATE])
+			if (typeof callback !== "function") {
+				throw new TypeError("getOrInsertComputed callback is not a function")
+			}
+			if (this.has(key)) {
+				return this.get(key)
+			}
+			const canonicalKey = Object.is(key, -0) ? 0 : key
+			const value = callback(canonicalKey)
+			this.set(canonicalKey, value)
+			return value
+		}
+	}
+
+	nativeMapUpsertMethods.forEach(method =>
+		Object.defineProperty(DraftMap.prototype, method, {
+			value: draftMapUpsertMethods[method],
+			writable: true,
+			configurable: true
+		})
+	)
+
 	function iteratorFrom<T, TReturn>(
 		iterable: Iterator<T, TReturn>
 	): IterableIterator<T> {

src/utils/common.ts

@@ -273,6 +273,10 @@ export function freeze<T>(obj: any, deep: boolean = false): T {
 			clear: dontMutateMethodOverride,
 			delete: dontMutateMethodOverride
 		})
+		if (isMap(obj))
+			nativeMapUpsertMethods.forEach(method =>
+				O.defineProperty(obj, method, dontMutateMethodOverride)
+			)
 	}
 	O.freeze(obj)
 	if (deep)
@@ -296,6 +300,10 @@ const dontMutateMethodOverride = {
 	[VALUE]: dontMutateFrozenCollections
 }
 
+export const nativeMapUpsertMethods = (
+	["getOrInsert", "getOrInsertComputed"] as const
+).filter(method => method in Map.prototype)
+
 export function isFrozen(obj: any): boolean {
 	// Fast path: primitives and null/undefined are always "frozen"
 	if (obj === null || !isObjectish(obj)) return true