MeekValueMap getOrInsert and getOrInsertComputed

Author: AlexanderOMara

valuemap.test.ts

@@ -94,6 +94,48 @@ Deno.test('MeekValueMap: get', () => {
 	assert(pairs);
 });
 
+Deno.test('MeekValueMap: getOrInsert', () => {
+	const pairs: readonly [number, { i: number }][] = new Array(100).fill(0)
+		.map((_, i) => [i, { i }]);
+	const map = new MeekValueMap();
+	for (let i = 0; i < pairs.length; i++) {
+		const [k, v] = pairs[i];
+		assertStrictEquals(map.has(k), false);
+		assertStrictEquals(map.getOrInsert(k, v), v);
+		assertStrictEquals(map.has(k), true);
+		assertStrictEquals(map.get(k), v);
+		const o = { i: -i };
+		assertStrictEquals(map.getOrInsert(k, o), v);
+		assertStrictEquals(map.get(k), v);
+	}
+	assert(pairs);
+});
+
+Deno.test('MeekValueMap: getOrInsertComputed', () => {
+	const uncalled = () => {
+		throw new Error('Should not be called');
+	};
+	const pairs: readonly [number, { i: number }][] = new Array(100).fill(0)
+		.map((_, i) => [i, { i }]);
+	const map = new MeekValueMap();
+	for (let i = 0; i < pairs.length; i++) {
+		const [k, v] = pairs[i];
+		assertStrictEquals(map.has(k), false);
+		assertStrictEquals(
+			map.getOrInsertComputed(k, (key) => {
+				assertStrictEquals(key, k);
+				return v;
+			}),
+			v,
+		);
+		assertStrictEquals(map.has(k), true);
+		assertStrictEquals(map.get(k), v);
+		assertStrictEquals(map.getOrInsertComputed(k, uncalled), v);
+		assertStrictEquals(map.get(k), v);
+	}
+	assert(pairs);
+});
+
 Deno.test('MeekValueMap: set', () => {
 	const pairs: readonly [number, { i: number }][] = new Array(100).fill(0)
 		.map((_, i) => [i, { i }]);
@@ -242,36 +284,22 @@ Deno.test('MeekValueMap: Symbol.toStringTag', () => {
 	);
 });
 
-Deno.test('MeekValueMap: GC', async () => {
+Deno.test('MeekValueMap: GC set', async () => {
 	await forceGC((map, key, value) => {
 		map.set(key, value);
 	});
 });
 
-Deno.test('MeekValueMap: modify while iterating', () => {
-	const pairs: readonly [number, { i: number }][] = new Array(100).fill(0)
-		.map((_, i) => [i, { i }]);
-	const mapExpt = new Map(pairs.slice(0, 60));
-	const mapTest = new MeekValueMap(pairs.slice(0, 60));
-	const readWhileModify = (map: Map<number, WeakKey> | MeekValueMap) => {
-		const r: number[] = [];
-		let i = 0;
-		for (const [k] of map) {
-			r.push(k);
-			map.delete(pairs[k + 1][0]);
-			if (i++ === 10) {
-				map.clear();
-				for (const [k, v] of pairs.slice(50)) {
-					map.set(k, v);
-				}
-			}
-		}
-		return r;
-	};
-	const valExpt = readWhileModify(mapExpt);
-	const valTest = readWhileModify(mapTest);
-	assertEquals(valExpt, valTest);
-	assert(pairs);
+Deno.test('MeekValueMap: GC getOrInsert', async () => {
+	await forceGC((map, key, value) => {
+		map.getOrInsert(key, value);
+	});
+});
+
+Deno.test('MeekValueMap: GC getOrInsertComputed', async () => {
+	await forceGC((map, key, value) => {
+		map.getOrInsertComputed(key, () => value);
+	});
 });
 
 /*

valuemap.ts

@@ -24,6 +24,11 @@ interface Pri<K = any, V extends WeakKey = WeakKey> {
 
 const pri = new WeakMap<MeekValueMap, Pri>();
 
+const unregister = (
+	fr: FinalizationRegistry<unknown>,
+	wr?: WeakRef<any>,
+): unknown => wr && fr.unregister(wr);
+
 /**
  * Like WeakValueMap.
  *
@@ -45,11 +50,8 @@ export class MeekValueMap<K = any, V extends WeakKey = WeakKey> {
 		const kwv = new Map<K, WeakRef<V>>();
 		const fr = new FinalizationRegistry(kwv.delete.bind(kwv));
 		for (const [key, value] of iterable ?? []) {
-			let ref = kwv.get(key);
-			if (ref) {
-				fr.unregister(ref);
-			}
-			ref = new WeakRef(value);
+			unregister(fr, kwv.get(key));
+			const ref = new WeakRef(value);
 			fr.register(value, key, ref);
 			kwv.set(key, ref);
 		}
@@ -139,6 +141,53 @@ export class MeekValueMap<K = any, V extends WeakKey = WeakKey> {
 		return (pri.get(this) as Pri<K, V>).kwv.get(key)?.deref();
 	}
 
+	/**
+	 * Get the value for a key from this map or insert the default value.
+	 *
+	 * @param key Key to get.
+	 * @param defaultValue Default value.
+	 * @returns Value for the key.
+	 */
+	public getOrInsert(key: K, defaultValue: V): V {
+		const { fr, kwv } = pri.get(this) as Pri<K, V>;
+		const old = kwv.get(key);
+		if (old) {
+			const r = old.deref();
+			if (r) {
+				return r;
+			}
+		}
+		const ref = new WeakRef(defaultValue);
+		unregister(fr, old);
+		fr.register(defaultValue, key, ref);
+		kwv.set(key, ref);
+		return defaultValue;
+	}
+
+	/**
+	 * Get the value for a key from this map or insert the default value.
+	 *
+	 * @param key Key to get.
+	 * @param callback Compute the default value.
+	 * @returns Value for the key.
+	 */
+	public getOrInsertComputed(key: K, callback: (key: K) => V): V {
+		const { fr, kwv } = pri.get(this) as Pri<K, V>;
+		const old = kwv.get(key);
+		if (old) {
+			const r = old.deref();
+			if (r) {
+				return r;
+			}
+		}
+		const value = callback(key);
+		const ref = new WeakRef(value);
+		unregister(fr, old);
+		fr.register(value, key, ref);
+		kwv.set(key, ref);
+		return value;
+	}
+
 	/**
 	 * Has a key in this map.
 	 *
@@ -172,10 +221,7 @@ export class MeekValueMap<K = any, V extends WeakKey = WeakKey> {
 	public set(key: K, value: V): this {
 		const { fr, kwv } = pri.get(this) as Pri<K, V>;
 		const ref = new WeakRef(value);
-		const old = kwv.get(key);
-		if (old) {
-			fr.unregister(old);
-		}
+		unregister(fr, kwv.get(key));
 		fr.register(value, key, ref);
 		kwv.set(key, ref);
 		return this;