Memoize decorator based on WeakMaps (respects non-primitive arguments by reference)

memoize.decorator.ts

const globalCache = new WeakMap();

export function Memoize() {
  return (target: any, propertyName: string, descriptor: TypedPropertyDescriptor<any>) => {
    if (descriptor.value != null) {
      descriptor.value = getNewFunction(descriptor.value);
    } else if (descriptor.get != null) {
      descriptor.get = getNewFunction(descriptor.get);
    } else {
      throw new Error('Only put a Memoize decorator on a method or get accessor.');
    }
  };
}

class MixedMap {

  map: Map<any, any> = new Map();
  weakMap: WeakMap<any, any> = new WeakMap();

  has(key: any) {
    return this.getCorrespondingMap(key).has(key);
  }

  get(key: any) {
    return this.getCorrespondingMap(key).get(key);
  }

  set(key: any, value: any) {
    this.getCorrespondingMap(key).set(key, value);

    return this;
  }

  clear() {
    this.map.clear();
    this.map = new Map();
    this.weakMap = new WeakMap();
  }

  private getCorrespondingMap(key: any) {
    return key instanceof Object ? this.weakMap : this.map;
  }

}

function getNewFunction(originalFunction: (...params: any[]) => void) {
  return function (this: any, ...args: any[]) {
    const thisCache = cache(globalCache, this, () => new WeakMap());
    const fnCache = cache(thisCache, originalFunction, () => new MixedMap());

    return cache(fnCache, args[0], () => originalFunction.apply(this, args));
  };
}

function cache<V>(wm: MixedMap | WeakMap<any, any>, key: any, createValue: () => V) {
  let value: V;

  if (!wm.has(key)) {
    wm.set(key, value = createValue());
  } else {
    value = wm.get(key);
  }

  return value;
}

memoize.spec.ts

import { Memoize } from './memoize.decorator';

describe('Memoize', () => {
  class TestClass {
    @Memoize() getNumber(multiplier = 1) {
      return multiplier * Math.random();
    }

    @Memoize() getObjectNumber(object?: TestSubclass) {
      return object ? object.getNumber() : null;
    }

    @Memoize() get value() {
      return Math.random();
    }
  }

  class TestSubclass {
    getNumber() {
      return Math.random();
    }
  }

  const [a, b] = [new TestClass(), new TestClass()];

  it('method should be memoized', () => {
    expect(a.getNumber()).toBe(a.getNumber());
    expect(a.getNumber(2)).toBe(a.getNumber(2));
    expect(a.getNumber(4)).toBe(a.getNumber(4));
    expect(a.getNumber()).not.toBe(a.getNumber(4));
    expect(a.getNumber(2)).not.toBe(a.getNumber(4));
  });

  it('method should be memoized by a primitive argument', () => {
    expect(a.getNumber(2)).toBe(a.getNumber(2));
    expect(a.getNumber(4)).toBe(a.getNumber(4));
    expect(a.getNumber(null)).toBe(a.getNumber(null));
    expect(a.getNumber()).not.toBe(a.getNumber(4));
    expect(a.getNumber(2)).not.toBe(a.getNumber(4));
    expect(a.getNumber(null)).not.toBe(a.getNumber());
  });

  it('method should be memoized by a non-primitive argument', () => {
    const [npa, npb] = [new TestSubclass(), new TestSubclass()];

    expect(a.getObjectNumber(npa)).toBe(a.getObjectNumber(npa));
    expect(a.getObjectNumber(npa)).not.toBe(a.getObjectNumber());
    expect(a.getObjectNumber(npa)).not.toBe(a.getObjectNumber(null));
    expect(a.getObjectNumber(npa)).not.toBe(a.getObjectNumber(npb));
  });

  it('accessor should be memoized', () => {
    expect(a.value).toBe(a.value);
  });

  it('multiple instances shouldn\'t share values for methods', () => {
    expect(a.getNumber()).not.toBe(b.getNumber());
  });

  it('multiple instances shouldn\'t share values for accessors', () => {
    expect(a.getNumber()).not.toBe(b.getNumber());
  });
});