memo.ts

import * as lru from "lru-cache";

const DEFAULT_LRU_MAX = 1000;

// need this for default arguments
const getArgumentCount = (fn: Function) => {
  const str = fn.toString();
  const params = str.match(/\(.*?\)/)[0];
  const args = params
    .substr(1, params.length - 2)
    .split(/\s*,\s*/)
    .filter((arg) => arg.substr(0, 3) !== "...");

  return args.length;
};

export const memoize = <TFunc extends (...args: any[]) => any>(
  fn: TFunc,
  lruMax: number = DEFAULT_LRU_MAX,
  argumentCount: number = getArgumentCount(fn)
) => {
  if (argumentCount == Infinity || isNaN(argumentCount)) {
    throw new Error(`Argument count cannot be Infinity, 0, or NaN.`);
  }

  if (!argumentCount) {
    console.error(`Argument count should not be 0. Defaulting to 1.`);
    argumentCount = 1;
  }

  return compilFastMemoFn(argumentCount, lruMax > 0)(
    fn,
    lru({ max: lruMax })
  ) as TFunc;
};

let _memoFns = {};

const compilFastMemoFn = (argumentCount: number, acceptPrimitives: boolean) => {
  const hash = "" + argumentCount + acceptPrimitives;
  if (_memoFns[hash]) {
    return _memoFns[hash];
  }

  const args = Array.from({ length: argumentCount }).map((v, i) => `arg${i}`);

  let buffer = `
  return function(fn, keyMemo) {
    var memo = new WeakMap();
    return function(${args.join(", ")}) {
      var currMemo = memo, prevMemo, key;
  `;

  for (let i = 0, n = args.length - 1; i < n; i++) {
    const arg = args[i];
    buffer += `
      prevMemo = currMemo;
      key      = ${arg};
      ${
        acceptPrimitives
          ? `if ((typeof key !== "object" || !key) && !(key = keyMemo.get(${arg}))) {
        keyMemo.set(${arg}, key = {});
      }`
          : ""
      }
      if (!(currMemo = currMemo.get(key))) {
        prevMemo.set(key, currMemo = new WeakMap());
      }
    `;
  }

  const lastArg = args[args.length - 1];

  buffer += `
      key = ${lastArg};
      ${
        acceptPrimitives
          ? `
      if ((typeof key !== "object" || !key) && !(key = keyMemo.get(${lastArg}))) {
        keyMemo.set(${lastArg}, key = {});
      }`
          : ""
      }
      if (!currMemo.has(key)) {
        currMemo.set(key, fn(${args.join(", ")}));
      }
      return currMemo.get(key);
    };
  };
  `;

  return (_memoFns[hash] = new Function(buffer)());
};