memoize.ts

type CacheType<T> = {
    has: (key: string) => boolean;
    get: (key: string) => T | undefined;
    set: (key: string, result: T, now?: number) => void;
    clear?: () => void;
};

const createCache = <T>(maxAge = 60_000): CacheType<T> => {
    const values = new Map<string, T>();
    const expiry = new Map<string, number>();

    const set = (key: string, result: T, now = Date.now()) => {
        values.set(key, result);
        expiry.set(key, now + maxAge);
    };

    const clear = (now = Date.now()) => {
        for (const [key, expiryTime] of expiry.entries()) {
            if (now >= expiryTime) {
                values.delete(key);
                expiry.delete(key);
            }
        }
    };

    return {
        has: values.has,
        get: values.get,
        clear,
        set,
    };
};

const createTimeoutCache = <T>(maxAge = 60_000): CacheType<T> => {
    const values = new Map<string, T>();

    const set = (key: string, result: T) => {
        values.set(key, result);

        const timer = setTimeout(() => {
            values.delete(key);
        }, maxAge);

        timer.unref?.();
    };

    return {
        has: values.has,
        get: values.get,
        set,
        clear: () => {},
    };
};

interface MemoizeOptions<T, A extends unknown[]> {
    maxAge?: number;
    getKey?: (args: A) => string;
    /**
     * Override
     */
    cache?: CacheType<T>;
}

const memoize = <T, A extends unknown[]>(
    fn: (...args: A) => T,
    options?: MemoizeOptions<T, A>,
) => {
    const {
        maxAge = 60_000,
        getKey = JSON.stringify,
        cache = createCache(maxAge),
    } = options || {};

    return (...args: A) => {
        const key = getKey(args);
        const now = Date.now();

        cache.clear?.();

        if (cache.has(key)) {
            return cache.get(key) as T;
        }

        const result = fn(...args);

        cache.set(key, result, now);

        return result;
    };
};

const memoizePromise = <T, A extends unknown[]>(
    fn: (...args: A) => Promise<T>,
    options?: MemoizeOptions<T, A>,
) => {
    const {
        maxAge = 60_000,
        getKey = JSON.stringify,
        cache = createCache(maxAge),
    } = options || {};

    return async (...args: A) => {
        const key = getKey(args);
        const now = Date.now();

        cache.clear?.();

        if (cache.has(key)) {
            return cache.get(key) as T;
        }

        const result = await fn(...args);

        cache.set(key, result, now);

        return result;
    };
};

// Usage example:
const memoizedFunction = memoizePromise(
    async (query: string) => console.log(query),
    {
        // maxAge: 5000,
        cache: createTimeoutCache(),
    },
);