MultiMap.ts

function* map<T, S>(iterable: Iterable<T>, callback: (param: T) => S): IterableIterator<S> {
  for (const item of iterable) {
    yield callback(item);
  }
}

function* flatMap<T, S>(iterable: Iterable<T>, callback: (param: T) => Iterable<S>): IterableIterator<S> {
  for (const item of iterable) {
    yield* callback(item);
  }
}

function first<T>(iterable: Iterable<T>): T | undefined {
  const [first] = iterable;
  return first;
}

function count(iterable: Iterable<any>): number {
  return [...iterable].length;
}

/**
 * Created on 2019-03-14.
 *
 * @author Konstantin Simon Maria Möllers
 */
export class MultiMap<K, V> implements Iterable<[K, V]> {
  private readonly keyToValues: Map<K, Set<V>>;

  constructor(entries?: ReadonlyArray<[K, V]> | null) {
    this.keyToValues = new Map();
    if (entries) {
      for (const [key, value] of entries) {
        if (this.keyToValues.has(key)) {
          this.keyToValues.get(key)!.add(value);
        } else {
          this.keyToValues.set(key, new Set<V>([value]));
        }
      }
    }
  }

  get size(): number {
    return count(this.values());
  }

  /**
   * Returns an iterable of entries in the multi map.
   */
  [Symbol.iterator](): IterableIterator<[K, V]> {
    return this.entries();
  }

  /**
   * Returns an iterable of key, value pairs for every entry in the multi map.
   */
  entries(): IterableIterator<[K, V]> {
    return flatMap(this.keyToValues.entries(), ([key, values]) => map(values, value => [key, value] as [K, V]));
  }

  /**
   * Returns an iterable of keys in the multi map.
   */
  keys(): IterableIterator<K> {
    return this.keyToValues.keys();
  }

  /**
   * Returns an iterable of values in the multi map.
   */
  values(): IterableIterator<V> {
    return flatMap(this.keyToValues.values(), value => value.keys());
  }

  /**
   * Clears all entries of the multi map.
   */
  clear(): void {
    this.keyToValues.clear();
  }

  /**
   * Deletes an entry from the multi map.
   */
  delete(key: K): boolean {
    return this.keyToValues.delete(key);
  }

  forEach(callbackfn: (value: V, key: K, map: MultiMap<K, V>) => void, thisArg?: any): void {
    for (const [key, value] of this.entries()) {
      callbackfn.call(thisArg, value, key, this);
    }
  }

  map<L, M>(callbackfn: (value: V, key: K, map: MultiMap<K, V>) => [L, M], thisArg?: any): MultiMap<L, M> {
    return new MultiMap<L, M>([...this.entries()].map(([key, value]) => callbackfn.call(thisArg, value, key, this)));
  }

  filter(predicate: (value: V, key: K, map: MultiMap<K, V>) => boolean, thisArg?: any): MultiMap<K, V> {
    return new MultiMap<K, V>([...this.entries()].filter(([key, value]) => predicate.call(thisArg, value, key, this)));
  }

  get(key: K): V | undefined {
    return first(this.getAll(key));
  }

  getAll(key: K): Set<V> {
    const value = this.keyToValues.get(key);
    if (!value) {
      const newSet = new Set<V>();
      this.keyToValues.set(key, newSet);

      return newSet;
    }

    return value;
  }

  has(key: K): boolean {
    return this.keyToValues.has(key);
  }

  add(key: K, value: V): this {
    const set = this.getAll(key);
    set.add(value);

    return this;
  }

  set(key: K, value: V): this {
    const set = this.getAll(key);
    set.clear();
    set.add(value);

    return this;
  }

  setAll(key: K, value: Iterable<V>): this {
    const set = this.getAll(key);
    set.clear();
    for (const item of value) {
      set.add(item);
    }

    return this;
  }
}