Implementation of a Multimap type, similar to Map, but each key can have multiple values.

multimap.ts

export interface Multimap<K, V> extends Map<K, V> {
	clear(): void;
	delete(key: K): boolean;
	delete(key: K, value: V): boolean;
	forEach(callbackfn: (value: V, key: K, map: this) => void, thisArg?: any): void;
	get(key: K): V | undefined;
	get(key: K, value: V): V | undefined;
	getAll(key: K): Set<V> | undefined;
	has(key: K): boolean;
	has(key: K, value: V): boolean;
	set(key: K, value: V): this;
	move(from: K, to: K, value: V): void;
	readonly size: number;

	[Symbol.iterator](): IterableIterator<[K, V]>;
	entries(): IterableIterator<[K, V]>;
	keys(): IterableIterator<K>;
	values(): IterableIterator<V>;
	[Symbol.toStringTag]: string;
}

export interface ReadonlyMultimap<K, V> extends ReadonlyMap<K, V> {
	forEach(callbackfn: (value: V, key: K, map: this) => void, thisArg?: any): void;
	get(key: K): V | undefined;
	get(key: K, value: V): V | undefined;
	getAll(key: K): Set<V> | undefined;
	has(key: K): boolean;
	has(key: K, value: V): boolean;
	readonly size: number;

	[Symbol.iterator](): IterableIterator<[K, V]>;
	entries(): IterableIterator<[K, V]>;
	keys(): IterableIterator<K>;
	values(): IterableIterator<V>;
	[Symbol.toStringTag]: string;
}

export interface MultimapConstructor {
	new(): Multimap<any, any>;
	new <K, V>(entries?: ReadonlyArray<readonly [K, V]> | null, persistentSets?: boolean): Multimap<K, V>;
	new <K, V>(iterable: Iterable<readonly [K, V]>, persistentSets?: boolean): Multimap<K, V>;
	readonly prototype: Multimap<any, any>;
}

/*
 * We implement Multimap as an extension of Map<K, Set<V>>. Unfortunately,
 * typescript cannot privately extend from an object. So it always exposes all
 * of the inherited methods, and disallows completely overriding the method
 * types.
 *
 * Thus we create a surrogate superclass type that is type-less. We'll inherit
 * from it in MultimapImpl and implement the correct Multimap interface.
 */
const MultimapSuperclass: {
	new(): any;
	new<K, V>(entries?: ReadonlyArray<readonly [K, Set<V>]> | null): any;
	new <K, V>(iterable: Iterable<readonly [K, Set<V>]>): any;
	readonly prototype: any;
} = Map;

/**
 * @ignore
 *
 * Beware! ALL super.fn below are untyped. Care must be taken to ensure that we
 * don't secretly violate the type constraints of Multimap<K, V>.
 */
class MultimapImpl<K, V> extends MultimapSuperclass<K, V> implements Multimap<K, V> {
	private persistent: Map<K, Set<V>>;

	constructor(iterable?: null | ReadonlyArray<readonly [K, V]> | Iterable<readonly [K, V]>, persistentSets?: boolean) {
		const persistent = new Map<K, Set<V>>();
		super((function* () {
			if(!iterable) return;

			for(let [key, value] of iterable) {
				let set = persistent.get(key);
				if(typeof set === "undefined") {
					persistent.set(key, set = new Set([value]));
					yield [key, set] as const;
				}
				else {
					set.add(value);
				}
			}
		})());

		this.persistent = persistent;
	}

	clear(): void {
		// We must clear each set, otherwise we get odd behavior when clear()
		// is called during iteration.
		for(let [key, set] of super.entries()) {
			set.clear();
		}
		super.clear();
	}

	delete(key: K, value?: V): boolean {
		const set: Set<V> | undefined = super.get(key);
		if(typeof set === "undefined") return false;
		else if(arguments.length === 1) {
			set.clear();
			return super.delete(key);
		}
		else if(set.delete(value!) && set.size === 0) {
			set.clear();
			return super.delete(key);
		}
		return false;
	}

	forEach(callbackfn: (value: V, key: K, map: this) => void, thisArg?: any): void {
		return super.forEach(
			(set: Set<V>, key: K) => set.forEach(
				(value) => callbackfn.call(thisArg, value, key, this)
			)
		);
	}

	get(key: K, value?: V): V | undefined {
		let set: Set<V> | undefined = super.get(key);
		if(set === void 0) return;
		let result = set.values().next().value;
		if(arguments.length === 1 && result !== value) return;
		return value;
	}

	getAll(key: K): Set<V> | undefined {
		return super.get(key);
	}

	has(key: K, value?: V): boolean {
		const set: Set<V> | undefined = super.get(key);
		if(set === void 0) return false;
		if(arguments.length === 1) return true;
		return set.has(value!);
	}

	set(key: K, value: V): this {
		let set: Set<V> | undefined = super.get(key);
		if(set === void 0 && (set = this.persistent.get(key)) === void 0) {
			this.persistent.set(key, set = new Set([value]));
			super.set(key, set);
		}
		else {
			set.add(value);
		}

		return this;
	}

	move(from: K, to: K, value: V): void {
		let fromSet: Set<V> | undefined = super.get(from);
		if(fromSet) fromSet.delete(value);
		let toSet: Set<V> | undefined = super.get(to);
		if(toSet) toSet.add(value);
	}

	*entries(): IterableIterator<[K, V]> {
		for(let [key, set] of super.entries() as IterableIterator<[K, Set<V>]> ) {
			for(let value of set.values()) {
				yield [key, value];
			}
		}
	}

	*values(): IterableIterator<V> {
		for(let [key, set] of super.entries() as IterableIterator<[K, Set<V>]> ) {
			yield* set.values();
		}
	}
}
interface MultimapImpl<K, V> {
	// inherited from hidden Map<K, V>
	readonly size: number;
	clear(): void;
	keys(): IterableIterator<K>;

	// Set on prototype
	[Symbol.iterator](): IterableIterator<[K, V]>;
	[Symbol.toStringTag]: string;
}
MultimapImpl.prototype[Symbol.iterator as any] = MultimapImpl.prototype.entries;
MultimapImpl.prototype[Symbol.toStringTag as any] = "Multimap";

export const Multimap: MultimapConstructor = MultimapImpl;