type MemoizedFunction<T extends (...args: any[]) => Promise<any> | any> = (...args: Parameters<T>) => Promise<ReturnType<T>> | ReturnType<T>;

export function memoize<T extends (...args: any[]) => Promise<any> | any>(func: T): MemoizedFunction<T> {
  const cache = new Map<string, ReturnType<T>>();

  return (...args: Parameters<T>): Promise<ReturnType<T>> | ReturnType<T> => {
    const key = JSON.stringify(args);
    if (cache.has(key)) {
      return cache.get(key)!;
    }
    const result = func(...args);
    if (result instanceof Promise) {
      return result.then(resolved => {
        return resolved;
      });
    } else {
      cache.set(key, result);
      return result;
    }
  };
}