Nested weak map

demo.ts

import { NestedWeakMap } from './nested-weak-map';

const key1 = {};
const key2 = {};
const map = new NestedWeakMap();

map.get([key1, key2]); // undefined
map.getOrCreate([key1, key2], () => 'default'); // "default"
map.get([key1, key2]); // "default"

nested-weak-map.ts

export class NestedWeakMap<T> {
  private content = new WeakMap<any, any>();

  get(keys: object[]): T | undefined {
    let current: any = this.content;
    for (const key of keys) {
      current = current.get(key);
      if (!current) {
        return;
      }
    }
    return current;
  }

  set(keys: object[], value: T) {
    let current = this.content;
    const initialKeys = keys.slice(0, -1);
    const last = keys[keys.length - 1];
    for (const key of initialKeys) {
      let next = current.get(key);
      if (!next) {
        next = new WeakMap();
        current.set(key, next);
      }
      current = next;
    }
    current.set(last, value);
  }

  getOrCreate(keys: object[], getDefaultValue: () => T) {
    let value = this.get(keys);
    if (!value) {
      value = getDefaultValue();
      this.set(keys, value);
    }
    return value;
  }
}