Transparent Map<K, WeakRef<V extends WeakKey>>

map-weak.ts

const retainment = new WeakMap<WeakRef<WeakKey>, unknown>();

export class MapWeak<K, V extends WeakKey> implements Map<K, V> {
	#map: Map<K, WeakRef<V>>;
	#finalizer = new FinalizationRegistry<K>(key => {
		this.#map.delete(key);
	});

	constructor(entries?: Iterable<readonly [K, V]> | null | undefined) {
		const entriesWeak: ConstructorParameters<typeof Map<K, WeakRef<V>>>[0] =
			entries &&
			Array.from(entries, ([key, value]) => {
				this.#finalizer.register(value, key);
				return [key, new WeakRef(value)];
			});

		this.#map = new Map(entriesWeak);
	}

	retain(key: K) {
		const ref = this.#map.get(key);
		if (ref) retainment.set(ref, ref.deref());
	}

	release(key: K) {
		const ref = this.#map.get(key);
		if (ref) retainment.delete(ref);
	}

	clear() {
		for (const value of this.#map.values()) {
			this.#finalizer.unregister(value.deref()!);
		}
		this.#map.clear();
	}

	delete(key: K) {
		const value = this.#map.get(key);
		if (this.#map.delete(key)) {
			this.#finalizer.unregister(value!.deref()!);
			return true;
		}
		return false;
	}

	forEach<This = undefined>(callbackfn: (this: This, value: V, key: K, map: this) => void, thisArg?: This) {
		this.#map.forEach((value, key) => callbackfn.call(thisArg!, value.deref()!, key, this));
	}

	get(key: K) {
		return this.#map.get(key)?.deref();
	}

	has(key: K) {
		return this.#map.has(key);
	}

	set(key: K, value: V) {
		this.#map.set(key, new WeakRef(value));
		return this;
	}

	get size() {
		return this.#map.size;
	}

	*entries(): MapIterator<[K, V]> {
		for (const [key, ref] of this.#map.entries()) yield [key, ref.deref()!];
	}

	keys(): MapIterator<K> {
		return this.#map.keys();
	}

	*values(): MapIterator<V> {
		for (const ref of this.#map.values()) yield ref.deref()!;
	}

	*[Symbol.iterator](): MapIterator<[K, V]> {
		for (const [key, ref] of this.#map.entries()) yield [key, ref.deref()!];
	}

	get [Symbol.toStringTag]() {
		return "MapWeak";
	}
}