multi-key-map.ts

type NestedMap<K extends any[], V> = Map<K[number], NestedMap<K, V>>;

export class MultiKeyMap<K extends any[], V> implements Map<K, V> {
  #data: NestedMap<K, V> = new Map();
  #depth = 0;

  constructor(iterable?: [K, V][]) {
    if (iterable) {
      for (const [keys, value] of iterable) {
        this.set(keys, value);
      }
    }
  }

  get size() {
    let size = 0;
    for (const _ of this) {
      size++;
    }
    return size;
  }

  clear(): void {
    this.#data.clear();
  }

  delete(keys: K): boolean {
    const [map, key] = this.#traverse(keys);
    return map?.delete(key) ?? false;
  }

  get(keys: K): V | undefined {
    const [map, key] = this.#traverse(keys);
    return map?.get(key);
  }

  has(keys: K): boolean {
    const [map, key] = this.#traverse(keys);
    return map?.has(key) ?? false;
  }

  set(keys: K, value: V): this {
    this.#depth = keys.length;
    const [map, key] = this.#traverse(keys, (map, key) => {
      map.set(key, new Map());
    });
    map.set(key, value);

    return this;
  }

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

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

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

  entries(): IterableIterator<[K, V]> {
    return this.#entries();
  }

  [Symbol.iterator](): IterableIterator<[K, V]> {
    return this.#entries();
  }

  get [Symbol.toStringTag]() {
    return 'MultiKeyMap';
  }

  *#entries(
    keys: K[number][] = [],
    map = this.#data,
    depth = this.#depth,
  ): IterableIterator<[K, V]> {
    for (const [key, value] of map.entries()) {
      if (depth > 1) {
        yield* {
          [Symbol.iterator]: () => {
            return this.#entries([...keys, key], value, depth - 1);
          },
        };
      } else {
        yield [[...keys, key] as K, value as unknown as V];
      }
    }
  }

  #traverse(
    keys: K,
    onMissing: (map: NestedMap<K, V>, key: K[number]) => void,
  ): [Map<K[-1], V>, K[-1]];
  #traverse(keys: K): [undefined] | [Map<K[-1], V>, K[-1]];
  #traverse(
    keys: K,
    onMissing?: (map: NestedMap<K, V>, key: K[number]) => void,
  ): [undefined] | [Map<K[-1], V>, K[-1]] {
    let map = this.#data;
    for (const key of keys.slice(0, -1)) {
      if (!map.has(key)) {
        if (onMissing) {
          onMissing(map, key);
        } else {
          return [undefined];
        }
      }
      map = map.get(key)!;
    }
    return [map as unknown as Map<K[-1], V>, keys.slice(-1)[0]];
  }
}

// multi-key-map.spec.ts

import { MultiKeyMap } from '../src/multi-key-map';

test('can set and get items using multiple keys', () => {
  const map = createMap();

  expect(map.get(['a', 1])).toBe('A');
  expect(map.get(['a', 2])).toBe('B');
  expect(map.get(['b', 1])).toBe('C');
  expect(map.get(['b', 2])).toBe('D');
});

test('can create a map from an iterable', () => {
  const map = new MultiKeyMap<[string, number], string>([
    [['a', 1], 'A'],
    [['a', 2], 'B'],
    [['b', 1], 'C'],
    [['b', 2], 'D'],
  ]);

  expect(map.size).toBe(4);
});

test('size should be the count of all items', () => {
  const map = createMap();

  expect(map.size).toBe(4);
});

test('can clear all items from the map', () => {
  const map = createMap();

  map.clear();

  expect(map.size).toBe(0);
});

test('returns true when an item was deleted from the map', () => {
  const map = createMap();

  expect(map.delete(['a', 1])).toBe(true);
  expect(map.size).toBe(3);
});

test('returns false when an item was not deleted from the map', () => {
  const map = createMap();

  expect(map.delete(['c', 1])).toBe(false);
  expect(map.size).toBe(4);
});

test('can check if an item is in the map', () => {
  const map = createMap();

  expect(map.has(['a', 1])).toBe(true);
  expect(map.has(['a', 3])).toBe(false);
  expect(map.has(['c', 1])).toBe(false);
});

test('can overwrite the values of an existing key', () => {
  const map = createMap();

  map.set(['a', 1], 'Z');

  expect(map.get(['a', 1])).toBe('Z');
});

test('can iterate over all the values using forEach', () => {
  const map = createMap();
  const callback = jest.fn();

  map.forEach(callback);

  expect(callback).toHaveBeenCalledWith('A', ['a', 1], map);
  expect(callback).toHaveBeenCalledWith('B', ['a', 2], map);
  expect(callback).toHaveBeenCalledWith('C', ['b', 1], map);
  expect(callback).toHaveBeenCalledWith('D', ['b', 2], map);
});

test('can iterate over all the keys', () => {
  const map = createMap();

  const [...keys] = map.keys();

  expect(keys).toEqual([
    ['a', 1],
    ['a', 2],
    ['b', 1],
    ['b', 2],
  ]);
});

test('can iterate over all the values', () => {
  const map = createMap();

  const [...values] = map.values();

  expect(values).toEqual(['A', 'B', 'C', 'D']);
});

test('can iterate over all the entries', () => {
  const map = createMap();

  const [...entries] = map.entries();

  expect(entries).toEqual([
    [['a', 1], 'A'],
    [['a', 2], 'B'],
    [['b', 1], 'C'],
    [['b', 2], 'D'],
  ]);
});

test('map is iterable', () => {
  const map = createMap();

  const [...entries] = map;

  expect(entries).toEqual([
    [['a', 1], 'A'],
    [['a', 2], 'B'],
    [['b', 1], 'C'],
    [['b', 2], 'D'],
  ]);
});

test('can create deep maps', () => {
  const map = new MultiKeyMap([
    [['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'], 'A'],
  ]);

  expect(map.get(['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j'])).toBe('A');
});

function createMap() {
  const map = new MultiKeyMap<[string, number], string>();

  map.set(['a', 1], 'A');
  map.set(['a', 2], 'B');
  map.set(['b', 1], 'C');
  map.set(['b', 2], 'D');

  return map;
}