Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Memoize decorator based on WeakMaps (respects non-primitive arguments by reference)
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;
}
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());
});
});
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment