map.ts

/**
 * @module
 *
 * MeekMap data collection.
 */

/**
 * Private data.
 *
 * @template K Key type.
 * @template V Value type.
 */
interface Pri<K extends WeakKey = WeakKey, V = any> {
	/**
	 * Finalization registry.
	 */
	readonly fr: FinalizationRegistry<WeakRef<K>>;

	/**
	 * Map of keys to weak references to keys.
	 */
	kwk: WeakMap<K, WeakRef<K>>;

	/**
	 * Set of weak references to keys.
	 */
	readonly wk: Set<WeakRef<K>>;

	/**
	 * Map of keys to values.
	 */
	kv: WeakMap<K, V>;
}

const pri = new WeakMap<MeekMap, Pri>();

/**
 * Like WeakMap.
 *
 * @template K Key type.
 * @template V Value type.
 */
export class MeekMap<K extends WeakKey = WeakKey, V = any> {
	/**
	 * Type string.
	 */
	declare public readonly [Symbol.toStringTag]: string;

	/**
	 * Create a new MeekMap.
	 *
	 * @param iterable Initial pairs.
	 */
	constructor(iterable?: Iterable<readonly [K, V]> | null) {
		const kwk = new WeakMap<K, WeakRef<K>>();
		const wk = new Set<WeakRef<K>>();
		const fr = new FinalizationRegistry(wk.delete.bind(wk));
		const kv = new WeakMap<K, V>();
		for (const [key, value] of iterable ?? []) {
			let ref = kwk.get(key);
			if (!ref) {
				ref = new WeakRef(key);
				fr.register(key, ref, key);
				kwk.set(key, ref);
				wk.add(ref);
			}
			kv.set(key, value);
		}
		pri.set(this, { fr, kwk, wk, kv });
	}

	/**
	 * Iterator for key-value pairs in this map.
	 *
	 * @returns Key-value iterator.
	 */
	public *[Symbol.iterator](): MapIterator<[K, V]> {
		const p = pri.get(this) as Pri<K, V>;
		for (const ref of p.wk) {
			const key = ref.deref();
			if (key) {
				yield [key, p.kv.get(key) as V];
			}
		}
	}

	/**
	 * Clear this map.
	 */
	public clear(): void {
		const p = pri.get(this) as Pri<K, V>;
		const map = new WeakMap<K, WeakRef<K>>();
		const values = new WeakMap<K, V>();
		p.wk.clear();
		p.kwk = map;
		p.kv = values;
	}

	/**
	 * Delete a key from this map.
	 *
	 * @param key Key to delete.
	 * @returns Whether the key was deleted.
	 */
	public delete(key: K): boolean {
		const { fr, kv, kwk, wk } = pri.get(this) as Pri<K, V>;
		const ref = kwk.get(key);
		if (ref) {
			fr.unregister(key);
			kwk.delete(key);
			kv.delete(key);
			return wk.delete(ref);
		}
		return false;
	}

	/**
	 * Iterator for key-value pairs in this map.
	 *
	 * @returns Key-value iterator.
	 */
	public *entries(): MapIterator<[K, V]> {
		const p = pri.get(this) as Pri<K, V>;
		for (const ref of p.wk) {
			const key = ref.deref();
			if (key) {
				yield [key, p.kv.get(key) as V];
			}
		}
	}

	/**
	 * Call a function for each pair in this map.
	 *
	 * @param callbackfn Callback function.
	 * @param thisArg This argument.
	 */
	public forEach(
		callbackfn: (value: V, key: K, map: MeekMap<K, V>) => void,
		thisArg?: any,
	): void {
		const p = pri.get(this) as Pri<K, V>;
		for (const ref of p.wk) {
			const key = ref.deref();
			if (key) {
				callbackfn.call(thisArg, p.kv.get(key) as V, key, this);
			}
		}
	}

	/**
	 * Get the value for a key from this map.
	 *
	 * @param key Key to get.
	 * @returns Value for the key.
	 */
	public get(key: K): V | undefined {
		return (pri.get(this) as Pri<K, V>).kv.get(key);
	}

	/**
	 * 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, kv, kwk, wk } = pri.get(this) as Pri<K, V>;
		if (kv.has(key)) {
			return kv.get(key)!;
		}
		const ref = new WeakRef(key);
		fr.register(key, ref, key);
		kwk.set(key, ref);
		wk.add(ref);
		kv.set(key, defaultValue);
		return defaultValue;
	}

	/**
	 * Get the value for a key from this map or insert computed 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, kv, kwk, wk } = pri.get(this) as Pri<K, V>;
		if (kv.has(key)) {
			return kv.get(key)!;
		}
		const value = callback(key);
		const ref = new WeakRef(key);
		fr.register(key, ref, key);
		kwk.set(key, ref);
		wk.add(ref);
		kv.set(key, value);
		return value;
	}

	/**
	 * Has a key in this map.
	 *
	 * @param key Key to check.
	 * @returns Whether the key is in this map.
	 */
	public has(key: K): boolean {
		return (pri.get(this) as Pri<K, V>).kv.has(key);
	}

	/**
	 * Iterator for keys in this map.
	 *
	 * @returns Key iterator.
	 */
	public *keys(): MapIterator<K> {
		const { wk } = pri.get(this) as Pri<K, V>;
		for (const ref of wk) {
			const key = ref.deref();
			if (key) {
				yield key;
			}
		}
	}

	/**
	 * Set a value for a key in this map.
	 *
	 * @param key Key to set.
	 * @param value Value to set.
	 * @returns This map.
	 */
	public set(key: K, value: V): this {
		const { fr, kv, kwk, wk } = pri.get(this) as Pri<K, V>;
		let ref = kwk.get(key);
		if (!ref) {
			ref = new WeakRef(key);
			fr.register(key, ref, key);
			kwk.set(key, ref);
		}
		wk.add(ref);
		kv.set(key, value);
		return this;
	}

	/**
	 * The number of keys in this map.
	 * Can be greater than number of active keys.
	 */
	public get size(): number {
		return (pri.get(this) as Pri<K, V>).wk.size;
	}

	/**
	 * Iterator for values in this map.
	 *
	 * @returns Value iterator.
	 */
	public *values(): MapIterator<V> {
		const p = pri.get(this) as Pri<K, V>;
		for (const ref of p.wk) {
			const key = ref.deref();
			if (key) {
				yield p.kv.get(key) as V;
			}
		}
	}

	static {
		Object.defineProperty(this.prototype, Symbol.toStringTag, {
			value: 'MeekMap',
			configurable: true,
			enumerable: false,
			writable: false,
		});
	}
}

/**
 * Readonly MeekMap.
 *
 * @template K Key type.
 * @template V Value type.
 */
export type ReadonlyMeekMap<K extends WeakKey = WeakKey, V = any> = Omit<
	MeekMap<K, V>,
	'clear' | 'delete' | 'getOrInsert' | 'getOrInsertComputed' | 'set'
>;