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