iterable weak map

iterable-weak-map.ts

export default class IterableWeakMap<K extends WeakKey, V> implements Iterable<[K, V]> {
	public constructor() {
		this.finalizationRegistry = new FinalizationRegistry((reference: WeakRef<K>) => {
			this.references.delete(reference);
			this.referenceToValue.delete(reference);
		});
	}

	public set(key: K, value: V) {
		const reference = new WeakRef(key);
		this.references.add(reference);
		this.referenceToValue.set(reference, value);
		this.finalizationRegistry.register(reference, reference, reference);
		return this;
	}

	public get(key: K): V | undefined {
		for (const reference of this.references)
			if (reference.deref() === key) return this.referenceToValue.get(reference);

		return;
	}

	public has(key: K) {
		for (const reference of this.references) if (reference.deref() === key) return true;
		return false;
	}

	public delete(key: K) {
		for (const reference of this.references)
			if (reference.deref() === key) {
				this.references.delete(reference);
				this.referenceToValue.delete(reference);
				this.finalizationRegistry.unregister(reference);
				return true;
			}

		return false;
	}

	public clear() {
		for (const reference of this.references) this.finalizationRegistry.unregister(reference);
		this.references.clear();
		this.referenceToValue.clear();
	}

	public get size() {
		let count = 0;
		// eslint-disable-next-line @typescript-eslint/no-unused-vars
		for (const _ of this) count += 1;
		return count;
	}

	public *entries(): IterableIterator<[K, V]> {
		yield* this[Symbol.iterator]();
	}
	public *keys(): IterableIterator<K> {
		for (const reference of this.references) {
			const key = reference.deref();
			if (key !== undefined) yield key;
		}
	}
	public *values(): IterableIterator<V> {
		const referenceToValue = this.referenceToValue;
		for (const reference of this.references) {
			const key = reference.deref();
			if (key !== undefined) {
				const value = referenceToValue.get(reference);
				if (value !== undefined) yield value;
			}
		}
	}

	public *[Symbol.iterator](): IterableIterator<[K, V]> {
		const referenceToValue = this.referenceToValue;
		for (const reference of this.references) {
			const key = reference.deref();
			if (key !== undefined) {
				const value = referenceToValue.get(reference);
				if (value !== undefined) yield [key, value];
			}
		}
	}

	private readonly references = new Set<WeakRef<K>>();
	private readonly referenceToValue = new Map<WeakRef<K>, V>();
	private readonly finalizationRegistry: FinalizationRegistry<WeakRef<K>>;
}