custom ts map with eq and cmp

map.ts

export enum Ordering {
  Less = -1,
  Equal = 0,
  Greater = 1,
}

export interface Eq<T> {
  eq(other: T): boolean;
}

export interface Cmp<T> {
  cmp(other: T): Ordering;
}

export class TravvyMap<K, V> extends Map<K,V> {
  private equalityFn: (a: K, b: K) => boolean;
  private comparatorFn: (a: K, b: K) => number;
  private _sortedEntries?: Array<[K, V]>;

  constructor(
    equalityFn?: (a: K, b: K) => boolean,
    comparatorFn?: (a: K, b: K) => number
  ) {
    super();
    this.equalityFn = equalityFn || defaultEqualityFn;
    this.comparatorFn = comparatorFn || defaultComparatorFn;
  }
  
  private findKey(key: K): K | undefined {
    for (const existingKey of super.keys()) {
      if (this.equalityFn(existingKey, key)) {
        return existingKey;
      }
    }
    return undefined;
  }

  set(key: K, value: V): this {
    this.invalidate();
    return super.set(this.findKey(key) ?? key, value);
  }

  get(key: K): V | undefined {
    const existing = this.findKey(key);
    return existing !== undefined ? super.get(existing) : undefined;
  }

  has(key: K): boolean {
    return this.findKey(key) !== undefined;
  }

  delete(key: K): boolean {
    const existingKey = this.findKey(key);
    if (existingKey !== undefined) {
      this.invalidate();
      return super.delete(existingKey);
    }
    return false;
  }

  invalidate() {
    this._sortedEntries = undefined;
  }

  get sortedEntries() {
    if (this._sortedEntries === undefined) {
      this._sortedEntries = Array.from(super.entries());
      this._sortedEntries.sort(([a], [b]) => this.comparatorFn(a, b));
    }
    return this._sortedEntries;
  }

  clear(): void {
    this.invalidate();
    super.clear();
  }

  *keys(): IterableIterator<K> {
    for (const [k] of this.sortedEntries) {
      yield k;
    }
  }

  *values(): IterableIterator<V> {
    for (const [_, v] of this.sortedEntries) {
      yield v;
    }
  }

  *entries(): IterableIterator<[K, V]> {
    for (const entry of this.sortedEntries) {
      yield entry;
    }
  }

  forEach(fn: (value: V, key: K, map: this) => void): void {
    for (const [k, v] of this.sortedEntries) {
      fn(v, k, this);
    }
  }
  
  [Symbol.iterator](): IterableIterator<[K, V]> {
    return this.entries();
  }
}

function defaultEqualityFn<K>(a: K, b: K): boolean {
  if (implementsEq(a)) {
    return a.eq(b);
  }
  return a === b;
}

function defaultComparatorFn<K>(a: K, b: K): Ordering {
  if (implementsCmp(a)) {
    return a.cmp(b);
  }
  if (a < b) return Ordering.Less;
  if (a > b) return Ordering.Greater;
  return Ordering.Equal;
}

function implementsEq<T>(x: T | Eq<T>): x is Eq<T> {
  return (
    typeof x === "object"
    && x !== null
    && "eq" in x
    && typeof x.eq === "function"
    && x.eq.length === 1
  );
}

function implementsCmp<T>(x: T | Cmp<T>): x is Cmp<T> {
  return (
    typeof x === "object"
    && x !== null
    && "cmp" in x
    && typeof x.cmp === "function"
    && x.cmp.length === 1
  );
}



// example
class Person implements Eq<Person>, Cmp<Person> {
  constructor(public id: number, public name: string) {}

  eq(other: Person) {
    return this.id === other.id;
  }

  cmp(other: Person) {
    return this.name.localeCompare(other.name);
  }
}

const team = new TravvyMap<Person, string>();

team.set(new Person(1, "Dennis"), "The Looks");
team.set(new Person(2, "Mac"), "The Brains");
team.set(new Person(3, "Charlie"), "The Wildcard");
team.set(new Person(4, "Frank"), "The Muscle");
team.set(new Person(5, "Dee"), "Bird");

console.log(team.get(new Person(2, "Mac"))); // "The Brains"
console.log(team.has(new Person(3, "Charlie"))); // true

console.log(Array.from(team.entries()));

set.ts

export class CustomSet<T> extends Set<T> {
  private equalityFn: (a: T, b: T) => boolean;
  private comparatorFn: (a: T, b: T) => number;
  private _sortedElements?: Array<T>;

  constructor(
    equalityFn?: (a: T, b: T) => boolean,
    comparatorFn?: (a: T, b: T) => number,
    iterable?: Iterable<T>
  ) {
    super(iterable);
    this.equalityFn = equalityFn || defaultEqualityFn;
    this.comparatorFn = comparatorFn || defaultComparatorFn;
  }

  private findElement(element: T): T | undefined {
    for (const existingElement of super.values()) {
      if (this.equalityFn(existingElement, element)) {
        return existingElement;
      }
    }
    return undefined;
  }

  add(element: T): this {
    const existing = this.findElement(element);
    if (existing === undefined) {
      super.add(element);
      this.invalidate();
    }
    return this;
  }

  has(element: T): boolean {
    return this.findElement(element) !== undefined;
  }

  delete(element: T): boolean {
    const existingElement = this.findElement(element);
    if (existingElement !== undefined) {
      const result = super.delete(existingElement);
      this.invalidate();
      return result;
    }
    return false;
  }

  invalidate() {
    this._sortedElements = undefined;
  }

  get sortedElements() {
    if (this._sortedElements === undefined) {
      this._sortedElements = Array.from(super.values());
      this._sortedElements.sort(this.comparatorFn);
    }
    return this._sortedElements;
  }

  clear(): void {
    this.invalidate();
    super.clear();
  }

  *values(): IterableIterator<T> {
    yield* this.sortedElements;
  }

  *keys(): IterableIterator<T> {
    yield* this.values();
  }

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

  forEach(fn: (value: T, value2: T, set: this) => void): void {
    for (const element of this.sortedElements) {
      fn(element, element, this);
    }
  }

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

function defaultEqualityFn<T>(a: T, b: T): boolean {
  if (implementsEq(a)) {
    return a.eq(b);
  }
  return a === b;
}

function defaultComparatorFn<T>(a: T, b: T): number {
  if (implementsCmp(a)) {
    return a.cmp(b);
  }
  if (a < b) return -1;
  if (a > b) return 1;
  return 0;
}

function implementsEq<T>(x: T | Eq<T>): x is Eq<T> {
  return (
    typeof x === "object"
    && x !== null
    && "eq" in x
    && typeof x.eq === "function"
    && x.eq.length === 1
  );
}

function implementsCmp<T>(x: T | Cmp<T>): x is Cmp<T> {
  return (
    typeof x === "object"
    && x !== null
    && "cmp" in x
    && typeof x.cmp === "function"
    && x.cmp.length === 1
  );
}

// Example usage
const team = new CustomSet<Person>();

team.add(new Person(1, "Dennis"));
team.add(new Person(2, "Mac"));
team.add(new Person(3, "Charlie"));
team.add(new Person(4, "Frank"));
team.add(new Person(5, "Dee"));

console.log(team.has(new Person(2, "Mac"))); // true
console.log(Array.from(team)); // Sorted by name: Charlie, Dee, Dennis, Frank, Mac