Caching for Javascript/Typescript using arbitrary objects as keys

cache.ts

import * as refs from "./refs";

class ById {
    constructor(public param: any) { }
}

// wraps an object in a special flag so that the cache uses the object's id instead of hashing its keys and values
export function byId(param: any) {
    return new ById(param)
}

function makeSeed(): number {
    return Math.floor(Math.random() * Number.MAX_SAFE_INTEGER)
}

let _cache = new Map()
let _seed = makeSeed()

/**
    The keys passed are hashed in such a way to avoid conflicts.

    For example, if you pass the [97, 98] it will has to a different value than if you pass the string "ab"
    This is because the type information is also used in the hash.

    To help the cache be more robust in this way, don't pass results of calls to `refs.objectId`, instead, pass the object wrapped by `cache.byId`, and if you want to pass a function, pass it directly.
*/

export function get<F extends (...args: any[]) => any>(key: any, compute: F, ...p: Parameters<F>): ReturnType<F> {
    let hash = hashAny(key, _seed);
    let cached = _cache.get(hash)
    if (cached !== undefined) {
        return cached
    } else {
        let computed = compute.apply(null, p);
        _cache.set(hash, computed);
        return computed;
    }
}

export function peek<T>(key: any): T | undefined {
    let hash = hashAny(key, _seed);
    return _cache.get(hash)
}

/**
    Clears all references held by this cache to allow them to be garbage collected
*/
export function clear() {
    _cache.clear()
    _seed = makeSeed()
}

export function fn<P, T>(fn: (p: P) => T, p: P): T {
    return get([fn, byId(p)], fn, p)
}

export function fnPeek<P, T>(fn: (p: P) => T, p: P): T | undefined {
    return peek([fn, byId(p)])
}

// based on cyrb53
// can be used to hash an object or a list
// must always be composed of numbers, booleans, or strings.
// not supported: bigint, symbol
// functions are hashed by id (using refs.objectId)
// special ById class is used to hash objects by id
export function hashAny(input: any, seed = 0): number {

    // We use these consts to identify data structure and make it robust against accidental similarity
    // For example, if we have the array ['A', 'B'] and the object {'A': 'B'} we don't
    // want to just digest the string 'A' then 'B' because then they will produce the same hash
    // so we also digest information that describes the structure of the object

    // ['A', 'B'] => ARRAY 2 ITEM 1 STRING 1 'A' ITEM 2 STRING 1 'B'
    // {'A': 'B'} => OBJECT 1 KEY 1 'A' VALUE STRING 1 'B'

    const NUM = 1;
    const BOOL = 2;
    const STRING = 3;
    const ARRAY = 4;
    const OBJECT = 5;
    const NULL = 6;
    const UNDEFINED = 7;
    const KEY = 8;
    const VALUE = 9;
    const ITEM = 10;

    const FUNC = 11;
    const REF = 12;

    let h1 = 0xdeadbeef ^ seed, h2 = 0x41c6ce57 ^ seed;
    function digest(n: number) {
        h1 = Math.imul(h1 ^ n, 2654435761);
        h2 = Math.imul(h2 ^ n, 1597334677);
    }

    function digestString(str: string) {
        digest(str.length)
        for (let i = 0; i < str.length; i++) {
            let ch = str.charCodeAt(i);
            digest(ch)
        }
    }

    function digestObject(obj: any) {
        let keys = Object.keys(obj);
        // potential bottle neck but absolutely necessary for stability
        keys.sort();
        digest(OBJECT);
        digest(keys.length);
        for (let k of keys) {
            digest(KEY);
            digestString(k)
            digest(VALUE);
            digestAny(obj[k])
        }
    }

    function digestArray(arr: any[]) {
        digest(ARRAY);
        digest(arr.length)
        for (let i = 0; i < arr.length; i++) {
            digest(ITEM);
            digest(i);
            digestAny(arr[i])
        }
    }

    function digestAny(obj: any) {
        switch (typeof obj) {
            case 'number':
                digest(NUM);
                digest(obj);
                break;
            case 'boolean':
                digest(BOOL);
                digest(obj ? 1 : 0);
                break;
            case 'string':
                digest(STRING);
                digestString(obj);
                break;
            case 'object':
                if (obj === null) {
                    digest(NULL);
                } else if (Array.isArray(obj)) {
                    digestArray(obj);
                } else if (obj instanceof ById) {
                    digest(REF);
                    digest(refs.objectId(obj.param))
                } else {
                    digestObject(obj);
                }
                break;
            case 'function':
                digest(FUNC);
                digest(refs.objectId(obj))
            case 'undefined':
                digest(UNDEFINED);
                break;
            default:
                throw 'Unexpected type'
        }
    }

    digestAny(input);

    h1 = Math.imul(h1 ^ (h1 >>> 16), 2246822507) ^ Math.imul(h2 ^ (h2 >>> 13), 3266489909);
    h2 = Math.imul(h2 ^ (h2 >>> 16), 2246822507) ^ Math.imul(h1 ^ (h1 >>> 13), 3266489909);
    return 4294967296 * (2097151 & h2) + (h1 >>> 0);
}

refs.ts

function _sequentialGenerator() {
    let _id = 0;
    return (): number => {
        let res = _id;
        _id += 1;
        return res;
    }
}

const nextObjectId = _sequentialGenerator();

let idmap = new Map<unknown, number>()
let objmap = new Map<number, unknown>()

export function objectId(obj: unknown): number {
    let v = idmap.get(obj)
    if (v === undefined) {
        v = nextObjectId()
        idmap.set(obj, v)
        objmap.set(v, obj)
    }
    return v
}

export function objectById<T = any>(id: number): T | null {
    return objmap.get(id) as T ?? null;
}

export function clear() {
    idmap.clear()
    objmap.clear()
}

declare const _ref_type: unique symbol;
export type Ref<T = any> = RawRef & {
    [_ref_type]: T
}

export type RawRef<T = any> = {
    obj: T;
    key: keyof T;
}

export function ref<T, K extends keyof T>(obj: T, key: K): Ref<T[K]> {
    return { obj, key } as Ref<T[K]>
}

export function get<T>(r: Ref<T>): T {
    return r.obj[r.key]
}

export function set<T>(r: Ref<T>, value: T) {
    r.obj[r.key] = value
}