KeyedSet -- set with a "key", like the Python set

keyed_set.ts

export class KeyedSet<T, K> implements Set<T> {
  private readonly keyValMap: Map<K, T> = new Map();

  constructor(private readonly makeKey: (value: T) => K) {}

  add(value: T) {
    const key = this.makeKey(value);
    this.keyValMap.set(key, value);
    return this;
  }

  clear() {
    this.keyValMap.clear();
  }

  delete(value: T) {
    const key = this.makeKey(value);
    return this.keyValMap.delete(key);
  }

  forEach(callbackfn: (v1: T, v2: T, set: Set<T>) => void, thisArg: any) {
    this.keyValMap.forEach(value => callbackfn(value, value, this), thisArg);
  }

  *entries(): IterableIterator<[T, T]> {
    for (const v of this.keyValMap.values()) {
      yield [v, v];
    }
  }

  has(value: T) {
    return this.keyValMap.has(this.makeKey(value));
  }

  [Symbol.iterator](): IterableIterator<T> {
    return this.keyValMap.values();
  }

  get [Symbol.toStringTag](): string {
    return '[object KeyedSet]';
  }

  get size() {
    return this.keyValMap.size;
  }

  get keys() {
    return this.keyValMap.values;
  }

  get values() {
    return this.keyValMap.values;
  }
}

describe('KeyedSet', () => {
  it('behaves like a native set', () => {
    type Box<T> = { value: T };
    const box = <T>(value: T) => ({ value } as Box<T>);

    const keyedSet = new KeyedSet<Box<string>, string>(box => box.value);

    keyedSet.add(box('hello'));
    keyedSet.add(box('world'));

    expect(keyedSet.has(box('hello'))).toBe(true);
    expect(keyedSet.has(box('world'))).toBe(true);
  });
});