Naive memoize gist

memoize.ts

/* tslint:disable:ban-types */

const BOUND = Symbol('BOUND_FUNCTION')
const UNDEFINED = Symbol('UNDEFINED')
const NULL = Symbol('NULL')
const WM_STORE = Symbol('WM_STORE')
const RESULT = Symbol('RESULT')
interface IST {
    [WM_STORE]?: WeakMap<any, IST>;
    [RESULT]?: any;
    [key : string]: undefined|IST;
    [NULL]?: IST;
    [BOUND]?: IST;
    [UNDEFINED]?: IST;
}

const store : IST = {[WM_STORE] : new WeakMap<any, IST>()}

function key(p : any) :any {
    if (p === undefined) { return UNDEFINED; }
    if (p === null) { return NULL; }
    return p
}

function isObject(p : any) : boolean {
    return ['function', 'object'].includes(typeof p)
}

function getIn(path : any[]) : any {
    let r : undefined | IST = store;

    for (let p of path) {
        if (r === undefined) { return r; }
        p = key(p)

        if (!isObject(p)) {
            r = r[p === undefined ? UNDEFINED : p]
        } else {
            const wm = r[WM_STORE]
            if (!wm) { return undefined; }
            r = wm.get(p)
        }
    }

    if (r === undefined) { return r; }
    return r[RESULT];
}

function setIn(path : any[], value: any) : void {
    let r : any = store;

    for (let p of path) {
        p = key(p)
        if (!isObject(p)) {
            if (!r[p]) {
                r[p] = {}
            }
            r = r[p]
        } else {
            if (!r[WM_STORE]) {
                r[WM_STORE] = new WeakMap()
            }
            r = r[WM_STORE]
            if (!r.has(p)) {
                r.set(p, {})
            }
            r = r.get(p)
        }
    }
    r[RESULT] = value
}

function intercept(x : any) : any {
    if (x === UNDEFINED) { return undefined; }
    return x
}

export default function memoize<R, T extends (...args: any[]) => R>(f : T, isBound : boolean = false) : T {
    const testF = getMemoized(f);
    if (testF) { return testF; }
    isBound = isBound && !f.prototype

    const memoized : any = function (this:any, ...args : any[]) : R {
        const that : any = isBound ? BOUND : this
        const path = [f, that, ...args]
        const result = getIn(path)
        if (result !== undefined) { return intercept(result) }
        const e = f.apply(that, args)
        setIn(path, e)
        return e
    }
    const memf =  isBound ? memoized.bind(null) : memoized
    setMemoized(f, memf)
    return memf
}

const memFunctions = new WeakMap<Function, any>()
const genByMem = new WeakSet<Function>()

function getMemoized<T extends Function>(f : T) : undefined|T {
    const testF = memFunctions.get(f)
    if (testF) { return testF }
    if (genByMem.has(f)) { return f }
    return undefined
}

function setMemoized<T extends Function>(f : T, mem: T) : void {
    memFunctions.set(f, mem)
    genByMem.add(mem)
}